From a919c79cb2587544b59861bf38b7d9e82940d800 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 18:02:01 +0300 Subject: [PATCH] =?UTF-8?q?feat(ai-chat):=20=D0=BF=D0=B5=D1=80=D1=81=D0=B8?= =?UTF-8?q?=D1=81=D1=82=20step-=D0=BC=D0=B0=D1=80=D0=BA=D0=B5=D1=80=D0=B0?= =?UTF-8?q?=20(#491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Персист отставал от 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) --- .../core/ai-chat/ai-chat-run.service.spec.ts | 73 ++++++++++++ .../src/core/ai-chat/ai-chat-run.service.ts | 34 ++++++ .../src/core/ai-chat/ai-chat.service.ts | 40 +++++++ .../core/ai-chat/ai-chat.step-marker.spec.ts | 104 ++++++++++++++++++ 4 files changed, 251 insertions(+) create mode 100644 apps/server/src/core/ai-chat/ai-chat.step-marker.spec.ts diff --git a/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts index ade4563c..41640cb0 100644 --- a/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts @@ -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(); + }); +}); diff --git a/apps/server/src/core/ai-chat/ai-chat-run.service.ts b/apps/server/src/core/ai-chat/ai-chat-run.service.ts index 1530c3b5..37ab3827 100644 --- a/apps/server/src/core/ai-chat/ai-chat-run.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat-run.service.ts @@ -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( diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index f1cda9b4..deb5b127 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -2721,6 +2721,34 @@ export function rowToUiMessage(row: AiChatMessage): Omit & { 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 | 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 + @@ -2891,6 +2919,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). diff --git a/apps/server/src/core/ai-chat/ai-chat.step-marker.spec.ts b/apps/server/src/core/ai-chat/ai-chat.step-marker.spec.ts new file mode 100644 index 00000000..98c53e4b --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat.step-marker.spec.ts @@ -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>; + 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>; + 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); + }); +});