feat(ai-chat): персист step-маркера (#491)

Персист отставал от live-стрима на целый текущий шаг, а run.stepCount как источник границы шага НЕгоден: recordStep — fire-and-forget, не атомарен с записью parts (рассинхрон сид↔маркер). Пишем маркер ТЕМ ЖЕ flush'ем, что и parts: - flushAssistant стампует metadata.stepsPersisted = число ЗАВЕРШЁННЫХ шагов, чьи parts лежат в ЭТОЙ строке (оба выводятся из finished → маркер не может разойтись с persisted parts). In-progress хвост (частичный шаг при error/abort или mid-stream flush) НЕ считается. Это step-alignment якорь, на котором строится resume-стек (ротация кольца по подтверждённому шагу N — коммит 3; attach режет хвост по «шаг > N»). - Контракт AiChatRunService.reconstructRunParts(runId) → { parts, stepsPersisted } — единственный интерфейс чтения ЖИВОГО рана (run → assistantMessageId → строка → чистый reconstructPartsFromRow). null при отсутствии рана / связанной строки / удалённой строки; маркер 0 у pre-#491 строки — безопасный пол. Потребители: attach (коммит 3), дельта (rows уже несут маркер в metadata), экспорт. messageRepo — опциональный 3-й параметр конструктора (2-арг тест-конструкции компилируются без изменений). - /messages и дельта отдают маркер в row автоматически (он внутри metadata). Тесты: property «stepsPersisted == число завершённых шагов для любого N + parts согласованы»; частичный хвост не инкрементит маркер; reconstructPartsFromRow (маркер, дефолт 0, фолбэк на content); reconstructRunParts (резолв + null-кейсы). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 18:02:01 +03:00
parent 1e28151f7a
commit 108d2c6bef
4 changed files with 251 additions and 0 deletions
@@ -771,3 +771,76 @@ describe('#487 AiChatRunService.supersede (CAS)', () => {
    expect(svc.hasZombie('run-1')).toBe(false);
  });
 });
+
+describe('AiChatRunService.reconstructRunParts (#491)', () => {
+  const ws = 'ws-1';
+
+  function makeMsgRepo(row: unknown) {
+    return { findById: jest.fn(async () => row) };
+  }
+
+  it('resolves run -> assistant row -> { parts, stepsPersisted }', async () => {
+    const repo = makeRepo({
+      findById: jest.fn(async () => ({
+        id: 'run-1',
+        assistantMessageId: 'm1',
+      })),
+    });
+    const msgRepo = makeMsgRepo({
+      content: '',
+      metadata: {
+        parts: [{ type: 'text', text: 'hi' }],
+        stepsPersisted: 3,
+      },
+    });
+    const svc = new AiChatRunService(
+      repo as never,
+      makeEnv() as never,
+      msgRepo as never,
+    );
+    const res = await svc.reconstructRunParts('run-1', ws);
+    expect(res).toEqual({
+      parts: [{ type: 'text', text: 'hi' }],
+      stepsPersisted: 3,
+    });
+    expect(msgRepo.findById).toHaveBeenCalledWith('m1', ws);
+  });
+
+  it('returns null when the run has no linked assistant row yet (seed window)', async () => {
+    const repo = makeRepo({
+      findById: jest.fn(async () => ({ id: 'run-1', assistantMessageId: null })),
+    });
+    const msgRepo = makeMsgRepo({ content: 'x', metadata: null });
+    const svc = new AiChatRunService(
+      repo as never,
+      makeEnv() as never,
+      msgRepo as never,
+    );
+    expect(await svc.reconstructRunParts('run-1', ws)).toBeNull();
+    expect(msgRepo.findById).not.toHaveBeenCalled();
+  });
+
+  it('returns null when the run does not exist', async () => {
+    const repo = makeRepo({ findById: jest.fn(async () => undefined) });
+    const msgRepo = makeMsgRepo({});
+    const svc = new AiChatRunService(
+      repo as never,
+      makeEnv() as never,
+      msgRepo as never,
+    );
+    expect(await svc.reconstructRunParts('nope', ws)).toBeNull();
+  });
+
+  it('returns null when the linked assistant row was deleted', async () => {
+    const repo = makeRepo({
+      findById: jest.fn(async () => ({ id: 'run-1', assistantMessageId: 'm1' })),
+    });
+    const msgRepo = makeMsgRepo(undefined);
+    const svc = new AiChatRunService(
+      repo as never,
+      makeEnv() as never,
+      msgRepo as never,
+    );
+    expect(await svc.reconstructRunParts('run-1', ws)).toBeNull();
+  });
+});
@@ -1,6 +1,9 @@
 import { Injectable, Logger, OnModuleInit } from '@nestjs/common';
 import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo';
+import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
 import { AiChatRun } from '@docmost/db/types/entity.types';
+import { reconstructPartsFromRow } from './ai-chat.service';
+import type { UIMessage } from 'ai';
 import { isUniqueViolation, violatedConstraint } from '@docmost/db/utils';
 import { EnvironmentService } from '../../integrations/environment/environment.service';

@@ -212,6 +215,10 @@ export class AiChatRunService implements OnModuleInit {
  constructor(
    private readonly runRepo: AiChatRunRepo,
    private readonly environment: EnvironmentService,
+    // #491: OPTIONAL so the many 2-arg test constructions of this service compile
+    // unchanged; Nest always injects the real repo in production. Only touched by
+    // reconstructRunParts (the live-run read interface).
+    private readonly messageRepo?: AiChatMessageRepo,
  ) {}

  /**
@@ -784,6 +791,33 @@ export class AiChatRunService implements OnModuleInit {
    return this.runRepo.findById(runId, workspaceId);
  }

+  /**
+   * #491 CONTRACT — the SINGLE interface for reading a LIVE run's output: resolve
+   * the run to its assistant-message projection (#183) and return its persisted
+   * `parts` plus `stepsPersisted` — the count of FINISHED steps confirmed on disk,
+   * written atomically with the parts (see {@link flushAssistant} / the step
+   * marker). Consumers: tail-only attach (commit 3 — slices at "step > N"), the
+   * degraded-poll delta (rows already carry the marker in metadata), and export.
+   *
+   * Returns `null` when the run does not exist, has no linked assistant row yet
+   * (the seed window), or the row was deleted — a consumer treats null as "nothing
+   * confirmed" and stays safe (attach 204 / full seed). `stepsPersisted` is 0 for a
+   * pre-#491 row with no marker, which is likewise the safe floor.
+   */
+  async reconstructRunParts(
+    runId: string,
+    workspaceId: string,
+  ): Promise<{ parts: UIMessage['parts']; stepsPersisted: number } | null> {
+    const run = await this.runRepo.findById(runId, workspaceId);
+    if (!run?.assistantMessageId || !this.messageRepo) return null;
+    const row = await this.messageRepo.findById(
+      run.assistantMessageId,
+      workspaceId,
+    );
+    if (!row) return null;
+    return reconstructPartsFromRow(row);
+  }
+
  /** The active run on a chat, if any (used to reject a concurrent start with a
   *  clean 409 before committing to the stream). */
  getActiveForChat(
@@ -2709,6 +2709,34 @@ export function rowToUiMessage(row: AiChatMessage): Omit<UIMessage, 'id'> & {
  return { id: row.id, role, parts: parts as UIMessage['parts'] };
 }

+/**
+ * The persisted parts + step marker of a run's assistant row (#491). This is the
+ * pure core of the `AiChatRunService.reconstructRunParts(runId)` contract — the
+ * SINGLE interface for reading a LIVE run's output — given the already-resolved
+ * row. Returns the `metadata.parts` (falling back to a single text part from
+ * `content` for a pre-#183 row) and `stepsPersisted` — the count of FINISHED steps
+ * whose parts are present, written atomically with the parts by
+ * {@link flushAssistant}. A missing marker (an old row) reads as 0, so a consumer
+ * treats it as "nothing confirmed" and stays safe (attach 204 / full seed). Pure.
+ */
+export function reconstructPartsFromRow(
+  row: Pick<AiChatMessage, 'content' | 'metadata'> | null | undefined,
+): { parts: UIMessage['parts']; stepsPersisted: number } {
+  const meta = (row?.metadata ?? {}) as {
+    parts?: UIMessage['parts'];
+    stepsPersisted?: number;
+  };
+  const parts =
+    Array.isArray(meta.parts) && meta.parts.length > 0
+      ? meta.parts
+      : textPart(row?.content ?? '');
+  const stepsPersisted =
+    typeof meta.stepsPersisted === 'number' && meta.stepsPersisted >= 0
+      ? meta.stepsPersisted
+      : 0;
+  return { parts: parts as UIMessage['parts'], stepsPersisted };
+}
+
 /**
 * The persisted-row patch shape produced by {@link flushAssistant}. It is the
 * SAME shape the assistant repo insert/update consume (content + toolCalls +
@@ -2879,6 +2907,18 @@ export function flushAssistant(
    // `parts`). Old rows have no marker and the legacy { output } shape; a
    // dual-shape query branches on this. Old rows are deliberately NOT migrated.
    toolTraceVersion: 2,
+    // #491 STEP MARKER: the number of FINISHED steps whose parts are in THIS row,
+    // written by the SAME flush that builds `parts` (atomically — they are both
+    // derived from `finished`, so the marker can NEVER disagree with the persisted
+    // parts). This is the step-alignment anchor the resume stack builds on:
+    // - the registry rotates its retention ring only on a CONFIRMED persist of
+    //   step N (commit 3);
+    // - attach slices the tail at "step > N" from the client's persisted seed.
+    // It is NOT `run.stepCount`: recordStep is fire-and-forget and NOT atomic with
+    // the parts write, so stepCount could race ahead of the persisted parts
+    // (seed↔marker drift). The in-progress trailing text (an error/abort partial,
+    // or a mid-stream flush) is NOT a finished step and is excluded from the count.
+    stepsPersisted: finished.length,
  };
  // finishReason: prefer an explicit one; else derive a sensible value from the
  // terminal status (so onError/onAbort records keep their historical reason).
@@ -0,0 +1,104 @@
+import {
+  flushAssistant,
+  reconstructPartsFromRow,
+} from './ai-chat.service';
+import type { AiChatMessage } from '@docmost/db/types/entity.types';
+
+/**
+ * #491 STEP MARKER — `metadata.stepsPersisted` is written by the SAME flush that
+ * builds `metadata.parts`, so the marker can never disagree with the persisted
+ * parts (the step-alignment anchor the resume stack builds on). These are
+ * PROPERTY tests: they assert the marker tracks the number of FINISHED steps for
+ * every flush shape, and that `reconstructPartsFromRow` reads it back safely.
+ */
+
+// A finished step carrying one line of text and one tool call/result.
+function step(i: number) {
+  return {
+    text: `step ${i}`,
+    toolCalls: [
+      { toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
+    ],
+    toolResults: [
+      { toolCallId: `c${i}`, toolName: 'getPage', output: { title: `T${i}` } },
+    ],
+  };
+}
+
+describe('flushAssistant step marker (#491)', () => {
+  it('seed (no steps) → stepsPersisted 0', () => {
+    const f = flushAssistant([], '', 'streaming');
+    expect(f.metadata.stepsPersisted).toBe(0);
+  });
+
+  it('PROPERTY: stepsPersisted equals the number of FINISHED steps, for any N', () => {
+    for (let n = 0; n <= 6; n++) {
+      const steps = Array.from({ length: n }, (_, i) => step(i));
+      const f = flushAssistant(steps, '', 'streaming');
+      expect(f.metadata.stepsPersisted).toBe(n);
+      // ...and the parts actually contain those N steps' text (marker agrees with
+      // the persisted parts — the atomicity the whole design relies on).
+      const parts = f.metadata.parts as Array<Record<string, unknown>>;
+      const textParts = parts.filter((p) => p.type === 'text');
+      expect(textParts).toHaveLength(n);
+    }
+  });
+
+  it('an in-progress trailing partial does NOT increment the marker', () => {
+    // 2 finished steps + a partial (not-yet-finished) trailing text: the marker
+    // counts only the CONFIRMED step boundaries, not the partial.
+    const f = flushAssistant([step(0), step(1)], 'partial third step', 'error', {
+      error: 'boom',
+    });
+    expect(f.metadata.stepsPersisted).toBe(2);
+    // The partial text IS persisted in parts (so the user sees it), but it is not a
+    // counted step.
+    const parts = f.metadata.parts as Array<Record<string, unknown>>;
+    expect(parts[parts.length - 1]).toEqual({
+      type: 'text',
+      text: 'partial third step',
+    });
+  });
+
+  it('terminal completed flush counts all finished steps', () => {
+    const f = flushAssistant([step(0), step(1), step(2)], '', 'completed', {
+      finishReason: 'stop',
+    });
+    expect(f.metadata.stepsPersisted).toBe(3);
+  });
+});
+
+describe('reconstructPartsFromRow (#491)', () => {
+  const row = (metadata: unknown, content = ''): AiChatMessage =>
+    ({ content, metadata }) as unknown as AiChatMessage;
+
+  it('reads parts + stepsPersisted from metadata', () => {
+    const f = flushAssistant([step(0), step(1)], '', 'streaming');
+    const r = reconstructPartsFromRow(row(f.metadata, f.content));
+    expect(r.stepsPersisted).toBe(2);
+    expect(r.parts).toEqual(f.metadata.parts);
+  });
+
+  it('defaults stepsPersisted to 0 for a pre-#491 row with no marker (safe floor)', () => {
+    const r = reconstructPartsFromRow(
+      row({ parts: [{ type: 'text', text: 'x' }] }),
+    );
+    expect(r.stepsPersisted).toBe(0);
+    expect(r.parts).toEqual([{ type: 'text', text: 'x' }]);
+  });
+
+  it('falls back to a single text part from content when no metadata.parts', () => {
+    const r = reconstructPartsFromRow(row(null, 'plain content'));
+    expect(r.stepsPersisted).toBe(0);
+    expect(r.parts).toEqual([{ type: 'text', text: 'plain content' }]);
+  });
+
+  it('null/undefined row → no parts, marker 0 (safe empty floor)', () => {
+    // textPart('') is empty, matching rowToUiMessage's fallback for an empty row.
+    expect(reconstructPartsFromRow(null)).toEqual({
+      parts: [],
+      stepsPersisted: 0,
+    });
+    expect(reconstructPartsFromRow(undefined).stepsPersisted).toBe(0);
+  });
+});