From 798a81abfe9f5d751fbbc9e6c2e1ba19ac0a29a8 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 17:55:06 +0300 Subject: [PATCH 1/7] =?UTF-8?q?perf(ai-chat):=20=D0=B4=D0=B5=D0=BB=D1=8C?= =?UTF-8?q?=D1=82=D0=B0-=D1=8D=D0=BD=D0=B4=D0=BF=D0=BE=D0=B8=D0=BD=D1=82?= =?UTF-8?q?=20=D0=BF=D0=BE=D0=BB=D0=BB=D0=B8=D0=BD=D0=B3=D0=B0=20+=20run-?= =?UTF-8?q?=D1=84=D0=B0=D0=BA=D1=82=20=D0=B2=20=D0=BE=D1=82=D0=B2=D0=B5?= =?UTF-8?q?=D1=82=D0=B5=20(#491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Degraded-poll рефетчил ВСЕ страницы infinite-query каждые 2.5 c с полными parts. Вводим дельта-эндпоинт «строки, изменённые после курсора»: - POST /ai-chat/messages/delta → { rows, cursor, run: {id,status}|null }. Курсор — таймстамп часов БД (now()), клиент эхом возвращает его каждый поллинг. Окно перекрытия now()−5s ловит строку, закоммиченную с updatedAt чуть раньше момента снятия предыдущего курсора на другом autocommit- соединении (одиночные UPDATE, длинных транзакций нет). Окно ГАРАНТИРУЕТ повторы → контракт: клиентский merge идемпотентен по id (mergeById). - Run-факт едет В дельте (отдельный /run-поллинг удвоил бы QPS — отвергнуто). - Все дельта-релевантные записи (message update/finalizeOwner/reconcile/sweep, run update/finalizeIfActive/markStopRequested/sweepRunning) стампуют updatedAt через SQL now(), а не app-clock new Date(): единая монотонная ось курсора, смешанные источники часов были независимым источником пропусков. - Клиент: сервис-функция getAiChatMessagesDelta + фиксация контракта идемпотентности mergeById тестом (свап degraded-поллинга на дельту — в коммите 3, вместе с re-seed-путём attach). Тесты (observable-property на живом Postgres, не моки): - дельта-семантика + монотонность курсора; - RACE «коммит позже с updatedAt раньше курсора» — ловится окном перекрытия (наивный updatedAt > cursor пропустил бы); - окно гарантирует повторы → идемпотентный merge; - updatedAt стампуется часами БД, а не app-clock (фейк системных часов в 2099 — стамп остаётся на времени БД); - контроллер: owner-gate, форма ответа, run-факт только {id,status}. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ai-chat/services/ai-chat-service.ts | 25 ++ .../ai-chat/utils/resume-helpers.test.ts | 33 +++ .../ai-chat/ai-chat.controller.delta.spec.ts | 108 +++++++++ .../src/core/ai-chat/ai-chat.controller.ts | 41 ++++ .../src/core/ai-chat/dto/ai-chat.dto.ts | 15 ++ .../ai-chat-message.repo.delta.spec.ts | 223 ++++++++++++++++++ .../repos/ai-chat/ai-chat-message.repo.ts | 103 +++++++- .../repos/ai-chat/ai-chat-run.repo.spec.ts | 15 +- .../repos/ai-chat/ai-chat-run.repo.ts | 25 +- 9 files changed, 570 insertions(+), 18 deletions(-) create mode 100644 apps/server/src/core/ai-chat/ai-chat.controller.delta.spec.ts create mode 100644 apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts diff --git a/apps/client/src/features/ai-chat/services/ai-chat-service.ts b/apps/client/src/features/ai-chat/services/ai-chat-service.ts index 34e6aa71..d5a6bfac 100644 --- a/apps/client/src/features/ai-chat/services/ai-chat-service.ts +++ b/apps/client/src/features/ai-chat/services/ai-chat-service.ts @@ -57,6 +57,31 @@ export async function stopRun( return req.data; } +/** + * Delta poll (#491): the chat's message rows changed since `cursor` (a DB-clock + * timestamp echoed from the previous poll) plus the current run fact, in ONE + * round-trip — the degraded-poll fallback's payload, replacing the old "refetch + * ALL infinite-query pages every 2.5s with full parts" poll. Omit `cursor` on the + * first poll (returns just a fresh cursor, no rows, to start the chain). The + * overlap window guarantees occasional REPEATS, so the caller MUST merge rows + * idempotently by id (mergeById). Owner-gated server-side. + */ +export async function getAiChatMessagesDelta( + chatId: string, + cursor?: string, +): Promise<{ + rows: IAiChatMessageRow[]; + cursor: string; + run: { id: string; status: string } | null; +}> { + const req = await api.post<{ + rows: IAiChatMessageRow[]; + cursor: string; + run: { id: string; status: string } | null; + }>("/ai-chat/messages/delta", { chatId, cursor }); + return req.data; +} + /** * #488: the run-fact — "is a run active on this chat?" — first-class from the * server (POST /ai-chat/run). Called on mount to seed the client FSM's run-fact diff --git a/apps/client/src/features/ai-chat/utils/resume-helpers.test.ts b/apps/client/src/features/ai-chat/utils/resume-helpers.test.ts index 8a4dc7ef..9278d004 100644 --- a/apps/client/src/features/ai-chat/utils/resume-helpers.test.ts +++ b/apps/client/src/features/ai-chat/utils/resume-helpers.test.ts @@ -109,4 +109,37 @@ describe("mergeById", () => { expect(mergeById(prev, null)).toBe(prev); expect(mergeById(prev, undefined)).toBe(prev); }); + + // #491 CONTRACT: the delta poll's overlap window GUARANTEES the same row is + // re-delivered across close polls, so merging must be IDEMPOTENT by id — merging + // the same row (or an equal-length list of rows) twice must not duplicate or + // reorder. This is the property the whole delta-poll design leans on; a + // regression here would re-introduce duplicate assistant bubbles on every poll. + it("is idempotent by id: re-merging the same row does not duplicate or reorder", () => { + const seed = [makeMsg("u1", "hi"), makeMsg("a1", "step 1")]; + const repeat = makeMsg("a1", "step 1"); // the SAME row the overlap re-delivers + const once = mergeById(seed, repeat); + const twice = mergeById(once, repeat); + const thrice = mergeById(twice, repeat); + // Length is stable (no growth), order is stable (user then assistant). + expect(once.map((m) => m.id)).toEqual(["u1", "a1"]); + expect(twice.map((m) => m.id)).toEqual(["u1", "a1"]); + expect(thrice.map((m) => m.id)).toEqual(["u1", "a1"]); + // The repeated merge converges: the row is replaced in place, never appended. + expect(twice[1]).toBe(repeat); + }); + + it("is idempotent across a batch of repeated + grown rows (delta re-delivery)", () => { + // A delta poll re-delivers a1 (unchanged) and a2 (grown one step). Applying the + // batch twice must equal applying it once — the poll can re-send either. + const start = [makeMsg("u1", "hi"), makeMsg("a1", "done")]; + const batch = [makeMsg("a1", "done"), makeMsg("a2", "grown step 2")]; + const apply = (list: typeof start) => + batch.reduce((acc, row) => mergeById(acc, row), list); + const once = apply(start); + const twice = apply(once); + expect(once.map((m) => m.id)).toEqual(["u1", "a1", "a2"]); + expect(twice.map((m) => m.id)).toEqual(["u1", "a1", "a2"]); + expect(twice).toEqual(once); + }); }); diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.delta.spec.ts b/apps/server/src/core/ai-chat/ai-chat.controller.delta.spec.ts new file mode 100644 index 00000000..2d418766 --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat.controller.delta.spec.ts @@ -0,0 +1,108 @@ +import { ForbiddenException } from '@nestjs/common'; +import { AiChatController } from './ai-chat.controller'; +import type { User, Workspace } from '@docmost/db/types/entity.types'; + +/** + * Wiring spec for the #491 delta-poll endpoint (`POST /ai-chat/messages/delta`). + * Owner-gated via assertOwnedChat (same gate as the other reads), NOT flag-gated. + * The run fact rides IN the delta response (no separate /run poll). Hand-rolled + * mocks — no Nest graph, no DB. Constructor order: (aiChatService, + * aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo). + */ +describe('AiChatController POST /ai-chat/messages/delta (#491)', () => { + const user = { id: 'u1' } as User; + const workspace = { id: 'ws1' } as Workspace; + + function makeController(opts: { + chat?: unknown; + delta?: { rows: unknown[]; cursor: string }; + run?: unknown; + }) { + const aiChatRunService = { + getLatestForChat: jest.fn().mockResolvedValue(opts.run), + }; + const aiChatRepo = { + findById: jest.fn().mockResolvedValue(opts.chat), + }; + const aiChatMessageRepo = { + findByChatUpdatedAfter: jest + .fn() + .mockResolvedValue(opts.delta ?? { rows: [], cursor: 'C1' }), + }; + const controller = new AiChatController( + {} as never, + aiChatRunService as never, + aiChatRepo as never, + aiChatMessageRepo as never, + {} as never, + {} as never, + ); + return { controller, aiChatRunService, aiChatRepo, aiChatMessageRepo }; + } + + it('owner-gates: a chat the user does not own throws, never reaching the repo', async () => { + const { controller, aiChatMessageRepo, aiChatRunService } = makeController({ + chat: { id: 'c1', creatorId: 'someone-else' }, + }); + await expect( + controller.getMessagesDelta({ chatId: 'c1' }, user, workspace), + ).rejects.toBeInstanceOf(ForbiddenException); + expect(aiChatMessageRepo.findByChatUpdatedAfter).not.toHaveBeenCalled(); + expect(aiChatRunService.getLatestForChat).not.toHaveBeenCalled(); + }); + + it('returns { rows, cursor, run:{id,status} } with the run fact inlined', async () => { + const rows = [{ id: 'm1' }]; + const { controller } = makeController({ + chat: { id: 'c1', creatorId: 'u1' }, + delta: { rows, cursor: 'C2' }, + run: { id: 'r1', status: 'running', error: 'ignored', stepCount: 3 }, + }); + const res = await controller.getMessagesDelta( + { chatId: 'c1', cursor: 'C1' }, + user, + workspace, + ); + expect(res).toEqual({ + rows, + cursor: 'C2', + // ONLY id + status — never the whole run row. + run: { id: 'r1', status: 'running' }, + }); + }); + + it('run is null when the chat has never had a run', async () => { + const { controller } = makeController({ + chat: { id: 'c1', creatorId: 'u1' }, + run: undefined, + }); + const res = await controller.getMessagesDelta( + { chatId: 'c1' }, + user, + workspace, + ); + expect(res.run).toBeNull(); + }); + + it('passes cursor through, defaulting a missing cursor to null (first poll)', async () => { + const { controller, aiChatMessageRepo } = makeController({ + chat: { id: 'c1', creatorId: 'u1' }, + }); + await controller.getMessagesDelta({ chatId: 'c1' }, user, workspace); + expect(aiChatMessageRepo.findByChatUpdatedAfter).toHaveBeenCalledWith( + 'c1', + 'ws1', + null, + ); + await controller.getMessagesDelta( + { chatId: 'c1', cursor: 'CX' }, + user, + workspace, + ); + expect(aiChatMessageRepo.findByChatUpdatedAfter).toHaveBeenLastCalledWith( + 'c1', + 'ws1', + 'CX', + ); + }); +}); diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.ts b/apps/server/src/core/ai-chat/ai-chat.controller.ts index f9b58b98..49267968 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.ts @@ -51,6 +51,7 @@ import { ChatIdDto, ExportChatDto, GeneratePageTitleDto, + GetChatDeltaDto, GetChatMessagesDto, GetRunDto, RenameChatDto, @@ -149,6 +150,46 @@ export class AiChatController { ); } + /** + * Delta poll (#491) — the degraded-poll fallback's payload. Returns the chat's + * message rows changed since `cursor` (a DB-clock timestamp from the previous + * poll), a FRESH cursor, AND the current run fact `{ id, status } | null`. This + * replaces the old degraded poll that refetched ALL infinite-query pages (full + * parts) every 2.5s: the client seeds once and thereafter merges only the + * deltas by id (the overlap window guarantees repeats — the merge is idempotent, + * see mergeById). The run fact rides IN the delta (a separate /run poll would + * double the poll QPS), so the client FSM gets the run's status on the same tick. + * Owner-gated via assertOwnedChat (same gate as the other read endpoints). + */ + @HttpCode(HttpStatus.OK) + @Post('messages/delta') + async getMessagesDelta( + @Body() dto: GetChatDeltaDto, + @AuthUser() user: User, + @AuthWorkspace() workspace: Workspace, + ): Promise<{ + rows: AiChatMessage[]; + cursor: string; + run: { id: string; status: string } | null; + }> { + await this.assertOwnedChat(dto.chatId, user, workspace); + const { rows, cursor } = + await this.aiChatMessageRepo.findByChatUpdatedAfter( + dto.chatId, + workspace.id, + dto.cursor ?? null, + ); + const run = await this.aiChatRunService.getLatestForChat( + dto.chatId, + workspace.id, + ); + return { + rows, + cursor, + run: run ? { id: run.id, status: run.status } : null, + }; + } + /** * Export a chat to Markdown (#183). The DB is the single source of truth: the * whole transcript is loaded (oldest -> newest) and rendered server-side. Now diff --git a/apps/server/src/core/ai-chat/dto/ai-chat.dto.ts b/apps/server/src/core/ai-chat/dto/ai-chat.dto.ts index 1fd3c8ff..ec09c6aa 100644 --- a/apps/server/src/core/ai-chat/dto/ai-chat.dto.ts +++ b/apps/server/src/core/ai-chat/dto/ai-chat.dto.ts @@ -37,6 +37,21 @@ export class GetChatMessagesDto { cursor?: string; } +/** + * Delta poll (#491): pull the chat's rows changed since `cursor` (a DB-clock + * timestamp from the previous poll) plus the current run fact — the degraded-poll + * fallback's payload, replacing the full infinite-query refetch. Omit `cursor` on + * the first poll (returns just a fresh cursor to start the chain). + */ +export class GetChatDeltaDto { + @IsString() + chatId: string; + + @IsOptional() + @IsString() + cursor?: string; +} + /** Resolve the chat bound to a document (the page's most-recent owned chat). */ export class BoundChatDto { @IsString() diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts new file mode 100644 index 00000000..e3086de8 --- /dev/null +++ b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts @@ -0,0 +1,223 @@ +import { randomUUID } from 'crypto'; +import { CamelCasePlugin, Kysely, sql } from 'kysely'; +import { PostgresJSDialect } from 'kysely-postgres-js'; +import postgres from 'postgres'; +import { AiChatMessageRepo } from './ai-chat-message.repo'; +import { AiChatRunRepo } from './ai-chat-run.repo'; + +/** + * #491 delta-poll — OBSERVABLE-PROPERTY tests against a LIVE Postgres (the local + * gitmost test DB, docker `gitmost-test-pg` on :5432), not "rows through a mock" + * (a mock cannot observe the DB clock nor the overlap-window race — the very + * things that matter here). Drives the REAL repo methods (`findByChatUpdatedAfter`, + * the now()-stamped `update`) and asserts: + * 1. delta-relevant writes stamp `updatedAt` from the DB clock, not the app clock + * (proven by faking the process clock far into the future and observing the + * stamp stays on real DB time); + * 2. the poll returns only rows changed after the cursor, ordered, with a fresh + * DB-clock cursor; + * 3. the "committed late but stamped earlier than the cursor" RACE is caught by + * the overlap window (a naive `updatedAt > cursor` would MISS it); + * 4. the overlap GUARANTEES repeats across close polls — the contract behind the + * client's idempotent merge (mergeById). + * + * FK triggers are bypassed (`session_replication_role = replica`) so synthetic + * chat/workspace ids need no parent fixtures; a single pooled connection (max 1) + * keeps that session setting for every query. SKIPS cleanly when the DB is + * unreachable so a DB-less CI never breaks. + */ +const CONN = + process.env.WAL_TEST_DATABASE_URL ?? + 'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost'; + +let db: Kysely; +let sqlClient: ReturnType; +let msgRepo: AiChatMessageRepo; +let runRepo: AiChatRunRepo; +let reachable = false; + +const workspaceId = randomUUID(); +const chatId = randomUUID(); + +beforeAll(async () => { + try { + sqlClient = postgres(CONN, { max: 1, onnotice: () => {} }); + db = new Kysely({ + dialect: new PostgresJSDialect({ postgres: sqlClient }), + plugins: [new CamelCasePlugin()], + }); + // Single connection keeps this session-scoped bypass for the whole suite. + await sql`set session_replication_role = replica`.execute(db); + await sql`select 1`.execute(db); + reachable = true; + } catch { + reachable = false; + } + msgRepo = new AiChatMessageRepo(db as never); + runRepo = new AiChatRunRepo(db as never); +}); + +afterAll(async () => { + if (db) { + try { + await db + .deleteFrom('aiChatMessages') + .where('workspaceId', '=', workspaceId) + .execute(); + await db + .deleteFrom('aiChatRuns') + .where('workspaceId', '=', workspaceId) + .execute(); + } catch { + /* best-effort cleanup */ + } + await db.destroy(); + } +}); + +afterEach(() => { + jest.useRealTimers(); +}); + +async function seedMessage(overrides: Record = {}) { + return msgRepo.insert({ + chatId, + workspaceId, + userId: null as never, + role: 'assistant', + content: 'x', + status: 'streaming', + ...overrides, + } as never); +} + +async function dbNow(): Promise { + const r = await sql<{ now: Date }>`select now() as now`.execute(db); + return r.rows[0].now.toISOString(); +} + +const maybe = (name: string, fn: () => Promise) => + it(name, async () => { + if (!reachable) { + console.warn(`SKIP (${name}): test DB unreachable at ${CONN}`); + return; + } + await fn(); + }); + +describe('AiChatMessageRepo.findByChatUpdatedAfter (#491 delta poll)', () => { + maybe('null cursor returns no rows and a fresh DB-clock cursor', async () => { + const before = await dbNow(); + const { rows, cursor } = await msgRepo.findByChatUpdatedAfter( + chatId, + workspaceId, + null, + ); + expect(rows).toEqual([]); + expect(new Date(cursor).getTime()).toBeGreaterThanOrEqual( + new Date(before).getTime(), + ); + }); + + maybe('returns only rows changed after the cursor', async () => { + const { cursor: c0 } = await msgRepo.findByChatUpdatedAfter( + chatId, + workspaceId, + null, + ); + const m = await seedMessage(); + const { rows, cursor: c1 } = await msgRepo.findByChatUpdatedAfter( + chatId, + workspaceId, + c0, + ); + expect(rows.map((r) => r.id)).toContain(m.id); + // Cursor is monotonic (advances). + expect(new Date(c1).getTime()).toBeGreaterThanOrEqual( + new Date(c0).getTime(), + ); + }); + + maybe( + 'RACE: a row stamped BEFORE the cursor but seen after is caught by the overlap', + async () => { + // Cursor taken now; then a row appears whose updatedAt is 2s in the PAST + // (committed late on another connection but stamped earlier). A naive + // `updatedAt > cursor` would MISS it; the 5s overlap window catches it. + const cursor = await dbNow(); + const m = await seedMessage(); + await sql`update ai_chat_messages set updated_at = now() - interval '2 seconds' where id = ${m.id}`.execute( + db, + ); + const { rows } = await msgRepo.findByChatUpdatedAfter( + chatId, + workspaceId, + cursor, + ); + expect(rows.map((r) => r.id)).toContain(m.id); + }, + ); + + maybe( + 'overlap GUARANTEES repeats across close polls (idempotent-merge contract)', + async () => { + const { cursor: c0 } = await msgRepo.findByChatUpdatedAfter( + chatId, + workspaceId, + null, + ); + const m = await seedMessage(); + const first = await msgRepo.findByChatUpdatedAfter( + chatId, + workspaceId, + c0, + ); + expect(first.rows.map((r) => r.id)).toContain(m.id); + // Immediately re-poll with the JUST-returned cursor: the row is still within + // the overlap window, so it is returned AGAIN — the client MUST dedupe by id. + const second = await msgRepo.findByChatUpdatedAfter( + chatId, + workspaceId, + first.cursor, + ); + expect(second.rows.map((r) => r.id)).toContain(m.id); + }, + ); + + maybe( + 'update() stamps updatedAt from the DB clock, not the app clock', + async () => { + const m = await seedMessage(); + // Fake the PROCESS clock ~73 years into the future. If the stamp came from + // `new Date()` the row would read year 2099; sql now() keeps it on DB time. + jest.useFakeTimers().setSystemTime(new Date('2099-01-01T00:00:00Z')); + const updated = await msgRepo.update(m.id, workspaceId, { + content: 'y', + }); + jest.useRealTimers(); + expect(updated).toBeDefined(); + expect(new Date(updated!.updatedAt).getFullYear()).toBeLessThan(2099); + }, + ); + + maybe( + 'run update() also stamps updatedAt from the DB clock', + async () => { + const run = await runRepo.insert({ + chatId, + workspaceId, + createdBy: null as never, + trigger: 'user', + status: 'running', + stepCount: 0, + } as never); + jest.useFakeTimers().setSystemTime(new Date('2099-01-01T00:00:00Z')); + const updated = await runRepo.update(run.id, workspaceId, { + stepCount: 1, + }); + jest.useRealTimers(); + expect(updated).toBeDefined(); + expect(new Date(updated!.updatedAt).getFullYear()).toBeLessThan(2099); + }, + ); +}); diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.ts b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.ts index 96070e77..57e82baa 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.ts @@ -25,6 +25,20 @@ const SWEEP_STREAMING_STALE_MS = 10 * 60 * 1000; // 10 minutes // into memory; far above any realistic transcript length. const FIND_ALL_BY_CHAT_LIMIT = 5000; +// Delta-poll overlap (#491): the poll query reaches this far BEHIND the client's +// echoed cursor, so a row that committed with an `updatedAt` marginally before the +// previous cursor was taken (on another autocommit connection) is still caught. +// Sized well above realistic single-row commit skew; the client merge is +// idempotent by id (mergeById), so the guaranteed repeats the overlap produces are +// harmless. +export const DELTA_POLL_OVERLAP_SECONDS = 5; + +// Hard cap on rows one delta poll returns — a safety bound (a poll should carry a +// handful of just-changed rows, never a whole transcript). Ordered by (updatedAt, +// id) asc, so on the pathological overflow the OLDEST changes win and the newest +// are picked up by the next poll (its cursor did not advance past them). +export const DELTA_POLL_MAX_ROWS = 500; + @Injectable() export class AiChatMessageRepo { private readonly logger = new Logger(AiChatMessageRepo.name); @@ -139,6 +153,72 @@ export class AiChatMessageRepo { .executeTakeFirst(); } + /** + * Delta read (#491) for the degraded poll: the chat's messages whose row + * changed AFTER the client's `cursor`, plus a FRESH cursor taken from the DB + * clock. Replaces the old "refetch ALL infinite-query pages every 2.5s with + * full parts" poll — the client seeds once (findByChat) and thereafter pulls + * only the deltas and merges them by id (mergeById). + * + * Cursor: a DB-clock timestamp (now()) the client echoes back each poll. All + * delta-relevant writes stamp `updatedAt` with now() (see `update` / + * `finalizeOwner`), so this is a SINGLE monotonic axis. The query overlaps the + * cursor by DELTA_POLL_OVERLAP_SECONDS to catch a row committed with an + * `updatedAt` marginally BEFORE the previous cursor was taken on another + * connection (single-row autocommit UPDATEs; no long transactions). The overlap + * GUARANTEES occasional REPEATS, so the client merge MUST be idempotent by id. + * + * `cursor === null` (first poll after the full seed) returns NO rows — there is + * nothing "new" relative to a just-loaded seed — only the fresh cursor to start + * the delta chain. The fresh cursor is read AFTER the rows, so it is >= every + * returned row's `updatedAt` (they were read strictly earlier) — a row that + * commits between the rows-read and the cursor-read is at most + * DELTA_POLL_OVERLAP_SECONDS behind the returned cursor, so the next poll's + * overlap window always re-includes it (no miss). + */ + async findByChatUpdatedAfter( + chatId: string, + workspaceId: string, + cursor: string | null, + ): Promise<{ rows: AiChatMessage[]; cursor: string }> { + if (cursor === null) { + const nowRow = await sql<{ now: Date }>`select now() as now`.execute( + this.db, + ); + return { rows: [], cursor: nowRow.rows[0].now.toISOString() }; + } + // Overlap the client cursor by DELTA_POLL_OVERLAP_SECONDS, computed in SQL off + // the echoed cursor so the whole comparison stays on the DB clock. + const rows = await this.db + .selectFrom('aiChatMessages') + .select(this.baseFields) + .where('chatId', '=', chatId) + .where('workspaceId', '=', workspaceId) + .where('deletedAt', 'is', null) + .where( + 'updatedAt', + '>', + sql`${cursor}::timestamptz - make_interval(secs => ${DELTA_POLL_OVERLAP_SECONDS})`, + ) + .orderBy('updatedAt', 'asc') + .orderBy('id', 'asc') + .limit(DELTA_POLL_MAX_ROWS) + .execute(); + // When the page filled (pathological overflow), DO NOT advance the cursor to + // now(): that would skip the changed rows past the cap that this poll did not + // return. Resume from the last returned row's updatedAt instead (the next + // poll's overlap re-includes ties by id). In the normal case the fresh DB-clock + // now() is the cursor. + if (rows.length === DELTA_POLL_MAX_ROWS) { + return { + rows, + cursor: rows[rows.length - 1].updatedAt.toISOString(), + }; + } + const nowRow = await sql<{ now: Date }>`select now() as now`.execute(this.db); + return { rows, cursor: nowRow.rows[0].now.toISOString() }; + } + async insert( insertable: InsertableAiChatMessage, trx?: KyselyTransaction, @@ -172,7 +252,13 @@ export class AiChatMessageRepo { const db = dbOrTx(this.db, opts?.trx); let query = db .updateTable('aiChatMessages') - .set({ ...(patch as Record), updatedAt: new Date() }) + // #491: stamp `updatedAt` from the DB clock (sql now()), NOT the app clock + // (new Date()). The delta-poll cursor (findByChatUpdatedAfter) is a single + // DB-clock axis; a per-step 'streaming' UPDATE stamped with the app clock + // would be a SECOND, skewed clock source and could leave a row's updatedAt + // just under a cursor taken from now() on another connection — an + // independent source of delta MISSES. All delta-relevant writes use now(). + .set({ ...(patch as Record), updatedAt: sql`now()` }) .where('id', '=', id) .where('workspaceId', '=', workspaceId); // Concurrency guard (#183 review): a per-step 'streaming' update must NEVER @@ -214,7 +300,9 @@ export class AiChatMessageRepo { const db = dbOrTx(this.db, trx); return db .updateTable('aiChatMessages') - .set({ ...(patch as Record), updatedAt: new Date() }) + // #491: DB-clock stamp (see `update`) — this terminal write flips the row's + // status, which the delta poll must observe on the shared now() cursor axis. + .set({ ...(patch as Record), updatedAt: sql`now()` }) .where('id', '=', id) .where('workspaceId', '=', workspaceId) .where((eb) => @@ -249,7 +337,9 @@ export class AiChatMessageRepo { .set({ status, metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`, - updatedAt: new Date(), + // #491: DB-clock stamp (see `update`) so a reconcile status flip lands on + // the same now() cursor axis the delta poll reads. + updatedAt: sql`now()`, }) .where('id', '=', id) .where('workspaceId', '=', workspaceId) @@ -307,7 +397,9 @@ export class AiChatMessageRepo { .set({ status: 'aborted', metadata: sql`coalesce(m.metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`, - updatedAt: new Date(), + // #491: DB-clock stamp (see `update`). The staleness WHERE below stays on + // the app clock — a >minutes window makes the ms-scale skew irrelevant. + updatedAt: sql`now()`, }) .where('m.status', '=', 'streaming') .where('m.updatedAt', '<', staleBefore) @@ -351,7 +443,8 @@ export class AiChatMessageRepo { .set({ status: 'aborted', metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`, - updatedAt: new Date(), + // #491: DB-clock stamp (see `update`). Staleness WHERE stays app-clock. + updatedAt: sql`now()`, }) .where('status', '=', 'streaming') .where('updatedAt', '<', staleBefore) diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.spec.ts b/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.spec.ts index d03a43e3..5d6faa9d 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.spec.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.spec.ts @@ -62,10 +62,17 @@ describe('AiChatRunRepo.sweepRunning', () => { // ...but a fresh 'running' run (updatedAt = now) must NOT be skipped: no // updatedAt predicate at all on the boot path. expect(rec.wheres.some(([col]) => col === 'updatedAt')).toBe(false); - // It flips to 'aborted' and stamps finishedAt. - expect(rec.set).toEqual( - expect.objectContaining({ status: 'aborted', finishedAt: expect.any(Date) }), - ); + // It flips to 'aborted' and stamps finishedAt + updatedAt. #491: the stamps + // are now DB-clock `sql now()` expressions (raw builders), NOT app-clock + // `new Date()`, so the run row shares the delta poll's single now() cursor axis + // — assert they are present and are the sql raw-builder objects (not a Date, + // not undefined). + expect(rec.set?.status).toBe('aborted'); + for (const stamp of ['finishedAt', 'updatedAt'] as const) { + expect(rec.set?.[stamp]).toBeDefined(); + expect(rec.set?.[stamp]).not.toBeInstanceOf(Date); + expect(typeof rec.set?.[stamp]).toBe('object'); + } }); it('phase-2 path: an explicit staleMs reintroduces the updatedAt window', async () => { diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.ts b/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.ts index c27480bf..5758867e 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.ts @@ -136,7 +136,11 @@ export class AiChatRunRepo { const db = dbOrTx(this.db, trx); return db .updateTable('aiChatRuns') - .set({ ...(patch as Record), updatedAt: new Date() }) + // #491: DB-clock stamp (sql now()) so the run row shares the delta poll's + // single now() cursor axis with the assistant message rows — a run-status + // change (the run fact the delta carries) must never sit on a skewed app + // clock relative to the message updatedAt cursor. + .set({ ...(patch as Record), updatedAt: sql`now()` }) .where('id', '=', id) .where('workspaceId', '=', workspaceId) .returning(this.baseFields) @@ -162,14 +166,15 @@ export class AiChatRunRepo { trx?: KyselyTransaction, ): Promise { const db = dbOrTx(this.db, trx); - const now = new Date(); return db .updateTable('aiChatRuns') .set({ status: patch.status, error: patch.error, - finishedAt: now, - updatedAt: now, + // #491: DB-clock stamps (finished_at + updated_at) so the terminal run + // fact lands on the delta poll's now() cursor axis. + finishedAt: sql`now()`, + updatedAt: sql`now()`, }) .where('id', '=', id) .where('workspaceId', '=', workspaceId) @@ -192,7 +197,8 @@ export class AiChatRunRepo { const db = dbOrTx(this.db, trx); return db .updateTable('aiChatRuns') - .set({ stopRequestedAt: new Date(), updatedAt: new Date() }) + // #491: DB-clock stamps (see `update`). + .set({ stopRequestedAt: sql`now()`, updatedAt: sql`now()` }) .where('id', '=', id) .where('workspaceId', '=', workspaceId) .where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[]) @@ -249,13 +255,14 @@ export class AiChatRunRepo { trx?: KyselyTransaction, ): Promise { const db = dbOrTx(this.db, trx); - const now = new Date(); let query = db .updateTable('aiChatRuns') .set({ status: 'aborted', - finishedAt: now, - updatedAt: now, + // #491: DB-clock stamps (see `update`). The staleness WHERE below stays on + // the app clock — a >minutes window makes the ms-scale skew irrelevant. + finishedAt: sql`now()`, + updatedAt: sql`now()`, error: sql`coalesce(error, ${'Run interrupted by a server restart.'})`, }) .where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[]); @@ -263,7 +270,7 @@ export class AiChatRunRepo { // sibling replica's live run is never aborted. Omitted on the phase-1 boot // sweep -> unconditional. if (typeof opts.staleMs === 'number') { - const staleBefore = new Date(now.getTime() - opts.staleMs); + const staleBefore = new Date(Date.now() - opts.staleMs); query = query.where('updatedAt', '<', staleBefore); } const rows = await query.returning('id').execute(); -- 2.52.0 From a919c79cb2587544b59861bf38b7d9e82940d800 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 18:02:01 +0300 Subject: [PATCH 2/7] =?UTF-8?q?feat(ai-chat):=20=D0=BF=D0=B5=D1=80=D1=81?= =?UTF-8?q?=D0=B8=D1=81=D1=82=20step-=D0=BC=D0=B0=D1=80=D0=BA=D0=B5=D1=80?= =?UTF-8?q?=D0=B0=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); + }); +}); -- 2.52.0 From daeeb1f3f2ce2f9807694532fe9dd9f00dd5f198 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 18:34:33 +0300 Subject: [PATCH 3/7] =?UTF-8?q?refactor(ai-chat):=20registry=20=E2=80=94?= =?UTF-8?q?=20step-aligned=20retention=20+=20tail-only=20attach=20(#491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Реестр ран-стримов больше не буферизует до 32МБ сырых SSE-кадров на активный ран и не выливает весь буфер в сокет синхронно при attach — это давало OOM на 1ГБ-контейнере при нескольких марафонских ранах. Теперь кольцо ограничено (env-настраиваемо, по умолчанию 4МБ) и держится в границах за счёт ротации по шагам. Серверная часть (суть коммита): - Штамповка кадров по шагам в ingestFrame. Штамп кадра = число `finish-step` кадров ДО него (с 0); сам finish-step несёт текущее значение, затем счётчик инкрементится. Так штамп совпадает с `metadata.stepsPersisted`: клиент с N персистнутыми шагами имеет 0..N-1 в сиде и просит хвост `stamp >= N`. Границу ловим дешёвым startsWith по `data: {"type":"finish-step"` — форма кадра проверена эмпирически против ai@6.0.207 (одна часть на кадр, type всегда первый ключ; кавычки в text-delta экранированы, ложных срабатываний нет). - Кольцо ротируется ТОЛЬКО на подтверждённом персисте шага N (`confirmPersistedStep`), сбрасывая кадры `stamp < N` (эти шаги уже на диске и придут в свежем сиде). `updateStreaming` теперь СИГНАЛИЗИРУЕТ исход (число персистнутых шагов или null), и ротация вызывается лишь при не-null возврате — провал персиста ничего не ротирует, кольцо покрывает БОЛЬШЕ (анти-инверсия: наивная ротация в .then() после НЕзаписанного шага дырявила бы гарантию). - Переполнение кольца сверх байтового капа вытесняет старейшие кадры; вытеснение ещё-не-персистнутого кадра открывает GAP. Гэп НЕ липкий: floor покрытия считается из кольца, поздний персист, проротировав дырявые шаги, его чистит. - attach(chatId, anchor, n): маркер шага N приходит ТОЛЬКО от клиента (сервер не читает строку — N из устаревшего сида дал бы тихую дыру в один шаг). Покрытие ОК ⟺ coverageFloor <= n; иначе 204 → клиент рефетчит (больший N) и переподключается. Хвост = синтетический `start`-кадр (ран-факт runId/chatId) + кадры `stamp >= n`. Инвариант 6 (нет кросс-ран реплея) сохранён через anchor; инвариант 4 (снапшот+регистрация в один синхронный тик) сохранён. N-срез применяется во ВСЕХ ветках, включая finished-retained: finished + N=N_final → пустой хвост + finish-кадр, клиент закрывает стрим. - Контроллер пишет хвост чанками с учётом drain (writeTailRespectingDrain), а не синхронным залпом (вторая половина OOM). Кап подписчика — производное 2× кап кольца, обе величины env-резолвятся на инстансе. Клиент: в тип строки добавлен `metadata.stepsPersisted` (источник N). PIN-SPEC трип-вайр на ai@6.0.207: `readUIMessageStream({ message })` продолжает последнее сообщение, `start`-кадр не сбрасывает parts, текст не пересекает finish-step — на этом держится продолжение при attach; апгрейд ai теперь падает громко. Тесты (observable-property против РЕАЛЬНОГО реестра/БД): детектор границы на реальной форме кадра, N-срез (в т.ч. посреди шага), ротация только на подтверждённом персисте, «персист провалился но кольцо влезло → attach успешен / провал + переполнение → 204», «устаревший N → 204 → после рефетча успех», очистка гэпа поздним персистом, finished-retained + N_final, memory-bound (5 параллельных марафонов сверх 32МБ, каждое кольцо ≤ кап). Обновлены registry/controller specs и DB-backed интеграционный attach-spec под новую сигнатуру/семантику. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../features/ai-chat/types/ai-chat.types.ts | 6 + .../utils/sdk-continuation-tripwire.test.ts | 83 +++ .../ai-chat-stream-registry.service.ts | 348 ++++++++++--- .../ai-chat/ai-chat-stream-registry.spec.ts | 483 ++++++++++++------ .../ai-chat/ai-chat.controller.attach.spec.ts | 22 +- .../src/core/ai-chat/ai-chat.controller.ts | 128 +++-- .../src/core/ai-chat/ai-chat.service.ts | 52 +- .../integration/ai-chat-attach.int-spec.ts | 71 ++- 8 files changed, 881 insertions(+), 312 deletions(-) create mode 100644 apps/client/src/features/ai-chat/utils/sdk-continuation-tripwire.test.ts diff --git a/apps/client/src/features/ai-chat/types/ai-chat.types.ts b/apps/client/src/features/ai-chat/types/ai-chat.types.ts index 2c987263..acb84de5 100644 --- a/apps/client/src/features/ai-chat/types/ai-chat.types.ts +++ b/apps/client/src/features/ai-chat/types/ai-chat.types.ts @@ -181,6 +181,12 @@ export interface IAiChatMessageRow { toolCalls?: unknown; metadata?: { parts?: UIMessage["parts"]; + // #491 step-alignment anchor: the count of FINISHED steps whose parts are in + // THIS row, written atomically with `parts` server-side (flushAssistant). The + // resume client reads it as its persisted step frontier N — the tail-only + // attach asks the run-stream registry for the frames of step N onward (the + // seed already carries steps 0..N-1). Absent on pre-#491 rows -> read as 0. + stepsPersisted?: number; // AI SDK v6 `totalUsage` persisted on assistant rows. Legacy cumulative // figure (sum of every step's usage for the turn); kept for back-compat and // as the fallback for older rows that have no `contextTokens`. diff --git a/apps/client/src/features/ai-chat/utils/sdk-continuation-tripwire.test.ts b/apps/client/src/features/ai-chat/utils/sdk-continuation-tripwire.test.ts new file mode 100644 index 00000000..c9fc19dc --- /dev/null +++ b/apps/client/src/features/ai-chat/utils/sdk-continuation-tripwire.test.ts @@ -0,0 +1,83 @@ +import { describe, it, expect } from "vitest"; +import { readUIMessageStream, type UIMessage } from "ai"; +import pkg from "../../../../package.json"; + +/** + * PIN-SPEC TRIP-WIRE (#491). The tail-only attach continuation relies on THREE + * behaviors of `ai@6.0.207`, verified line-by-line in the issue. Without this + * test, an `ai` bump could silently break attach (the client would append the + * live tail to the wrong message, or duplicate a step): + * + * 1. `readUIMessageStream({ message })` CONTINUES the passed message — it does + * not start a fresh one — so the tail streamed after a re-seed is appended to + * the seeded assistant row (the same DB id). + * 2. A `start` frame does NOT reset the existing message's parts (so the seeded + * steps 0..N-1 survive; the synthetic `start` the registry prepends only + * carries the run-fact metadata). + * 3. Text parts do NOT cross a `finish-step` boundary — a new `text-start` after + * `finish-step` is a NEW part — so the reconstructed steps stay separated and + * the step frontier stays meaningful. + * + * If an `ai` upgrade changes any of these, this test fails LOUD instead of the + * resume path silently corrupting. + */ +describe("ai SDK continuation trip-wire (#491, tail-only attach)", () => { + it("is pinned to the exact ai version the continuation was verified against", () => { + // A caret/range bump is exactly what would silently break attach — require an + // exact pin. Bumping ai MUST re-verify the behavior asserted below, then this. + expect((pkg as { dependencies: Record }).dependencies.ai).toBe( + "6.0.207", + ); + }); + + it("continues the seeded message: start does not reset parts, the tail appends as new parts", async () => { + // A seeded assistant row with ONE finished step already reconstructed. + const seeded: UIMessage = { + id: "assistant-1", + role: "assistant", + parts: [ + { type: "step-start" }, + { type: "text", text: "STEP0", state: "done" }, + ], + } as UIMessage; + + // The tail the registry delivers on re-attach: a synthetic start (run-fact), + // then step 1's frames, then finish. As UI-message chunks (what the SSE frames + // decode to). + const chunks = [ + { type: "start", messageMetadata: { runId: "r1", chatId: "c1" } }, + { type: "start-step" }, + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: "STEP1" }, + { type: "text-end", id: "t1" }, + { type: "finish-step" }, + { type: "finish" }, + ]; + const stream = new ReadableStream({ + start(c) { + for (const ch of chunks) c.enqueue(ch); + c.close(); + }, + }); + + let last: UIMessage | undefined; + for await (const msg of readUIMessageStream({ message: seeded, stream })) { + last = msg; + } + + expect(last).toBeDefined(); + // Same message id (continuation, not a fresh message). + expect(last!.id).toBe("assistant-1"); + // The seeded step-0 parts SURVIVED the `start` frame, and step 1 was appended + // as SEPARATE parts (text did not cross the finish-step boundary). + const shape = last!.parts.map((p) => `${p.type}:${(p as { text?: string }).text ?? ""}`); + expect(shape).toEqual([ + "step-start:", + "text:STEP0", + "step-start:", + "text:STEP1", + ]); + // The run-fact metadata from the synthetic start frame is applied. + expect(last!.metadata).toMatchObject({ runId: "r1", chatId: "c1" }); + }); +}); diff --git a/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts index 1ef476f5..5fb3c15a 100644 --- a/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts @@ -1,49 +1,132 @@ import { Injectable, Logger, OnModuleDestroy } from '@nestjs/common'; /** - * In-memory run-stream registry (#184 phase 1.5). A durable agent run tees its - * SSE frames here (via `pipeUIMessageStreamToResponse({ consumeSseStream })`) - * so a LATE tab — one that reloaded, or opened after the starter dropped — can - * attach through `GET /ai-chat/runs/:chatId/stream`, replay the frames buffered - * so far, and then follow the live tail as a normal streamer. + * In-memory run-stream registry (#184 phase 1.5, step-aligned retention #491). A + * durable agent run tees its SSE frames here (via + * `pipeUIMessageStreamToResponse({ consumeSseStream })`) so a LATE tab — one that + * reloaded, or opened after the starter dropped — can attach through + * `GET /ai-chat/runs/:chatId/stream`, be handed the TAIL past the step it already + * has persisted, and then follow the live tail as a normal streamer. * * This is deliberately single-process and best-effort: it holds nothing the DB * does not (the run + assistant row are the source of truth), so a process * restart simply drops in-flight entries and the client falls back to its * restore + degraded-poll path. The async `attach` return type is the seam for a * future phase-2 cross-process backend (Redis) — the interface does not change. + * + * ── #491 step-aligned retention (the OOM fix) ──────────────────────────────── + * The old registry buffered up to 32MB of raw SSE frames PER active run (V8 ~2× + * in memory) and, on attach, blasted the WHOLE buffer to the socket synchronously + * with no drain — a handful of marathon runs on a 1GB container OOM'd. #491 caps + * the ring at a few MB (env-tunable, default 4MB) and keeps it there by ROTATING: + * + * - Every buffered frame is STAMPED with a step number at tee (see ingestFrame). + * Convention: the stamp of a frame is the number of `finish-step` parts seen + * BEFORE it (starting at 0). The finish-step frame itself carries the current + * value, THEN the counter increments. So a frame stamped `s` is the content of + * the (s+1)-th step — 0-based step index `s` — and the stamp aligns EXACTLY + * with `metadata.stepsPersisted`: a client whose persisted `stepsPersisted` is + * N has steps 0..N-1 on disk (and in its seed) and needs the tail `stamp >= N`. + * + * - The ring rotates ONLY on a CONFIRMED persist of step N + * (`confirmPersistedStep`), dropping frames with `stamp < N` (those steps are + * now on disk and a fresh client seed carries them). A NON-confirmed step is + * never rotated away, so a persist FAILURE just makes the ring cover MORE + * (auto-safe). This is the anti-inversion rule: a naive "rotate in .then()" + * that rotated after an UNwritten step would drop a step nobody has → silent + * hole. Rotation is gated on a real, successful persist. + * + * - If the ring still exceeds its byte cap after rotation (a single fat step, or + * a lagging persist), the OLDEST frames are evicted to stay bounded. Evicting a + * not-yet-persisted frame opens a GAP: an attach whose N falls at or below an + * evicted step answers 204 and the client degrades to restore+poll. The gap is + * NOT sticky — the coverage floor is recomputed from the ring, so a later + * persist that rotates past the holey steps clears it. + * + * ── attach numbering / coverage (the wire convention) ──────────────────────── + * The step marker N comes ONLY FROM THE CLIENT (a query param). The server never + * reads the row to derive N — a server-side N from a stale seed would open a + * silent one-step hole. N is the client's persisted `stepsPersisted` (a COUNT): + * - the tail it needs = frames with `stamp >= N`; + * - coverage is OK ⟺ `coverageFloor(entry) <= N`, where coverageFloor is the + * smallest step FULLY present in the ring (its smallest retained stamp, bumped + * by one when that leading step was only partially evicted by overflow). If + * `coverageFloor > N` the ring starts AFTER the client's frontier (a hole, or + * the client's seed simply lagged behind a rotation) → 204 → the client + * refetches (a larger N) and re-attaches. + * The N cutoff is applied in ALL branches, INCLUDING the finished-retained replay. + * + * ── same-tick invariants (unchanged, still load-bearing) ───────────────────── + * invariant 1: only the matching run may mutate/observe an entry (runId check). + * invariant 2: retention deletes ONLY its own entry (a replacement may own the key). + * invariant 3: open() over a live entry mirrors the done-path (subscribers released). + * invariant 4: the tail SLICE + subscriber registration happen in ONE synchronous + * tick inside attach() — no await between them — so a concurrently + * ingested frame is EITHER in the snapshot (buffered before the sync + * block, and the just-added subscriber never sees it) OR fanned out to + * the paused subscriber's `pending` (ingested after) — never both and + * never neither: no loss, no duplication. NOTE (#491): the controller + * now AWAITS the drain-respecting tail write BEFORE calling start(), so + * frames ingested during that await accumulate in `pending`; this is + * bounded by the subscriber cap (an overflow degrades start() to an + * end(), a 204-equivalent). It is the SYNCHRONOUS snapshot+registration + * — not a same-tick start() — that makes this correct. + * invariant 5: the controller wires close-cleanup BEFORE any write. + * invariant 6: no cross-run replay — the `anchor` (the client's assistant row id) + * must match this run's assistant id, or a foreign run's transcript + * would be appended to the client's message. */ /** How long a finished entry is retained for late attach (replay + immediate end). */ export const RUN_STREAM_RETAIN_FINISHED_MS = 30_000; /** - * Per-run replay buffer cap. Past this the buffer is dropped (attach -> 204, and - * the client falls back to its restore + degraded-poll path, #430). - * - * Raised from 4MB to 32MB (#430): marathon autonomous runs (11-25 min observed) - * stream far more than 4MB of SSE frames, so a live disconnect mid-run would find - * an already-overflowed buffer and could only degrade-poll instead of re-attaching - * to the live tail. 32MB comfortably covers those runs while staying bounded. - * - * Memory cost: this is the WORST-CASE retained size PER ACTIVE run (the buffer is - * freed on finish + retention, or dropped immediately on overflow). With the small - * number of concurrent autonomous runs a single workspace realistically has, 32MB - * each is an acceptable ceiling; the overflow->204->degraded-poll fallback remains - * the backstop for anything larger, so correctness never depends on this bound. + * DEFAULT per-run replay ring cap (#491, down from 32MB). SSE frames carry + * UNcompacted tool outputs + framing overhead (×1.5–2 vs the persisted parts), so + * a "2–3 large reads + reasoning" step routinely blows past 2MB; 4MB comfortably + * holds a step or two of TAIL, which is all a resuming client needs (steps below + * its persisted frontier come from the seed, not the ring). The ring stays bounded + * because it rotates on every confirmed persist; this cap is only the ceiling for + * the un-persisted tail between rotations. Env-tunable via + * RUN_STREAM_MAX_BUFFER_BYTES (bytes); a 0/invalid value falls back to this. */ -export const RUN_STREAM_MAX_BUFFER_BYTES = 32 * 1024 * 1024; +export const RUN_STREAM_MAX_BUFFER_BYTES = 4 * 1024 * 1024; -// 2x the replay cap: a just-written full-replay burst alone can never trip the -// per-subscriber cap (see controller); only a genuinely stalled socket can. +// 2× the ring cap: a just-written full-tail burst alone can never trip the +// per-subscriber cap (see controller); only a genuinely stalled socket can. This +// derivative relationship is preserved even when the ring cap is env-overridden. export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * RUN_STREAM_MAX_BUFFER_BYTES; +/** + * A finish-step boundary frame is exactly `data: {"type":"finish-step"...}\n\n` + * (verified empirically against ai@6.0.207 — each UI-message-stream part is a + * single `data: {json}\n\n` event, never split across `data:` lines, and `type` + * is always the first key). A prefix match is cheaper than JSON.parse-per-frame + * and has no false positives: a literal `"type":"finish-step"` inside a text + * delta is JSON-escaped (`\"type\":...`), and the frame would start with + * `data: {"type":"text-delta"` anyway. + */ +const FINISH_STEP_FRAME_PREFIX = 'data: {"type":"finish-step"'; + +/** Resolve the ring cap from the environment, falling back to the default. */ +function resolveMaxBufferBytes(): number { + const raw = process.env.RUN_STREAM_MAX_BUFFER_BYTES; + if (!raw) return RUN_STREAM_MAX_BUFFER_BYTES; + const parsed = Number(raw); + return Number.isFinite(parsed) && parsed > 0 + ? Math.floor(parsed) + : RUN_STREAM_MAX_BUFFER_BYTES; +} + export interface RunStreamCallbacks { onFrame: (frame: string) => void; onEnd: () => void; } export interface RunStreamAttachment { + // The synthetic `start` frame (carrying { runId, chatId }) followed by the + // buffered TAIL filtered to `stamp >= N`. The controller writes these to the + // socket in chunks respecting drain, then calls start(). replay: string[]; finished: boolean; start(): void; // drain pending frames (order preserved) and go live @@ -53,14 +136,19 @@ export interface RunStreamAttachment { interface Subscriber extends RunStreamCallbacks { started: boolean; pending: string[]; - // Byte size of `pending`, capped at SUBSCRIBER_MAX_BUFFERED_BYTES. `start()` is - // called in the SAME tick as `attach()` today (see attach), so `pending` never - // holds more than one microtask of frames — but the async `attach` signature is - // a phase-2 seam: an await between attach and start would let a stalled paused - // subscriber buffer the WHOLE run here. The cap is the structural backstop. + // Byte size of `pending`, capped at the subscriber cap. `start()` is called in + // the SAME tick as `attach()` today, so `pending` never holds more than one + // microtask of frames — but the controller writes the (potentially large) tail + // respecting drain BEFORE start(), so a stalled socket can accumulate here; the + // cap is the structural backstop (an overflow degrades start() to an end()). pendingBytes: number; overflowed: boolean; pendingEnd: boolean; + // The client's step frontier N: this subscriber only receives frames with + // `stamp >= minStamp` (the tail past what it already persisted). Live frames + // always satisfy this (their stamp is the current, highest step), so it only + // filters the rare out-of-order below-frontier frame. + minStamp: number; } interface Entry { @@ -68,8 +156,20 @@ interface Entry { // The persisted assistant row id of this run (set at bind; undefined if the // seed failed). Used by the attach anchor check (invariant 6). assistantMessageId?: string; + // Parallel arrays: frames[i] is the SSE string, stamps[i] its step number. frames: string[]; + stamps: number[]; bytes: number; + // The running step counter used to stamp the NEXT frame (number of finish-step + // frames seen so far). + currentStamp: number; + // The highest confirmed `stepsPersisted`: frames with stamp < persistedFloor are + // on disk (safe to drop, never re-buffered). Monotonic (confirmPersistedStep). + persistedFloor: number; + // The highest stamp EVICTED by an overflow (unsafe) drop, -1 if none. Used to + // detect a partially-evicted leading step when computing the coverage floor. + overflowThroughStamp: number; + // Sticky-for-logging only: at least one unsafe (overflow) eviction happened. overflowed: boolean; finished: boolean; subscribers: Set; @@ -80,6 +180,10 @@ interface Entry { export class AiChatStreamRegistryService implements OnModuleDestroy { private readonly logger = new Logger(AiChatStreamRegistryService.name); private readonly entries = new Map(); // key: chatId + // Env-resolved caps (per instance) so a deployment can tune the ceiling without + // a code change. The subscriber cap keeps the documented 2× relationship. + readonly maxBufferBytes = resolveMaxBufferBytes(); + readonly subscriberMaxBufferedBytes = 2 * this.maxBufferBytes; /** * Register a fresh entry at the START of a run (before any frame), so a tab @@ -105,7 +209,11 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { this.entries.set(chatId, { runId, frames: [], + stamps: [], bytes: 0, + currentStamp: 0, + persistedFloor: 0, + overflowThroughStamp: -1, overflowed: false, finished: false, subscribers: new Set(), @@ -150,6 +258,34 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { void pump(); } + /** + * Confirm that step `stepsPersisted` (a COUNT: steps 0..stepsPersisted-1) is on + * disk for this run, and ROTATE the ring: drop the buffered frames of those + * now-persisted steps (stamp < stepsPersisted). This is the ONLY thing that + * rotates the ring, and it is called ONLY after a genuinely SUCCESSFUL per-step + * persist (see ai-chat.service updateStreaming). A failed persist never calls + * it, so the ring covers more (auto-safe). Identity-checked (invariant 1) and + * monotonic (a stale lower count is ignored). + */ + confirmPersistedStep( + chatId: string, + runId: string, + stepsPersisted: number, + ): void { + const entry = this.entries.get(chatId); + if (!entry || entry.runId !== runId) return; + if (!Number.isFinite(stepsPersisted) || stepsPersisted <= entry.persistedFloor) + return; + entry.persistedFloor = stepsPersisted; + // Clean rotation: drop the persisted steps from the head. These frames are on + // disk + carried by a fresh client seed, so this NEVER opens a gap. + while (entry.frames.length > 0 && entry.stamps[0] < stepsPersisted) { + entry.bytes -= Buffer.byteLength(entry.frames[0]); + entry.frames.shift(); + entry.stamps.shift(); + } + } + /** * Terminate a run's entry from the OUTER catch of the stream method (a failure * before/while wiring the pipe, so `done` will never arrive). Identity-checked @@ -162,36 +298,62 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { } /** - * Attach to a run's stream. Async only for the phase-2 Redis seam — the body - * runs synchronously so the replay snapshot and the subscriber registration - * happen in ONE tick with no await between them (invariant 4): a frame ingested - * concurrently cannot slip into the gap and be lost or duplicated. + * Attach to a run's stream from the client's step frontier `n` (its persisted + * `stepsPersisted`). Async only for the phase-2 Redis seam — the body runs + * synchronously so the tail SLICE and the subscriber registration happen in ONE + * tick with no await between them (invariant 4). * * Returns null (-> the caller answers 204) when: - * - there is no entry, or it overflowed (replay is gone); - * - expect=live with an anchor that does not match this run's assistant id - * (invariant 6: a stripped tab must never replay a FOREIGN run's transcript); - * - the run finished and the caller did not expect a live tail. - * A finished run with expect=live yields a replay-only attachment (no - * subscriber registered). Otherwise a paused subscriber is registered and the - * caller replays `replay`, then calls start() to drain and go live. + * - there is no entry; + * - the `anchor` does not match this run's assistant id (invariant 6); + * - the ring does not cover the client's frontier (coverageFloor > n): a hole + * from overflow, or the client's seed simply lagged behind a rotation. The + * client then refetches (a larger n) and re-attaches. + * + * Otherwise the attachment's `replay` is a synthetic `start` frame (the run-fact + * on re-attach) followed by the buffered tail filtered to `stamp >= n`. For a + * FINISHED run this is replay-only (no subscriber) and ends after the replay — + * with n = N_final that tail is just the run's `finish` frame, so the client + * closes the stream. For a LIVE run a paused subscriber is registered; the + * caller writes the replay (respecting drain) then calls start() to drain the + * pending frames and go live. */ async attach( chatId: string, - expectLive: boolean, anchor: string | undefined, + n: number, cb: RunStreamCallbacks, ): Promise { const entry = this.entries.get(chatId); - if (!entry || entry.overflowed) return null; + if (!entry) return null; // Invariant 6: cross-run replay is forbidden. Before bind, assistantMessageId // is undefined and mismatches any anchor -> 204 -> client restore+poll path. - if (expectLive && anchor && entry.assistantMessageId !== anchor) return null; - if (entry.finished && !expectLive) return null; - if (entry.finished && expectLive) { + if (anchor && entry.assistantMessageId !== anchor) return null; + // A finished entry with NOTHING in the ring (aborted before the first frame, + // or fully overflowed) has no tail to deliver -> 204 -> the client polls. + if (entry.finished && entry.frames.length === 0) return null; + const floor = this.coverageFloor(entry); + if (floor > n) { + this.logger.warn( + `run-stream attach gap for run=${entry.runId}: coverageFloor=${floor} ` + + `> client n=${n} -> 204 (client refetches + re-attaches)`, + ); + return null; + } + + const startFrame = this.buildStartFrame(chatId, entry.runId); + const sliceTail = (): string[] => { + const out: string[] = [startFrame]; + for (let i = 0; i < entry.frames.length; i++) { + if (entry.stamps[i] >= n) out.push(entry.frames[i]); + } + return out; + }; + + if (entry.finished) { // Replay-only: the run is done, no subscriber is registered. return { - replay: entry.frames.slice(), + replay: sliceTail(), finished: true, start: () => undefined, unsubscribe: () => undefined, @@ -206,15 +368,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { pendingBytes: 0, overflowed: false, pendingEnd: false, + minStamp: n, }; + // Register + snapshot in the SAME synchronous block (invariant 4). No await + // separates them, so a concurrently ingested frame cannot be lost/duplicated. entry.subscribers.add(sub); - // Snapshot in the SAME synchronous block as the registration (invariant 4). - const replay = entry.frames.slice(); - // CONTRACT: the caller MUST call start() in the SAME tick as this attach() - // returns — no await between them. While a subscriber is paused, every frame - // is buffered in sub.pending; a delayed start() lets a whole run accumulate - // there. The pendingBytes cap (see ingestFrame) is the structural backstop if - // that contract is ever broken (e.g. the phase-2 Redis await seam). + const replay = sliceTail(); return { replay, finished: false, @@ -263,24 +422,77 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { this.entries.clear(); } - /** Buffer + fan-out a single frame. See invariant/overflow semantics inline. */ + /** The synthetic `start` frame the tail is prefixed with — the source of the + * run-fact (runId/chatId) on re-attach. A `start` frame does NOT reset the + * client's message parts (ai@6.0.207 createStreamingUIMessageState), so it is + * safe to prepend even when the sliced tail begins mid-message. */ + private buildStartFrame(chatId: string, runId: string): string { + return `data: ${JSON.stringify({ + type: 'start', + messageMetadata: { runId, chatId }, + })}\n\n`; + } + + /** + * The smallest step FULLY present in the ring: its smallest retained stamp, or + * (when the leading step was only partially evicted by an overflow) one past it. + * When the ring is empty it is the current step (only the live tail is coming). + * An attach at frontier `n` is covered ⟺ coverageFloor <= n. + */ + private coverageFloor(entry: Entry): number { + if (entry.frames.length === 0) return entry.currentStamp; + const min = entry.stamps[0]; + return entry.overflowThroughStamp >= min ? min + 1 : min; + } + + /** + * Buffer (step-stamped) + fan-out a single frame. The stamp is the number of + * finish-step frames seen BEFORE this one; a finish-step frame carries the + * current value and THEN increments the counter (so its stamp equals the 0-based + * index of the step it closes). Only frames at/above persistedFloor are buffered + * (already-persisted steps are on disk); the ring is then trimmed to the byte + * cap, an unsafe eviction opening a gap. Fan-out is always live (filtered per + * subscriber by its frontier). + */ private ingestFrame(entry: Entry, frame: string): void { - entry.bytes += Buffer.byteLength(frame); - if (!entry.overflowed) { + const size = Buffer.byteLength(frame); + const stamp = entry.currentStamp; + if (frame.startsWith(FINISH_STEP_FRAME_PREFIX)) { + entry.currentStamp = stamp + 1; + } + + // Buffer for replay only if this step is not already persisted+rotated away. + if (stamp >= entry.persistedFloor) { entry.frames.push(frame); - if (entry.bytes > RUN_STREAM_MAX_BUFFER_BYTES) { - // The crossing frame was already counted AND (below) fanned out; only the - // replay buffer is dropped. After overflow no more frames are buffered, - // but live fan-out continues. - entry.overflowed = true; - entry.frames = []; - this.logger.warn( - `run-stream buffer overflow for run=${entry.runId}; ` + - `late attach will 204 until the run ends`, - ); + entry.stamps.push(stamp); + entry.bytes += size; + // Enforce the ring cap. Evicting a not-yet-persisted frame (stamp >= + // persistedFloor) opens a GAP; a leftover persisted frame (< floor) is a + // safe drop. Keep evicting until the ring is back under the cap. + while (entry.bytes > this.maxBufferBytes && entry.frames.length > 0) { + const evStamp = entry.stamps[0]; + entry.bytes -= Buffer.byteLength(entry.frames[0]); + entry.frames.shift(); + entry.stamps.shift(); + if (evStamp >= entry.persistedFloor) { + if (evStamp > entry.overflowThroughStamp) + entry.overflowThroughStamp = evStamp; + if (!entry.overflowed) { + entry.overflowed = true; + this.logger.warn( + `run-stream ring overflow for run=${entry.runId}: an un-persisted ` + + `step was evicted to stay under ${this.maxBufferBytes}B; a late ` + + `attach at an evicted step will 204 until a later persist confirms`, + ); + } + } } } + + // Fan out live, filtered to each subscriber's frontier (a subscriber only + // wants the tail past the step it already persisted). for (const sub of entry.subscribers) { + if (stamp < sub.minStamp) continue; if (sub.started) { try { sub.onFrame(frame); @@ -289,12 +501,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { } } else { sub.pending.push(frame); - sub.pendingBytes += Buffer.byteLength(frame); - if (sub.pendingBytes > SUBSCRIBER_MAX_BUFFERED_BYTES) { + sub.pendingBytes += size; + if (sub.pendingBytes > this.subscriberMaxBufferedBytes) { // The paused subscriber's buffer overflowed — only possible if start() - // was delayed past the same-tick contract (the phase-2 await seam). - // Drop it rather than buffer the whole run; on start() it degrades to an - // immediate end (a 204-equivalent) instead of replaying a partial. + // was delayed (the controller's drain-respecting tail write, or the + // phase-2 await seam). Drop it rather than buffer the whole run; on + // start() it degrades to an immediate end (a 204-equivalent). sub.overflowed = true; sub.pending = []; entry.subscribers.delete(sub); diff --git a/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts b/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts index 193b76b1..06df6d75 100644 --- a/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts @@ -2,18 +2,26 @@ import { AiChatStreamRegistryService, RUN_STREAM_MAX_BUFFER_BYTES, RUN_STREAM_RETAIN_FINISHED_MS, - SUBSCRIBER_MAX_BUFFERED_BYTES, RunStreamCallbacks, } from './ai-chat-stream-registry.service'; /** - * Unit tests for the in-memory run-stream registry (#184 phase 1.5). The registry - * is the whole of the resumable-transport contract: replay ordering, paused -> - * live hand-off, overflow, retention, the anchor check (invariant 6), and the - * mirror-the-done-path replace semantics (invariant 3). Every enumerated case in - * the issue's task 1.5 has a test here. + * Unit tests for the in-memory run-stream registry (#184 phase 1.5, step-aligned + * retention #491). The registry is the whole of the resumable-transport contract: + * step-stamped retention, tail-only attach at the client's frontier N, the + * confirmed-persist ring rotation (and the anti-inversion rule), the memory bound, + * the overflow gap, paused -> live hand-off, retention, the anchor check + * (invariant 6), and the mirror-the-done-path replace semantics (invariant 3). */ +// Real ai@6 UI-message-stream SSE frames are `data: {json}\n\n`, one part each. +const sse = (part: Record): string => + `data: ${JSON.stringify(part)}\n\n`; +const finishStep = (): string => sse({ type: 'finish-step' }); +const textDelta = (id: string, delta: string): string => + sse({ type: 'text-delta', id, delta }); +const finish = (): string => sse({ type: 'finish' }); + // A ReadableStream whose frames the test pushes explicitly, plus close/error. function makePushStream(): { stream: ReadableStream; @@ -58,6 +66,9 @@ function collector(): { }; } +// The tail past the synthetic start frame (replay[0] is always the start frame). +const tail = (replay: string[]): string[] => replay.slice(1); + describe('AiChatStreamRegistryService', () => { const CHAT = 'chat-1'; let registry: AiChatStreamRegistryService; @@ -71,7 +82,21 @@ describe('AiChatStreamRegistryService', () => { registry.onModuleDestroy(); }); - it('replays frames in arrival order (live attach)', async () => { + it('prepends a synthetic start frame carrying { runId, chatId }', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push('a'); + await flush(); + + const c = collector(); + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; + const start = JSON.parse(att.replay[0].replace(/^data: /, '').trim()); + expect(start.type).toBe('start'); + expect(start.messageMetadata).toEqual({ runId: 'run-1', chatId: CHAT }); + }); + + it('replays the buffered tail (from frontier 0) in arrival order (live attach)', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); @@ -81,13 +106,13 @@ describe('AiChatStreamRegistryService', () => { await flush(); const c = collector(); - const att = await registry.attach(CHAT, false, undefined, c.cb); + const att = await registry.attach(CHAT, 'assist-1', 0, c.cb); expect(att).not.toBeNull(); - expect(att!.replay).toEqual(['a', 'b', 'c']); + expect(tail(att!.replay)).toEqual(['a', 'b', 'c']); expect(att!.finished).toBe(false); }); - it('late attach gets the full prefix as replay plus the live tail', async () => { + it('late attach gets the buffered prefix as tail plus the live tail', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); @@ -96,17 +121,16 @@ describe('AiChatStreamRegistryService', () => { await flush(); const c = collector(); - const att = (await registry.attach(CHAT, false, undefined, c.cb))!; - expect(att.replay).toEqual(['a', 'b']); + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; + expect(tail(att.replay)).toEqual(['a', 'b']); att.start(); - // Live tail arrives after start(). src.push('c'); src.push('d'); await flush(); expect(c.frames).toEqual(['c', 'd']); }); - it('a paused subscriber receives frames buffered during pause in order, then live (no loss/reorder)', async () => { + it('a paused subscriber receives frames buffered during pause in order, then live', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); @@ -114,81 +138,45 @@ describe('AiChatStreamRegistryService', () => { await flush(); const c = collector(); - // Attach (paused). Frames that arrive BEFORE start() must queue, not drop. - const att = (await registry.attach(CHAT, false, undefined, c.cb))!; - expect(att.replay).toEqual(['a']); + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; + expect(tail(att.replay)).toEqual(['a']); src.push('b'); // arrives while paused -> pending src.push('c'); await flush(); expect(c.frames).toEqual([]); // nothing delivered yet (paused) - att.start(); // drains pending in order + att.start(); expect(c.frames).toEqual(['b', 'c']); - src.push('d'); // now live + src.push('d'); await flush(); expect(c.frames).toEqual(['b', 'c', 'd']); }); it('a run that finishes while a subscriber is paused ends it on start()', async () => { registry.open(CHAT, 'run-1'); + registry.bind(CHAT, 'run-1', 'assist-1', makePushStream().stream); const c = collector(); - const att = (await registry.attach(CHAT, false, undefined, c.cb))!; - // Terminate the run while the subscriber is still paused. + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; registry.abortEntry(CHAT, 'run-1'); expect(c.ended()).toBe(0); // paused: not ended yet att.start(); expect(c.ended()).toBe(1); // start() drains + ends }); - it('finished + expect=live returns a replay WITHOUT registering a subscriber', async () => { - registry.open(CHAT, 'run-1'); - const src = makePushStream(); - registry.bind(CHAT, 'run-1', 'assist-1', src.stream); - src.push('a'); - src.push('b'); - src.close(); - await flush(); - - const c = collector(); - const att = (await registry.attach(CHAT, true, undefined, c.cb))!; - expect(att.finished).toBe(true); - expect(att.replay).toEqual(['a', 'b']); - // No subscriber registered: start()/unsubscribe are no-ops and the entry has - // zero subscribers. - const entry = (registry as any).entries.get(CHAT); - expect(entry.subscribers.size).toBe(0); - att.start(); - expect(c.frames).toEqual([]); - }); - - it('finished WITHOUT expect=live returns null', async () => { - registry.open(CHAT, 'run-1'); - const src = makePushStream(); - registry.bind(CHAT, 'run-1', 'assist-1', src.stream); - src.push('a'); - src.close(); - await flush(); - - const c = collector(); - expect(await registry.attach(CHAT, false, undefined, c.cb)).toBeNull(); - }); - - it('anchor mismatch with expect=live returns null (and null before bind sets assistantMessageId)', async () => { + it('anchor mismatch returns null (and null before bind sets assistantMessageId)', async () => { registry.open(CHAT, 'run-1'); const c = collector(); // Before bind: assistantMessageId is undefined -> mismatches any anchor. - expect( - await registry.attach(CHAT, true, 'assist-1', c.cb), - ).toBeNull(); + expect(await registry.attach(CHAT, 'assist-1', 0, c.cb)).toBeNull(); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); src.push('a'); await flush(); // Wrong anchor -> null (cross-run replay forbidden, invariant 6). - expect(await registry.attach(CHAT, true, 'other-id', c.cb)).toBeNull(); + expect(await registry.attach(CHAT, 'other-id', 0, c.cb)).toBeNull(); }); - it('matching anchor with expect=live attaches', async () => { + it('matching anchor attaches', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); @@ -196,97 +184,60 @@ describe('AiChatStreamRegistryService', () => { await flush(); const c = collector(); - const att = await registry.attach(CHAT, true, 'assist-1', c.cb); + const att = await registry.attach(CHAT, 'assist-1', 0, c.cb); expect(att).not.toBeNull(); - expect(att!.replay).toEqual(['a']); + expect(tail(att!.replay)).toEqual(['a']); }); - it('overflow: attach returns null, but the LIVE subscriber keeps receiving (incl. the crossing frame)', async () => { + it('a throwing onFrame ejects only that subscriber; the ingest loop stays alive', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); - // A live (started) subscriber attached before the flood. + const bad = collector(); + const badAtt = (await registry.attach(CHAT, 'assist-1', 0, { + onFrame: () => { + throw new Error('boom'); + }, + onEnd: bad.cb.onEnd, + }))!; + badAtt.start(); + + const good = collector(); + const goodAtt = (await registry.attach(CHAT, 'assist-1', 0, good.cb))!; + goodAtt.start(); + + src.push('a'); // bad throws on this frame -> ejected + src.push('b'); // good still receives both + await flush(); + + const entry = (registry as any).entries.get(CHAT); + expect(entry.subscribers.size).toBe(1); + expect(good.frames).toEqual(['a', 'b']); + }); + + it('open() over a LIVE entry ends started subscribers once; a late done never touches the new entry (invariant 3)', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push('a'); + await flush(); + const c = collector(); - const att = (await registry.attach(CHAT, false, undefined, c.cb))!; + const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!; att.start(); - // Cap-relative so it survives a buffer-cap change (#430): a quarter-cap frame - // means 5 frames comfortably exceed the replay cap; the last one crosses. - const chunk = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 4)); - for (let i = 0; i < 5; i++) src.push(chunk + i); - await flush(); - - const entry = (registry as any).entries.get(CHAT); - expect(entry.overflowed).toBe(true); - expect(entry.bytes).toBeGreaterThan(RUN_STREAM_MAX_BUFFER_BYTES); - // The live subscriber received ALL 5 frames, including the crossing one. - expect(c.frames).toHaveLength(5); - expect(c.frames[4]).toBe(chunk + 4); - - // A NEW attach after overflow gets null (replay buffer is gone). - const c2 = collector(); - expect(await registry.attach(CHAT, false, undefined, c2.cb)).toBeNull(); - }); - - it('a paused subscriber whose pending buffer overflows is dropped and ends on start(); other subscribers keep receiving', async () => { - registry.open(CHAT, 'run-1'); - const src = makePushStream(); - registry.bind(CHAT, 'run-1', 'assist-1', src.stream); - - // A: paused (start() deliberately delayed to simulate the phase-2 await seam). - const a = collector(); - const attA = (await registry.attach(CHAT, false, undefined, a.cb))!; - // B: live (started) — its delivery must be unaffected by A's overflow. - const b = collector(); - const attB = (await registry.attach(CHAT, false, undefined, b.cb))!; - attB.start(); - - // Cap-relative so it survives a buffer-cap change (#430): a quarter-of-the- - // per-subscriber-cap frame means 5 frames exceed A's paused-pending cap while - // B streams every frame live. - const chunk = 'x'.repeat(Math.floor(SUBSCRIBER_MAX_BUFFERED_BYTES / 4)); - for (let i = 0; i < 5; i++) src.push(chunk + i); - await flush(); - - const entry = (registry as any).entries.get(CHAT); - // A was dropped from the subscriber set on overflow; B (started) remains. - expect(entry.subscribers.size).toBe(1); - expect(a.frames).toEqual([]); // paused + overflowed: nothing was delivered - // B received every frame live (delivery unaffected by A's overflow). - expect(b.frames).toHaveLength(5); - - // A's start() (arriving late) degrades to an immediate end, not a partial replay. - attA.start(); - expect(a.frames).toEqual([]); - expect(a.ended()).toBe(1); - }); - - it('open() over a LIVE entry ends started subscribers exactly once and a late done does not touch the new entry (invariant 3)', async () => { - registry.open(CHAT, 'run-1'); - const src = makePushStream(); - registry.bind(CHAT, 'run-1', 'assist-1', src.stream); - src.push('a'); - await flush(); - - const c = collector(); - const att = (await registry.attach(CHAT, false, undefined, c.cb))!; - att.start(); // started subscriber on run-1 - - // run-2 starts on the same chat while run-1's tee is still reading. registry.open(CHAT, 'run-2'); - expect(c.ended()).toBe(1); // exactly one onEnd from the replace + expect(c.ended()).toBe(1); const newEntry = (registry as any).entries.get(CHAT); expect(newEntry.runId).toBe('run-2'); expect(newEntry.finished).toBe(false); - // The old tee now completes: its late done must NOT double-end nor delete the - // new entry. src.push('b'); src.close(); await flush(); - expect(c.ended()).toBe(1); // still exactly one + expect(c.ended()).toBe(1); const still = (registry as any).entries.get(CHAT); expect(still).toBe(newEntry); expect(still.runId).toBe('run-2'); @@ -299,7 +250,6 @@ describe('AiChatStreamRegistryService', () => { src.push('a'); await flush(); const entry = (registry as any).entries.get(CHAT); - // Frames were NOT ingested (bind bailed), assistantMessageId untouched. expect(entry.frames).toEqual([]); expect(entry.assistantMessageId).toBeUndefined(); }); @@ -310,32 +260,253 @@ describe('AiChatStreamRegistryService', () => { const entry = (registry as any).entries.get(CHAT); expect(entry.finished).toBe(false); }); +}); - it('a throwing onFrame ejects only that subscriber; the ingest loop stays alive', async () => { +/** + * #491 step-stamped retention: the boundary detector, tail-only slicing at the + * client's frontier N, the confirmed-persist rotation (+ anti-inversion), the + * overflow gap, the memory bound, and the finished-retained tail. All observable + * against the REAL registry driven through open/bind/ingest. + */ +describe('AiChatStreamRegistryService step-aligned retention (#491)', () => { + const CHAT = 'chat-s'; + let registry: AiChatStreamRegistryService; + + beforeEach(() => { + registry = new AiChatStreamRegistryService(); + jest.spyOn((registry as any).logger, 'warn').mockImplementation(() => {}); + }); + afterEach(() => registry.onModuleDestroy()); + + const entryOf = () => (registry as any).entries.get(CHAT); + + it('stamps frames by finish-step count, aligned with stepsPersisted', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + // step 0 content, its finish-step, step 1 content, its finish-step, finish. + src.push(textDelta('t0', 'a')); // stamp 0 + src.push(finishStep()); // stamp 0 (the finish-step frame carries the pre value) + src.push(textDelta('t1', 'b')); // stamp 1 + src.push(finishStep()); // stamp 1 + src.push(finish()); // stamp 2 + await flush(); + const e = entryOf(); + expect(e.stamps).toEqual([0, 0, 1, 1, 2]); + expect(e.currentStamp).toBe(2); + }); - const bad = collector(); - const badAtt = (await registry.attach(CHAT, false, undefined, { - onFrame: () => { - throw new Error('boom'); - }, - onEnd: bad.cb.onEnd, - }))!; - badAtt.start(); + it('does NOT treat a text delta that merely quotes "finish-step" as a boundary', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + // A model that literally types "type":"finish-step" — JSON-escaped in the frame. + src.push(textDelta('t0', '"type":"finish-step"')); + await flush(); + expect(entryOf().currentStamp).toBe(0); // no false boundary + }); - const good = collector(); - const goodAtt = (await registry.attach(CHAT, false, undefined, good.cb))!; - goodAtt.start(); - - src.push('a'); // bad throws on this frame -> ejected - src.push('b'); // good still receives both + it('tail-only: attach at N slices frames with stamp >= N', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b')); // 1 + src.push(finishStep()); // 1 + src.push(textDelta('t2', 'c')); // 2 (in-progress) await flush(); - const entry = (registry as any).entries.get(CHAT); - expect(entry.subscribers.size).toBe(1); // bad ejected, good remains - expect(good.frames).toEqual(['a', 'b']); + const c = collector(); + // Client persisted 2 steps -> wants the tail from step 2. + const att = (await registry.attach(CHAT, 'assist-1', 2, c.cb))!; + expect(tail(att.replay)).toEqual([textDelta('t2', 'c')]); + }); + + it('attach in the MIDDLE of a step (N between finish-steps) slices from that step', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b1')); // 1 + src.push(textDelta('t1', 'b2')); // 1 (still step 1, no finish-step yet) + await flush(); + + const c = collector(); + const att = (await registry.attach(CHAT, 'assist-1', 1, c.cb))!; + // Step 0's frames are dropped from the tail; the whole in-progress step 1 is kept. + expect(tail(att.replay)).toEqual([textDelta('t1', 'b1'), textDelta('t1', 'b2')]); + }); + + it('rotates the ring ONLY on a confirmed persist (drops stamp < N)', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b')); // 1 + await flush(); + expect(entryOf().stamps).toEqual([0, 0, 1]); + + // Confirm step 0 persisted (stepsPersisted = 1) -> drop stamp < 1. + registry.confirmPersistedStep(CHAT, 'run-1', 1); + expect(entryOf().stamps).toEqual([1]); + expect(entryOf().persistedFloor).toBe(1); + }); + + it('persist FAILED but the ring still fits -> attach SUCCEEDS and the tail includes step N', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b')); // 1 (step 1's persist FAILED -> no confirm) + await flush(); + // No confirmPersistedStep for step 1: the ring still holds step 1. + + const c = collector(); + // Client's last successful persist was step 0 -> stepsPersisted = 1. + const att = await registry.attach(CHAT, 'assist-1', 1, c.cb); + expect(att).not.toBeNull(); + expect(tail(att!.replay)).toEqual([textDelta('t1', 'b')]); // includes step 1 + }); + + it('persist failed AND the ring overflowed past N -> 204 (coverage gap)', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + // Step 0: a fat step that blows past the cap with NO persist confirmation. + const big = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 2)); + src.push(textDelta('t0', big)); // 0 + src.push(textDelta('t0', big)); // 0 + src.push(textDelta('t0', big)); // 0 -> overflow evicts stamp-0 frames + await flush(); + const e = entryOf(); + expect(e.overflowed).toBe(true); + expect(e.bytes).toBeLessThanOrEqual(registry.maxBufferBytes); + + // A client at frontier 0 falls at/below an evicted step -> gap -> null. + const c = collector(); + expect(await registry.attach(CHAT, 'assist-1', 0, c.cb)).toBeNull(); + }); + + it('stale N (client seed lagged behind a rotation) -> 204; after a refetch (larger N) -> success', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(textDelta('t1', 'b')); // 1 + src.push(finishStep()); // 1 + src.push(textDelta('t2', 'c')); // 2 + await flush(); + // Server confirmed steps 0 and 1 -> rotate away stamp < 2. + registry.confirmPersistedStep(CHAT, 'run-1', 2); + expect(entryOf().stamps).toEqual([2]); + + // A client whose seed still says stepsPersisted = 1 -> below minStamp -> 204. + const stale = collector(); + expect(await registry.attach(CHAT, 'assist-1', 1, stale.cb)).toBeNull(); + + // It refetches (now stepsPersisted = 2) and re-attaches -> success. + const fresh = collector(); + const att = await registry.attach(CHAT, 'assist-1', 2, fresh.cb); + expect(att).not.toBeNull(); + expect(tail(att!.replay)).toEqual([textDelta('t2', 'c')]); + }); + + it('overflow gap CLEARS once a later persist rotates out the holey steps', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + const big = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 2)); + src.push(textDelta('t0', big)); // 0 + src.push(textDelta('t0', big)); // 0 + src.push(finishStep()); // 0 (still stamp 0) + src.push(textDelta('t1', 'small')); // 1 + src.push(finishStep()); // 1 + src.push(textDelta('t2', 'c')); // 2 + await flush(); + expect(entryOf().overflowed).toBe(true); + + // Late persist confirms steps 0..1 -> rotates out the holey step-0 frames. + registry.confirmPersistedStep(CHAT, 'run-1', 2); + // A client at frontier 2 is now cleanly covered (the hole was below it). + const c = collector(); + const att = await registry.attach(CHAT, 'assist-1', 2, c.cb); + expect(att).not.toBeNull(); + expect(tail(att!.replay)).toEqual([textDelta('t2', 'c')]); + }); + + it('finished-retained + N = N_final -> empty tail plus the finish frame', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(finish()); // 1 (N_final = 1) + src.close(); + await flush(); + // The last step's per-step persist confirmed stepsPersisted = 1. + registry.confirmPersistedStep(CHAT, 'run-1', 1); + + const c = collector(); + const att = (await registry.attach(CHAT, 'assist-1', 1, c.cb))!; + expect(att.finished).toBe(true); + // Empty step tail; just the finish frame so the client's SDK closes the stream. + expect(tail(att.replay)).toEqual([finish()]); + // No subscriber registered for a finished run. + expect(entryOf().subscribers.size).toBe(0); + }); + + it('confirmPersistedStep is monotonic and identity-checked', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); + src.push(finishStep()); + src.push(textDelta('t1', 'b')); + await flush(); + registry.confirmPersistedStep(CHAT, 'run-1', 1); + expect(entryOf().persistedFloor).toBe(1); + // A stale lower count is ignored. + registry.confirmPersistedStep(CHAT, 'run-1', 0); + expect(entryOf().persistedFloor).toBe(1); + // A foreign runId is ignored. + registry.confirmPersistedStep(CHAT, 'WRONG', 5); + expect(entryOf().persistedFloor).toBe(1); + }); + + it('MEMORY BOUND: 5 parallel marathon runs each stream well past 32MB; each ring stays <= the cap', async () => { + const cap = registry.maxBufferBytes; + const chats = ['m0', 'm1', 'm2', 'm3', 'm4']; + const srcs = chats.map((chat) => { + registry.open(chat, `run-${chat}`); + const s = makePushStream(); + registry.bind(chat, `run-${chat}`, `assist-${chat}`, s.stream); + return s; + }); + // ~256KB frames; 160 per chat = 40MB streamed each, well past the old 32MB. + // Interleave a finish-step every 8 frames so steps advance realistically. No + // persist confirmation -> the ONLY thing keeping memory bounded is the cap. + const frame = 'y'.repeat(256 * 1024); + for (let batch = 0; batch < 20; batch++) { + for (let i = 0; i < 8; i++) { + for (const s of srcs) s.push(textDelta('t', frame)); + } + for (const s of srcs) s.push(finishStep()); + await flush(); // drain the pump so queues never hold a whole run + } + let total = 0; + for (const chat of chats) { + const e = (registry as any).entries.get(chat); + expect(e.bytes).toBeLessThanOrEqual(cap); + total += e.bytes; + } + // Total retained across all 5 runs is bounded by 5x the per-run cap — the old + // registry would have retained ~5x40MB = 200MB here. + expect(total).toBeLessThanOrEqual(cap * chats.length); }); }); @@ -361,7 +532,7 @@ describe('AiChatStreamRegistryService retention timers', () => { it('a finished entry is removed after the retention window', () => { registry.open(CHAT, 'run-1'); - registry.abortEntry(CHAT, 'run-1'); // finalize -> retention armed + registry.abortEntry(CHAT, 'run-1'); expect((registry as any).entries.get(CHAT)).toBeDefined(); jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1); expect((registry as any).entries.get(CHAT)).toBeUndefined(); @@ -369,20 +540,18 @@ describe('AiChatStreamRegistryService retention timers', () => { it('retention deletes ONLY its own entry (invariant 2)', () => { registry.open(CHAT, 'run-1'); - registry.abortEntry(CHAT, 'run-1'); // arm retention for entry A - // Simulate the race where the key was replaced without clearing A's timer. + registry.abortEntry(CHAT, 'run-1'); const sentinel = { marker: true }; (registry as any).entries.set(CHAT, sentinel); jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1); - // A's timer saw entries.get(CHAT) !== A, so it did NOT delete the successor. expect((registry as any).entries.get(CHAT)).toBe(sentinel); }); it('open() over a retained entry clears its timer and the successor survives', () => { registry.open(CHAT, 'run-1'); - registry.abortEntry(CHAT, 'run-1'); // retained, timer armed + registry.abortEntry(CHAT, 'run-1'); const clearSpy = jest.spyOn(global, 'clearTimeout'); - registry.open(CHAT, 'run-2'); // must clear run-1's retain timer + registry.open(CHAT, 'run-2'); expect(clearSpy).toHaveBeenCalled(); jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1); const entry = (registry as any).entries.get(CHAT); diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts b/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts index 67aa7cbe..05bc19b8 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts @@ -8,10 +8,12 @@ import { SUBSCRIBER_MAX_BUFFERED_BYTES } from './ai-chat-stream-registry.service import type { User, Workspace } from '@docmost/db/types/entity.types'; /** - * Wiring spec for the #184 phase 1.5 attach endpoint + * Wiring spec for the #184 phase 1.5 attach endpoint (tail-only #491) * (`GET /ai-chat/runs/:chatId/stream`). Owner-gated via assertOwnedChat; the - * registry is mocked so this exercises ONLY the controller's replay/live/204/ - * cleanup wiring against a fake raw socket. Constructor order is (aiChatService, + * registry is mocked so this exercises ONLY the controller's tail-write/live/204/ + * cleanup wiring against a fake raw socket. The attach signature is now + * `(chatId, anchor, n, cb)` — the client hands its persisted step frontier `n` + * and its assistant row id `anchor`. Constructor order is (aiChatService, * aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo, * streamRegistry, environment). */ @@ -86,8 +88,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { attach: jest.fn( ( _chatId: string, - _live: boolean, _anchor: string | undefined, + _n: number, cb: RunStreamCallbacks, ) => { capturedCb = cb; @@ -156,7 +158,7 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { expect(res.hijack).not.toHaveBeenCalled(); }); - it('threads expect=live and anchor through to the registry', async () => { + it('threads anchor and the numeric frontier n through to the registry', async () => { const { controller, streamRegistry } = makeController({ chat: owned, attachment: null, @@ -165,8 +167,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { const { req } = makeReq(); await controller.attachRunStream( 'c1', - 'live', 'anchor-1', + '2', req, res, user, @@ -174,13 +176,13 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { ); expect(streamRegistry.attach).toHaveBeenCalledWith( 'c1', - true, 'anchor-1', + 2, // parsed to a number expect.anything(), ); }); - it('passes expect=false when the query is absent', async () => { + it('floors n to 0 when the query is absent/invalid', async () => { const { controller, streamRegistry } = makeController({ chat: owned, attachment: null, @@ -198,8 +200,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { ); expect(streamRegistry.attach).toHaveBeenCalledWith( 'c1', - false, undefined, + 0, expect.anything(), ); }); @@ -245,8 +247,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { const { req } = makeReq(); await controller.attachRunStream( 'c1', - 'live', 'a1', + '1', req, res, user, diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.ts b/apps/server/src/core/ai-chat/ai-chat.controller.ts index 49267968..b97ad6ff 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.ts @@ -64,6 +64,47 @@ import { SUBSCRIBER_MAX_BUFFERED_BYTES, } from './ai-chat-stream-registry.service'; import { startSseHeartbeat } from './sse-resilience'; + +/** + * Write the attach TAIL to the hijacked socket in chunks that RESPECT drain + * (#491): each `write()` that returns false (the kernel buffer is full) is awaited + * on the next 'drain' before continuing. The old code wrote the whole buffer + * synchronously, which — with the pre-#491 32MB ring — spiked memory (half the + * OOM). Bails immediately if the socket ended/errored mid-write. Frames that the + * paused registry subscriber buffers while this awaits are delivered by start(). + */ +async function writeTailRespectingDrain( + raw: { + write(chunk: string): boolean; + writableEnded?: boolean; + destroyed?: boolean; + once(event: string, cb: () => void): unknown; + removeListener?(event: string, cb: () => void): unknown; + }, + frames: string[], +): Promise { + for (const frame of frames) { + if (raw.writableEnded || raw.destroyed) return; + const ok = raw.write(frame); + if (!ok) { + // Kernel buffer full — wait for drain (or an early close/error) before the + // next chunk, so a slow reader never forces the whole tail into memory. + // Remove ALL three listeners once any fires, so a many-chunk tail with + // repeated backpressure never leaks (MaxListenersExceededWarning). + await new Promise((resolve) => { + const finish = (): void => { + raw.removeListener?.('drain', finish); + raw.removeListener?.('close', finish); + raw.removeListener?.('error', finish); + resolve(); + }; + raw.once('drain', finish); + raw.once('close', finish); + raw.once('error', finish); + }); + } + } +} import { EnvironmentService } from '../../integrations/environment/environment.service'; /** @@ -290,19 +331,25 @@ export class AiChatController { } /** - * Attach to a chat's live run stream (#184 phase 1.5). A late/reloaded tab - * replays the frames buffered so far and then follows the live tail as a normal - * streamer. Owner-gated via assertOwnedChat (same gate as getRun). When there is - * nothing to resume — no entry, a finished run without expect=live, an - * overflowed buffer, or an anchor that pins a DIFFERENT run — the endpoint - * answers 204, the ONLY "nothing to resume" signal the AI SDK's reconnect - * accepts (it maps 204 to a silent no-op). With AI_CHAT_RESUMABLE_STREAM off the - * registry is never populated, so attach always 204s. + * Attach to a chat's live run stream from the client's step frontier (#184 phase + * 1.5, tail-only #491). A late/reloaded tab hands the server the step count it + * has PERSISTED (`n` = the seeded row's `metadata.stepsPersisted`) and its + * assistant row id (`anchor`); the registry answers with the TAIL past step `n` + * (a synthetic `start` frame + the buffered frames stamped >= n) and then the + * live tail. Owner-gated via assertOwnedChat (same gate as getRun). When there + * is nothing to resume — no entry, a ring that does not cover the client's + * frontier (overflow gap, or the client's seed lagged a rotation), or an anchor + * that pins a DIFFERENT run (invariant 6) — the endpoint answers 204, the ONLY + * "nothing to resume" signal the AI SDK's reconnect accepts (it maps 204 to a + * silent no-op); the client then refetches (a larger n) and re-attaches. With + * AI_CHAT_RESUMABLE_STREAM off the registry is never populated, so attach always + * 204s. * - * `expect=live` opts into replaying a finished-but-retained run (safe only when - * the client stripped the streaming tail); `anchor` is the client's assistant - * row id, which must match this run's (invariant 6) or a foreign run's - * transcript would be replayed into the store. + * The step marker `n` comes ONLY from the client — the server never reads the + * row to derive it, because a server-side n from a stale seed would open a + * silent one-step hole. The tail is written to the socket in CHUNKS respecting + * drain (writeTailRespectingDrain): the old code synchronously blasted the whole + * buffer, which — with the old 32MB cap — was half the OOM. */ @SkipTransform() @UseGuards(JwtAuthGuard, UserThrottlerGuard) @@ -310,39 +357,41 @@ export class AiChatController { @Get('runs/:chatId/stream') async attachRunStream( @Param('chatId', new ParseUUIDPipe()) chatId: string, - @Query('expect') expect: string | undefined, @Query('anchor') anchor: string | undefined, + @Query('n') n: string | undefined, @Req() req: FastifyRequest, @Res() res: FastifyReply, @AuthUser() user: User, @AuthWorkspace() workspace: Workspace, ): Promise { await this.assertOwnedChat(chatId, user, workspace); // same gate as getRun + // The client's persisted step frontier. A missing/invalid value floors to 0 + // ("give me everything") which, past any rotation, safely 204s. + const frontier = Number.isFinite(Number(n)) ? Math.max(0, Number(n)) : 0; + // The per-subscriber backpressure cap tracks the (env-tunable) ring cap. + const subscriberCap = + this.streamRegistry?.subscriberMaxBufferedBytes ?? + SUBSCRIBER_MAX_BUFFERED_BYTES; let stopHeartbeat: () => void = () => undefined; - const attachment = await this.streamRegistry?.attach( - chatId, - expect === 'live', - anchor, - { - onFrame: (frame) => { - // Backpressure guard: 2x the replay cap, so the initial replay burst - // alone can never trip it; only a genuinely stalled socket can. - try { - if (res.raw.writableLength > SUBSCRIBER_MAX_BUFFERED_BYTES) { - res.raw.destroy(); // 'close' fires -> unsubscribe below - return; - } - if (!res.raw.writableEnded) res.raw.write(frame); - } catch { - res.raw.destroy(); + const attachment = await this.streamRegistry?.attach(chatId, anchor, frontier, { + onFrame: (frame) => { + // Backpressure guard: 2x the ring cap, so the initial tail burst alone + // can never trip it; only a genuinely stalled socket can. + try { + if (res.raw.writableLength > subscriberCap) { + res.raw.destroy(); // 'close' fires -> unsubscribe below + return; } - }, - onEnd: () => { - stopHeartbeat(); - if (!res.raw.writableEnded) res.raw.end(); - }, + if (!res.raw.writableEnded) res.raw.write(frame); + } catch { + res.raw.destroy(); + } }, - ); + onEnd: () => { + stopHeartbeat(); + if (!res.raw.writableEnded) res.raw.end(); + }, + }); if (!attachment) { res.status(204).send(); // the ONLY "nothing to resume" signal the SDK accepts return; @@ -371,13 +420,16 @@ export class AiChatController { // deliberately NO Connection/Keep-Alive (hop-by-hop; Safari/HTTP2) }); res.raw.flushHeaders?.(); - for (const frame of attachment.replay) res.raw.write(frame); + // Write the tail in chunks respecting drain (not a synchronous blast, which + // was half the OOM). Frames the paused subscriber buffers meanwhile are + // drained by start() below; its cap is the backstop for a stalled socket. + await writeTailRespectingDrain(res.raw, attachment.replay); if (attachment.finished) { - res.raw.end(); + if (!res.raw.writableEnded) res.raw.end(); return; } stopHeartbeat = startSseHeartbeat(res.raw, 15_000); - attachment.start(); // drain pending accumulated during replay, go live + attachment.start(); // drain pending accumulated during the tail write, go live } catch { attachment.unsubscribe(); stopHeartbeat(); 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 deb5b127..e58788b0 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -1543,30 +1543,39 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // Per-step (non-terminal) update: persist the finished steps the moment a // step ends. Tolerant — a failed update is logged and swallowed so it never // throws into the stream. Keeps status 'streaming'. - const updateStreaming = async (): Promise => { - if (!assistantId) return; + // + // #491: it now SIGNALS its outcome — the persisted `stepsPersisted` count on + // a CONFIRMED write, or null when it was skipped/failed. The caller rotates + // the run-stream registry ring ONLY on a non-null return (a confirmed + // persist), so a failed persist never rotates away a step nobody has (the + // classic inversion bug); a failure just makes the ring cover more. + const updateStreaming = async (): Promise => { + if (!assistantId) return null; // Cheap short-circuit once the turn is finalized (see `finalized` below). // The AUTHORITATIVE guard is `onlyIfStreaming` on the UPDATE: a late // fire-and-forget step update could still be in flight on another pool // connection when finalize runs, so the SQL `WHERE status='streaming'` // (not this flag) is what prevents it clobbering the terminal row. - if (finalized) return; + if (finalized) return null; + // Build the flush ONCE so the returned count is EXACTLY the persisted + // `stepsPersisted` (both derive from capturedSteps.length at this instant). + const flushed = flushAssistant(capturedSteps, '', 'streaming', { + pageChanged, + partsCache, + }); + const stepsPersisted = flushed.metadata.stepsPersisted as number; try { - await this.aiChatMessageRepo.update( - assistantId, - workspace.id, - flushAssistant(capturedSteps, '', 'streaming', { - pageChanged, - partsCache, - }), - { onlyIfStreaming: true }, - ); + await this.aiChatMessageRepo.update(assistantId, workspace.id, flushed, { + onlyIfStreaming: true, + }); + return stepsPersisted; } catch (err) { this.logger.warn( `Failed to update streaming assistant row: ${ err instanceof Error ? err.message : 'unknown error' }`, ); + return null; } }; @@ -1742,7 +1751,24 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // this point still recovers the step. Not awaited here (never block the // stream), but SERIALIZED via stepUpdateChain so the writes commit in // step order; updateStreaming is error-tolerant (logs + swallows). - stepUpdateChain = stepUpdateChain.then(() => updateStreaming()); + // #491: on a CONFIRMED persist, rotate the run-stream registry ring to + // drop the now-on-disk steps (stamp < stepsPersisted). Gated on the + // resumable flag (same as open/bind) and identity-checked in the + // registry; a null return (skipped/failed) rotates NOTHING (auto-safe). + stepUpdateChain = stepUpdateChain.then(async () => { + const persisted = await updateStreaming(); + if ( + persisted != null && + runId && + this.environment?.isAiChatResumableStreamEnabled?.() + ) { + this.streamRegistry?.confirmPersistedStep( + chatId, + runId, + persisted, + ); + } + }); // #184: persist the run's progress (finished-step count). Fire-and- // forget; the hook swallows its own errors. if (runId) runHooks?.onStep?.(runId, capturedSteps.length); diff --git a/apps/server/test/integration/ai-chat-attach.int-spec.ts b/apps/server/test/integration/ai-chat-attach.int-spec.ts index 70da3cbd..6e57199e 100644 --- a/apps/server/test/integration/ai-chat-attach.int-spec.ts +++ b/apps/server/test/integration/ai-chat-attach.int-spec.ts @@ -30,11 +30,13 @@ import { * tees the SSE frames into it via `consumeSseStream` while stamping the DB row id * via `generateMessageId` (both gated on runId + the resumable flag). * - * Proven here: a finished run's replay is the full frame sequence incl `[DONE]` - * with `start.messageId` == the seeded DB row id; the anchor check (invariant 6); - * an attach opened BEFORE the first frame follows the live stream from frame 0; an - * explicit stop surfaces `{"type":"abort"}` + `[DONE]` + end to the subscriber; - * and the legacy (non-run) path tees nothing. + * Proven here (tail-only #491): a finished run attached at its persisted frontier + * N_final delivers only the TAIL past N (a synthetic `start` carrying the run-fact + * + the terminal `finish`/`[DONE]`) — the step content below N lives in the seeded + * DB row, NOT the ring; the anchor check (invariant 6); an attach opened BEFORE the + * first frame follows the live stream from frame 0; an explicit stop surfaces + * `{"type":"abort"}` + `[DONE]` + end to the subscriber; and the legacy (non-run) + * path tees nothing. */ const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); @@ -133,14 +135,16 @@ function liveSink(): { }; } -// The SSE `start` frame carries the message id; pull it out of a `data: {...}`. -function parseStartMessageId(frames: string[]): string | undefined { +// Parse the first `start` frame's JSON out of a `data: {...}` sequence. +function parseStartFrame( + frames: string[], +): { messageId?: string; messageMetadata?: any } | undefined { for (const f of frames) { const m = /^data: (\{.*\})\s*$/m.exec(f.trim()); if (!m) continue; try { const json = JSON.parse(m[1]); - if (json.type === 'start') return json.messageId; + if (json.type === 'start') return json; } catch { /* not this frame */ } @@ -270,7 +274,7 @@ describe('AiChatService run-stream attach [integration]', () => { await destroyTestDb(); }); - it('run-wrapped: replay is the full frame sequence incl [DONE], start.messageId == the seeded DB row id', async () => { + it('run-wrapped, tail-only: a finished run at N_final delivers the run-fact start + finish/[DONE]; the step content lives in the seeded row', async () => { const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; const registry = new AiChatStreamRegistryService(); const runService = new AiChatRunService(runRepo, { @@ -300,27 +304,37 @@ describe('AiChatService run-stream attach [integration]', () => { ); }); const rowId = await assistantRowId(chatId); + // The client reads its persisted step frontier N from the seeded row. + const row: any = await msgRepo.findById(rowId, workspaceId); + const nFinal = row.metadata.stepsPersisted as number; + expect(nFinal).toBe(1); // a single finished step + // The step content is in the SEEDED row (parts/content), not the ring. + expect(JSON.stringify(row.metadata.parts)).toContain('Hello'); - // Finished-run replay with expect=live + the correct anchor. + // Attach at N_final with the correct anchor: the tail past step 1 is just + // the terminal frames; step 0's 'Hello' is BELOW the frontier (seeded). const sink = liveSink(); - const att = await registry.attach(chatId, true, rowId, sink.cb); + const att = await registry.attach(chatId, rowId, nFinal, sink.cb); expect(att).not.toBeNull(); expect(att!.finished).toBe(true); - // The tee captured frames (consumeSseStream was wired). - expect(att!.replay.length).toBeGreaterThan(0); - // generateMessageId stamped the DB row id onto the streamed start frame. - expect(parseStartMessageId(att!.replay)).toBe(rowId); - // The full sequence includes the streamed text and the terminal marker. - const joined = att!.replay.join(''); - expect(joined).toContain('Hello'); + // The synthetic start frame carries the run-fact (runId/chatId), the source + // of the run-fact on re-attach. + const start = parseStartFrame(att!.replay); + expect(start?.messageMetadata).toMatchObject({ + runId: box.runId, + chatId, + }); + // The terminal marker is delivered so the client's SDK closes the stream. expect(att!.replay.some((f) => f.includes('[DONE]'))).toBe(true); + // 'Hello' (step 0, below the frontier) is NOT re-streamed — it is seeded. + expect(att!.replay.some((f) => f.includes('Hello'))).toBe(false); } finally { registry.onModuleDestroy(); await cleanup(); } }); - it('anchor mismatch with expect=live returns null (invariant 6)', async () => { + it('anchor mismatch returns null (invariant 6)', async () => { const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; const registry = new AiChatStreamRegistryService(); const runService = new AiChatRunService(runRepo, { @@ -347,7 +361,7 @@ describe('AiChatService run-stream attach [integration]', () => { const sink = liveSink(); // A foreign anchor must NOT replay this run's transcript. expect( - await registry.attach(chatId, true, 'a-different-run-row', sink.cb), + await registry.attach(chatId, 'a-different-run-row', 1, sink.cb), ).toBeNull(); } finally { registry.onModuleDestroy(); @@ -376,8 +390,11 @@ describe('AiChatService run-stream attach [integration]', () => { try { // Attach while the entry exists (opened at begin) but before any frame. const sink = liveSink(); - const att = (await registry.attach(chatId, false, undefined, sink.cb))!; - expect(att.replay).toEqual([]); // nothing streamed yet -> replay from 0 + const att = (await registry.attach(chatId, undefined, 0, sink.cb))!; + // Nothing streamed yet -> the tail is just the synthetic start frame; the + // whole live stream (start..DONE) follows via onFrame after start(). + expect(att.replay).toHaveLength(1); + expect(att.replay[0]).toContain('"type":"start"'); att.start(); // go live (drains nothing, then follows) // Now emit the whole turn. @@ -448,7 +465,7 @@ describe('AiChatService run-stream attach [integration]', () => { }); try { const sink = liveSink(); - const att = (await registry.attach(chatId, false, undefined, sink.cb))!; + const att = (await registry.attach(chatId, undefined, 0, sink.cb))!; att.start(); // Give streamText a beat to begin consuming the partial output. @@ -526,7 +543,9 @@ describe('AiChatService run-stream attach [integration]', () => { expect(entry).toBeDefined(); expect(entry.finished).toBe(true); const sink = liveSink(); - expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull(); + // Finished with an EMPTY ring (aborted before any frame) -> null -> the + // client degrades to poll instead of hanging on an empty stream. + expect(await registry.attach(chatId, undefined, 0, sink.cb)).toBeNull(); } finally { registry.onModuleDestroy(); await cleanup(); @@ -556,8 +575,8 @@ describe('AiChatService run-stream attach [integration]', () => { }); const sink = liveSink(); // No entry was ever opened; attach always yields null. - expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull(); - expect(await registry.attach(chatId, true, 'anything', sink.cb)).toBeNull(); + expect(await registry.attach(chatId, undefined, 0, sink.cb)).toBeNull(); + expect(await registry.attach(chatId, 'anything', 1, sink.cb)).toBeNull(); } finally { registry.onModuleDestroy(); await cleanup(); -- 2.52.0 From daca5ce8d61f174792b90e1aa2fc807652813b98 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 19:16:59 +0300 Subject: [PATCH 4/7] =?UTF-8?q?fix(ai-chat):=20=D1=86=D0=B5=D0=BB=D0=BE?= =?UTF-8?q?=D1=81=D1=82=D0=BD=D0=BE=D1=81=D1=82=D1=8C=20delta-spec=20(?= =?UTF-8?q?=D0=BC=D0=BE=D0=BB=D1=87=D0=B0=D0=BB=D0=B8=D0=B2=D1=8B=D0=B9=20?= =?UTF-8?q?=D1=81=D0=BA=D0=B8=D0=BF)=20+=202=20hardening=20(#491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Внутреннее ре-ревью: DB-backed delta-spec молча СКИПАЛСЯ — нулевое покрытие инварианта DB-clock курсора. - Импорт `import postgres from 'postgres'` (default) при tsconfig module:commonjs без esModuleInterop компилился в `postgres_1.default(...)`, а CJS-`postgres` не имеет `.default` → TypeError в beforeAll → пустой catch → reachable=false → все 6 тестов уходили в console.warn('SKIP') и return БЕЗ ассертов (suite оставался бы зелёным даже при регрессии на new Date()). Фикс: `import * as postgres from 'postgres'` (как в рабочем int-харнессе). - Хардненг харнесса: реальная ошибка программирования в beforeAll больше НЕ маскируется под «DB unreachable» — скип легитимен только для сетевого отказа (ECONNREFUSED и т.п.), иначе rethrow → suite падает громко. - Два DB-clock теста использовали jest.useFakeTimers() целиком, что замораживало внутренние таймеры postgres.js → awaited DB round-trip зависал на 5s-капе (и вешал afterAll). Фейкаем ТОЛЬКО Date (doNotFake всех таймеров) — запрос резолвится, а инвариант «стамп от часов БД, не app-clock» по-прежнему доказан (скос процесс-часов в 2099 → стамп остаётся на времени БД). Теперь все 6 тестов РЕАЛЬНО исполняются и зелёные против живого Postgres. Два дешёвых hardening из ревью: - registry coverageFloor: пустая ветка возвращает max(currentStamp, persistedFloor) — инвариант «клиент с n=persistedFloor всегда покрыт» структурный, а не тайминг-зависимый. - GetChatDeltaDto.cursor: @IsString → @IsISO8601 — битый курсор отсекается 400 на уровне DTO, а не 500 на `::timestamptz`. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ai-chat-stream-registry.service.ts | 8 ++- .../src/core/ai-chat/dto/ai-chat.dto.ts | 13 +++- .../ai-chat-message.repo.delta.spec.ts | 59 +++++++++++++++++-- 3 files changed, 71 insertions(+), 9 deletions(-) diff --git a/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts index 5fb3c15a..743bad30 100644 --- a/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts @@ -440,7 +440,13 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { * An attach at frontier `n` is covered ⟺ coverageFloor <= n. */ private coverageFloor(entry: Entry): number { - if (entry.frames.length === 0) return entry.currentStamp; + // Empty ring: only the live tail is coming. The floor is the current step, + // but never below persistedFloor — a confirmed persist can rotate the ring + // empty while currentStamp still lags a beat behind on another connection, so + // max() keeps the invariant STRUCTURAL (a client with n = persistedFloor is + // always covered) rather than timing-dependent. + if (entry.frames.length === 0) + return Math.max(entry.currentStamp, entry.persistedFloor); const min = entry.stamps[0]; return entry.overflowThroughStamp >= min ? min + 1 : min; } diff --git a/apps/server/src/core/ai-chat/dto/ai-chat.dto.ts b/apps/server/src/core/ai-chat/dto/ai-chat.dto.ts index ec09c6aa..517f74ed 100644 --- a/apps/server/src/core/ai-chat/dto/ai-chat.dto.ts +++ b/apps/server/src/core/ai-chat/dto/ai-chat.dto.ts @@ -1,4 +1,10 @@ -import { IsOptional, IsString, MaxLength, MinLength } from 'class-validator'; +import { + IsISO8601, + IsOptional, + IsString, + MaxLength, + MinLength, +} from 'class-validator'; /** Identify a chat by id (workspace-scoped on the server). */ export class ChatIdDto { @@ -47,8 +53,11 @@ export class GetChatDeltaDto { @IsString() chatId: string; + // ISO-8601 timestamp echoed from the previous poll's response. Validated as + // ISO-8601 (not a bare string): a malformed cursor would otherwise reach the + // `::timestamptz` cast in findByChatUpdatedAfter and 500 instead of a clean 400. @IsOptional() - @IsString() + @IsISO8601() cursor?: string; } diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts index e3086de8..94ccddeb 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts @@ -1,7 +1,12 @@ import { randomUUID } from 'crypto'; import { CamelCasePlugin, Kysely, sql } from 'kysely'; import { PostgresJSDialect } from 'kysely-postgres-js'; -import postgres from 'postgres'; +// NOT a default import: the project tsconfig is `module: commonjs` with NO +// esModuleInterop, so `import postgres from 'postgres'` compiles to +// `postgres_1.default(...)` and the CJS `postgres` export has no `.default` — +// it threw in beforeAll, was swallowed as "DB unreachable", and SILENTLY voided +// all six tests. Mirror the working integration harness (test/integration/db.ts). +import * as postgres from 'postgres'; import { AiChatMessageRepo } from './ai-chat-message.repo'; import { AiChatRunRepo } from './ai-chat-run.repo'; @@ -50,8 +55,21 @@ beforeAll(async () => { await sql`set session_replication_role = replica`.execute(db); await sql`select 1`.execute(db); reachable = true; - } catch { + } catch (err) { reachable = false; + // A genuine connection failure (ECONNREFUSED etc.) is a legitimate skip on a + // DB-less CI. A PROGRAMMING error (bad import, typo, driver misuse) must NOT + // masquerade as "DB unreachable" and silently void the whole suite (that is + // exactly the bug that hid this spec's zero coverage) — rethrow it so the + // suite fails LOUDLY. + const msg = String((err as Error)?.message ?? err); + if ( + !/ECONNREFUSED|ENOTFOUND|ETIMEDOUT|EHOSTUNREACH|connect|terminating|password|authentication|role .* does not exist|database .* does not exist/i.test( + msg, + ) + ) { + throw err; + } } msgRepo = new AiChatMessageRepo(db as never); runRepo = new AiChatRunRepo(db as never); @@ -96,6 +114,34 @@ async function dbNow(): Promise { return r.rows[0].now.toISOString(); } +// Fake ONLY the Date object (so in-process `new Date()`/`Date.now()` jump), while +// leaving every TIMER function real. Faking timers wholesale freezes postgres.js's +// internal connection/query timers, so the awaited DB round-trip would hang the +// test (and the afterAll cleanup) at the jest 5s cap. With Date-only faking the +// query resolves normally, and we still prove the stamp is the DB clock (not the +// skewed process clock). +function fakeDateOnly(iso: string): void { + jest.useFakeTimers({ + doNotFake: [ + 'hrtime', + 'nextTick', + 'performance', + 'queueMicrotask', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'requestIdleCallback', + 'cancelIdleCallback', + 'setImmediate', + 'clearImmediate', + 'setInterval', + 'clearInterval', + 'setTimeout', + 'clearTimeout', + ], + now: new Date(iso), + }); +} + const maybe = (name: string, fn: () => Promise) => it(name, async () => { if (!reachable) { @@ -188,9 +234,10 @@ describe('AiChatMessageRepo.findByChatUpdatedAfter (#491 delta poll)', () => { 'update() stamps updatedAt from the DB clock, not the app clock', async () => { const m = await seedMessage(); - // Fake the PROCESS clock ~73 years into the future. If the stamp came from - // `new Date()` the row would read year 2099; sql now() keeps it on DB time. - jest.useFakeTimers().setSystemTime(new Date('2099-01-01T00:00:00Z')); + // Skew the PROCESS clock ~73 years into the future (Date only). If the stamp + // came from `new Date()` the row would read year 2099; sql now() keeps it on + // DB time. + fakeDateOnly('2099-01-01T00:00:00Z'); const updated = await msgRepo.update(m.id, workspaceId, { content: 'y', }); @@ -211,7 +258,7 @@ describe('AiChatMessageRepo.findByChatUpdatedAfter (#491 delta poll)', () => { status: 'running', stepCount: 0, } as never); - jest.useFakeTimers().setSystemTime(new Date('2099-01-01T00:00:00Z')); + fakeDateOnly('2099-01-01T00:00:00Z'); const updated = await runRepo.update(run.id, workspaceId, { stepCount: 1, }); -- 2.52.0 From d81781aa27d22531b921b20cea0477a8cf3c1169 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 19:40:19 +0300 Subject: [PATCH 5/7] =?UTF-8?q?feat(ai-chat):=20=D0=BA=D0=BB=D0=B8=D0=B5?= =?UTF-8?q?=D0=BD=D1=82=D1=81=D0=BA=D0=B0=D1=8F=20=D0=B0=D0=BA=D1=82=D0=B8?= =?UTF-8?q?=D0=B2=D0=B0=D1=86=D0=B8=D1=8F=20tail-only=20resume=20+=20delta?= =?UTF-8?q?-=D0=BF=D0=BE=D0=BB=D0=BB=D0=B8=D0=BD=D0=B3=20(#491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Переводит клиент на серверный tail-only контракт resume (задел commit 3), не трогая FSM run-fsm.ts — меняется только рантайм-обвязка в chat-thread.tsx. A. Убран STRIP-механизм. Seed теперь содержит ВСЕ персистнутые строки без изъятия хвоста (стриминговый хвост — это шаги 0..N-1, к которым SDK- продолжение дописывает tail). stripRef/strippedRowRef заменены на anchorRef { id, stepsPersisted } — персистнутая assistant-строка, питающая ?anchor=&n=. Восстановления stripped-строки на 204/NONE/starved удалены (строку никто не изымал — нечего восстанавливать); invalidateQueries + диспатчи FSM сохранены. Блок anchor-mismatch в reconcile сверяется по id из свежей персист-истории, а не по «живой» строке. B. Вход в attaching/reconnecting — ВСЕГДА через re-seed из персиста; «живой» стор НИКОГДА не база для tail-apply. На локальном FINISH_DISCONNECT (и на live-follow повторном дропе observer-а) сначала getRun(chatId) → замена «живой» частичной строки персистнутой по id (mergeById) + установка anchor, и лишь ПОСЛЕ этого диспатч RUN_FACT + FINISH_DISCONNECT (который планирует реконнект). Так attach не может продублировать частичный шаг N. Фильтр «живой» строки в resumeStream-эффекте убран (его заменяет re-seed). Инварианты FSM (I1 epoch-штамп, I4 honor-in-stopping, DISCONNECT-FIRST, сброс ownership на терминалах, render-gate) сохранены. C. URL attach: ?anchor=&n= при наличии якоря, без expect. D. Degraded-поллинг переведён с полного рефетча всех страниц на дельту: useAiChatMessagesQuery больше не поллит (seed один раз), а окно при degradedPoll раз в 2.5с зовёт getAiChatMessagesDelta(chatId, cursor) и идемпотентно по id мёржит строки в тот же infinite-query кэш через новый чистый хелпер mergeDeltaRowsIntoPages. Арминг/разарм (onResumeFallback) и idle-cap не тронуты. Хелперы: seedRows удалён; добавлены stepsPersistedOf и mergeDeltaRowsIntoPages (+ юнит-тесты на идемпотентность). Тесты chat-thread обновлены под новый URL, seed-без-стрипа, re-seed-из-персиста на дисконнекте (mutation-verify: падают без re-seed и при n мимо персиста) и 204→poll-без-restore. Весь ai-chat vitest зелёный (398), tsc без новых ошибок. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ai-chat/components/ai-chat-window.tsx | 66 ++++- .../ai-chat/components/chat-thread.test.tsx | 126 +++++++--- .../ai-chat/components/chat-thread.tsx | 233 ++++++++++-------- .../ai-chat/utils/resume-helpers.test.ts | 110 +++++++-- .../features/ai-chat/utils/resume-helpers.ts | 67 ++++- 5 files changed, 429 insertions(+), 173 deletions(-) diff --git a/apps/client/src/features/ai-chat/components/ai-chat-window.tsx b/apps/client/src/features/ai-chat/components/ai-chat-window.tsx index 4b8b07c0..dbb90e53 100644 --- a/apps/client/src/features/ai-chat/components/ai-chat-window.tsx +++ b/apps/client/src/features/ai-chat/components/ai-chat-window.tsx @@ -58,8 +58,11 @@ import ConversationList from "@/features/ai-chat/components/conversation-list.ts import ChatThread from "@/features/ai-chat/components/chat-thread.tsx"; import { exportAiChat, + getAiChatMessagesDelta, stopRun, } from "@/features/ai-chat/services/ai-chat-service.ts"; +import { mergeDeltaRowsIntoPages } from "@/features/ai-chat/utils/resume-helpers.ts"; +import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.ts"; import { useChatSession } from "@/features/ai-chat/hooks/use-chat-session.ts"; import { shouldCollapseOnOutsidePointer, @@ -269,17 +272,64 @@ export default function AiChatWindow() { const { data: messageRows, isLoading: messagesLoading } = useAiChatMessagesQuery( activeChatId ?? undefined, - // DELIBERATELY DUMB: poll every 2.5s WHILE ARMED, otherwise off. NO error - // checks (TanStack resets fetchFailureCount each fetch; the poll must survive - // a server restart), NO tail checks, NO cap here — the settled/stalled/idle-cap - // semantics all live in ChatThread's FSM, which disarms via onResumeFallback. - () => (degradedPoll === true ? 2500 : false), - // #344: gate on windowOpen too — no message history is fetched (and no - // degraded poll runs) while the window is closed; it loads when the window - // opens with an active chat. + // #491: the full infinite-query no longer POLLS. It seeds the thread ONCE; the + // degraded fallback now runs a DELTA poller (below) that augments THIS cache + // idempotently, instead of refetching every page (with full parts) every 2.5s. + false, + // #344: gate on windowOpen too — no message history is fetched while the window + // is closed; it loads when the window opens with an active chat. windowOpen, ); + // #491 degraded DELTA poll. While armed (degradedPoll) and the window is open on a + // chat, poll POST /ai-chat/messages/delta every 2.5s: it returns only the rows + // CHANGED since the previous cursor (+ the run fact) in ONE round-trip. We merge + // those rows into the SAME infinite-query cache the thread reads (idempotently by + // id — the delta's overlap window re-delivers rows), so the thread's reconcile + // effect follows the detached run to its terminal row from a fraction of the wire + // cost. The run-fact settle stays the thread FSM's job (row-status reconcile), so + // we do NOT double-poll /run here. Cursor resets when the chat changes / disarms. + const deltaCursorRef = useRef(undefined); + useEffect(() => { + deltaCursorRef.current = undefined; + }, [activeChatId, degradedPoll]); + useEffect(() => { + if (!degradedPoll || !windowOpen || !activeChatId) return; + const chatId = activeChatId; + let cancelled = false; + const tick = async (): Promise => { + try { + const res = await getAiChatMessagesDelta(chatId, deltaCursorRef.current); + if (cancelled) return; + deltaCursorRef.current = res.cursor; + if (res.rows.length > 0) { + queryClient.setQueryData( + AI_CHAT_MESSAGES_RQ_KEY(chatId), + ( + old: + | { + pages: { items: IAiChatMessageRow[]; meta: unknown }[]; + pageParams: unknown[]; + } + | undefined, + ) => + old + ? { ...old, pages: mergeDeltaRowsIntoPages(old.pages, res.rows) } + : old, + ); + } + } catch { + // Transient failure (e.g. a server restart mid-run): swallow and retry on + // the next tick — the poll must survive a bounce, like the old dumb refetch. + } + }; + const id = setInterval(() => void tick(), 2500); + return () => { + cancelled = true; + clearInterval(id); + }; + }, [degradedPoll, windowOpen, activeChatId, queryClient]); + // #184 reconnect-and-live-follow. Whether detached agent runs are enabled for // this workspace. When the feature is off no runs are ever created, so the // resume attempt would only ever 204; gating ChatThread's resume on it avoids a diff --git a/apps/client/src/features/ai-chat/components/chat-thread.test.tsx b/apps/client/src/features/ai-chat/components/chat-thread.test.tsx index ba4dcb7c..5fbcebaa 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.test.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.test.tsx @@ -172,9 +172,18 @@ function resetState() { h.state.getRun.mockResolvedValue({ run: null, message: null }); } +// #491: the streaming tail carries a persisted step frontier (metadata.stepsPersisted), +// which the tail-only attach reads as `n` in `?anchor=&n=`. Seeded WHOLE now. const streamingTail = () => [ row("u1", "user", undefined, "hi"), - row("a1", "assistant", "streaming", "partial"), + { + id: "a1", + role: "assistant", + content: "partial", + status: "streaming", + createdAt: "2026-01-01T00:00:00Z", + metadata: { stepsPersisted: 2 }, + } as IAiChatMessageRow, ]; const settledTail = () => [ row("u1", "user", undefined, "hi"), @@ -335,20 +344,24 @@ describe("ChatThread — send now", () => { expect(screen.getAllByLabelText("Remove queued message")).toHaveLength(1); }); - it("Stop then a REAL network-drop finish exits to idle (honor-in-stopping), NOT a false reconnect", () => { + it("Stop then a REAL network-drop finish exits to idle (honor-in-stopping), NOT a false reconnect", async () => { // Regression for the disconnect-first reorder: on the STOP path, even a drop- // form finish { isError:true, isDisconnect:true } arriving in `stopping` must be // HONORED (reducer) and exit to idle — it must NOT enter the reconnect ladder. startLocalStreamWithRun(); // live local stream, autonomous fireEvent.click(screen.getByLabelText("Stop")); // STOP_REQUESTED -> stopping h.state.error = { message: "Failed to fetch" }; - act(() => { + // #491: the disconnect re-seeds from persist (async getRun) before dispatching + // FINISH_DISCONNECT, which the reducer HONORS in `stopping` -> idle. Flush it. + await act(async () => { h.state.onFinish?.({ message: { id: "a1", role: "assistant", parts: [] }, isAbort: false, isDisconnect: true, isError: true, }); + await Promise.resolve(); + await Promise.resolve(); }); expect(screen.queryByText(/reconnecting/i)).toBeNull(); }); @@ -803,19 +816,24 @@ describe("ChatThread — resume (attach) machinery", () => { expect(h.state.resumeStream).not.toHaveBeenCalled(); }); - it("strips the streaming tail from the seed, keeps a user tail whole", () => { + it("#491 tail-only: seeds the streaming tail WHOLE (no strip), keeps a user tail whole", () => { renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() }); - expect(h.state.seededMessages).toHaveLength(1); + // MUTATION-VERIFY: re-introduce the seed-strip and this goes red — the streaming + // tail (steps 0..N-1) MUST be seeded so the SDK continuation appends the tail to + // the RIGHT message. Both rows (user + assistant) are seeded. + expect(h.state.seededMessages).toHaveLength(2); cleanup(); resetState(); renderThread({ autonomousRunsEnabled: true, initialRows: userTail() }); expect(h.state.seededMessages).toHaveLength(1); }); - it("builds the attach URL with expect=live&anchor only for a stripped streaming tail", () => { + it("#491 tail-only: builds the attach URL with ?anchor=&n= from the persisted step frontier", () => { renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() }); + // n=2 comes from a1's metadata.stepsPersisted (MUTATION-VERIFY: hardcode n=0 and + // this fails). No `expect=live` param anymore. expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe( - "/api/ai-chat/runs/c1/stream?expect=live&anchor=a1", + "/api/ai-chat/runs/c1/stream?anchor=a1&n=2", ); cleanup(); resetState(); @@ -839,39 +857,41 @@ describe("ChatThread — resume (attach) machinery", () => { }); } - it("204 on a streaming tail: restore + invalidate + onResumeFallback(true)", async () => { + it("204 on a streaming tail: NO restore (row kept) + invalidate + onResumeFallback(true)", async () => { const { onResumeFallback, invalidateSpy } = renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail(), }); await attachFetch({ status: 204, ok: false }); - expect(h.state.setMessages).toHaveBeenCalledTimes(1); // restore + // #491 tail-only: the anchor row was never stripped, so there is NOTHING to + // restore. MUTATION-VERIFY: re-add a restore setMessages here and it goes red. + expect(h.state.setMessages).not.toHaveBeenCalled(); expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ["ai-chat-messages", "c1"], }); expect(onResumeFallback).toHaveBeenCalledWith(true); }); - it("F7 restart-survival: a 500 attach failure restores the row AND arms the poll", async () => { + it("F7 restart-survival: a 500 attach failure arms the poll WITHOUT a restore", async () => { const { onResumeFallback, invalidateSpy } = renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail(), }); await attachFetch({ status: 500, ok: false }); - expect(h.state.setMessages).toHaveBeenCalledTimes(1); + expect(h.state.setMessages).not.toHaveBeenCalled(); expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ["ai-chat-messages", "c1"], }); expect(onResumeFallback).toHaveBeenCalledWith(true); }); - it("F7 restart-survival: a network throw restores the row AND arms the poll", async () => { + it("F7 restart-survival: a network throw arms the poll WITHOUT a restore", async () => { const { onResumeFallback, invalidateSpy } = renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail(), }); await attachFetch(new Error("network down"), true); - expect(h.state.setMessages).toHaveBeenCalledTimes(1); + expect(h.state.setMessages).not.toHaveBeenCalled(); expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ["ai-chat-messages", "c1"], }); @@ -931,7 +951,7 @@ describe("ChatThread — resume (attach) machinery", () => { expect(h.state.sendMessage).not.toHaveBeenCalled(); }); - it("an empty resumed message (starved replay) restores the row AND arms the poll", () => { + it("an empty resumed message (starved replay) arms the poll WITHOUT a restore", () => { h.state.status = "ready"; const { onResumeFallback } = renderThread({ autonomousRunsEnabled: true, @@ -947,7 +967,9 @@ describe("ChatThread — resume (attach) machinery", () => { isError: false, }); }); - expect(h.state.setMessages).toHaveBeenCalledTimes(1); // restore + // #491 tail-only: the seeded steps 0..N-1 are still on screen (the SDK + // continuation never wiped them), so there is nothing to restore — just poll. + expect(h.state.setMessages).not.toHaveBeenCalled(); expect(onResumeFallback).toHaveBeenCalledWith(true); // arm }); @@ -995,24 +1017,41 @@ describe("ChatThread — live reconnect + stalled", () => { cleanup(); }); + // #491: the authoritative PERSISTED assistant row `getRun` projects on a local + // disconnect — the re-seed source. Its metadata.stepsPersisted becomes `n`. + const persistedAnchor = (steps = 3) => ({ + run: { id: "run-1", status: "running" }, + message: { + id: "a2", + role: "assistant", + content: "persisted 0..N-1", + status: "streaming", + createdAt: "2026-01-01T00:00:00Z", + metadata: { stepsPersisted: steps }, + }, + }); + // A REAL live SSE drop. ai@6.0.207 emits BOTH { isError:true, isDisconnect:true } - // for a network TypeError AND sets useChat `error` — NOT the { isError:false, - // error:null } form the old tests fed. This is the form browser QA hit; with the - // buggy isError-first routing OR without the errorView render-gate these tests go - // red (a real drop surfaces the terminal error banner, masking the reconnect - // ladder). MUTATION-VERIFY of disconnect-first + the errorView phase-gate. - function disconnect(message: unknown = liveMsg) { + // for a network TypeError AND sets useChat `error`. #491: an autonomous local drop + // now RE-SEEDS from persist (async getRun) BEFORE entering the reconnect ladder, so + // this helper is async and flushes the getRun microtask before returning. + async function disconnect(message: unknown = liveMsg) { h.state.error = { message: "Failed to fetch" }; // the SDK sets error on the drop - act(() => { + await act(async () => { h.state.onFinish?.({ message, isAbort: false, isDisconnect: true, isError: true, }); + // Flush the getRun().then re-seed + the deferred FINISH_DISCONNECT dispatch. + await Promise.resolve(); + await Promise.resolve(); }); } function renderLive() { + // The persisted-anchor read the local disconnect performs to re-seed from persist. + h.state.getRun.mockResolvedValue(persistedAnchor()); const view = renderThread({ autonomousRunsEnabled: true, initialRows: settledTail(), @@ -1032,35 +1071,43 @@ describe("ChatThread — live reconnect + stalled", () => { }); } - it("a live disconnect starts a backoff reconnect (banner + resumeStream after backoff)", () => { + it("#491: a live disconnect RE-SEEDS from persist, then backs off to reconnect with ?anchor=&n=", async () => { renderLive(); - disconnect(); + await disconnect(); + // The re-seed read the authoritative persisted row and replaced the live partial. + // MUTATION-VERIFY: skip the getRun re-seed (send `n` off the live message) and the + // n below no longer matches the PERSISTED stepsPersisted. + expect(h.state.getRun).toHaveBeenCalledWith("c1"); + expect(h.state.setMessages).toHaveBeenCalled(); // re-seeded the store from persist expect(screen.getByText(/reconnecting/i)).toBeTruthy(); expect(h.state.resumeStream).not.toHaveBeenCalled(); advanceToAttempt(1); expect(h.state.resumeStream).toHaveBeenCalledTimes(1); + // n=3 is the PERSISTED row's stepsPersisted (from getRun), NOT the live store. expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe( - "/api/ai-chat/runs/c1/stream?expect=live&anchor=a2", + "/api/ai-chat/runs/c1/stream?anchor=a2&n=3", ); }); - it("#488 (browser QA): the reconnect banner is SHOWN, not masked by the residual useChat error", () => { + it("#488 (browser QA): the reconnect banner is SHOWN, not masked by the residual useChat error", async () => { // The drop sets useChat `error` (real SDK), and the terminal errorView describes // it ("Lost connection to the server"). The FSM phase-gate must let the // `reconnecting` banner WIN over that residual error. MUTATION-VERIFY: revert the // errorView phase-gate (show errorView whenever error is set) and the terminal // banner masks "reconnecting…" -> red. renderLive(); - disconnect(); + await disconnect(); expect(h.state.error).not.toBeNull(); // the SDK error IS set during recovery expect(screen.getByText(/reconnecting/i)).toBeTruthy(); // The terminal "Lost connection… reload" banner must NOT be showing. expect(screen.queryByText(/reload and try again/i)).toBeNull(); }); - it("#488 commit 2: a disconnect BEFORE the first assistant frame reconnects with NO anchor", () => { + it("#488 commit 2: a disconnect BEFORE the first assistant frame reconnects with NO anchor", async () => { renderLive(); - disconnect(null); // no assistant message yet (pre-first-frame break) + // No persisted assistant row for a pre-first-frame break -> no anchor. + h.state.getRun.mockResolvedValue({ run: null, message: null }); + await disconnect(null); // no assistant message yet (pre-first-frame break) expect(screen.getByText(/reconnecting/i)).toBeTruthy(); expect( screen.queryByText("Connection lost — the answer was interrupted."), @@ -1074,7 +1121,7 @@ describe("ChatThread — live reconnect + stalled", () => { it("a live re-attach (2xx) clears the reconnect banner", async () => { renderLive(); - disconnect(); + await disconnect(); advanceToAttempt(1); await reconnect({ status: 200, ok: true }); expect(screen.queryByText(/reconnecting/i)).toBeNull(); @@ -1082,7 +1129,7 @@ describe("ChatThread — live reconnect + stalled", () => { it("a 204 arms the degraded poll and backs off to the next attempt", async () => { const { onResumeFallback } = renderLive(); - disconnect(); + await disconnect(); advanceToAttempt(1); expect(h.state.resumeStream).toHaveBeenCalledTimes(1); await reconnect({ status: 204, ok: false }); @@ -1094,7 +1141,7 @@ describe("ChatThread — live reconnect + stalled", () => { it("exhausts the attempt limit into a manual Retry, which restarts the sequence", async () => { renderLive(); - disconnect(); + await disconnect(); for (let n = 1; n <= 5; n++) { advanceToAttempt(n); expect(h.state.resumeStream).toHaveBeenCalledTimes(n); @@ -1112,22 +1159,23 @@ describe("ChatThread — live reconnect + stalled", () => { it("#488 commit 3: two breaks in a row produce two reconnect cycles", async () => { renderLive(); // First break -> reconnect -> re-attach live. - disconnect(); + await disconnect(); advanceToAttempt(1); expect(h.state.resumeStream).toHaveBeenCalledTimes(1); await reconnect({ status: 200, ok: true }); expect(screen.queryByText(/reconnecting/i)).toBeNull(); - // The re-attached observer stream drops AGAIN -> a SECOND reconnect cycle - // (the old one-shot !wasResumed gate sent this to silent poll). - disconnect(); + // The re-attached observer (live-follow) stream drops AGAIN -> a SECOND reconnect + // cycle. #491: this too re-seeds from persist before re-attaching (never tail- + // applies over the live-follow partial). + await disconnect(); expect(screen.getByText(/reconnecting/i)).toBeTruthy(); advanceToAttempt(1); expect(h.state.resumeStream).toHaveBeenCalledTimes(2); }); - it("does NOT reconnect when autonomous runs are disabled", () => { + it("does NOT reconnect when autonomous runs are disabled", async () => { renderThread({ autonomousRunsEnabled: false, initialRows: settledTail() }); - disconnect(); + await disconnect(); expect(screen.queryByText(/reconnecting/i)).toBeNull(); expect( screen.getByText("Connection lost — the answer was interrupted."), @@ -1138,7 +1186,7 @@ describe("ChatThread — live reconnect + stalled", () => { it("#488 commit 4a: the poll idle cap surfaces a stalled banner + Retry (not silent)", async () => { renderLive(); - disconnect(); + await disconnect(); advanceToAttempt(1); await reconnect({ status: 204, ok: false }); // arms the poll (reconnecting) // No activity for the whole idle cap -> stalled. diff --git a/apps/client/src/features/ai-chat/components/chat-thread.tsx b/apps/client/src/features/ai-chat/components/chat-thread.tsx index e213da32..01e703ce 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.tsx @@ -42,7 +42,7 @@ import { assistantMessageHasVisibleContent } from "@/features/ai-chat/utils/mess import { isStreamingTail, isSettledAssistantTail, - seedRows, + stepsPersistedOf, mergeById, } from "@/features/ai-chat/utils/resume-helpers.ts"; import { getRun } from "@/features/ai-chat/services/ai-chat-service.ts"; @@ -266,25 +266,33 @@ export default function ChatThread({ // is NOT one of the lifecycle flags the FSM replaced. const mountedRef = useRef(true); - // attachStrategy DATA (behind the resumeStream effect; #491 swaps it to tail-only - // WITHOUT touching the FSM). The controller is effect-owned (aborted in cleanup, - // I5). `stripRef`/`strippedRowRef` are the current full-replay+strip anchor. + // attachStrategy DATA (behind the resumeStream effect; #491 tail-only, WITHOUT + // touching the FSM). The controller is effect-owned (aborted in cleanup, I5). + // `anchorRef` is the PERSISTED assistant row that pins the run (server invariant + // 6) and its persisted step frontier N: it feeds `?anchor=&n=` + // so the tail-only attach returns frames for steps >= N (the seed carries 0..N-1). + // It is NOT a "stripped" row — the seed keeps every row (tail-only replaces the + // old full-replay+strip). Null when there is no streaming/active tail to resume. const attachAbortRef = useRef(null); - const stripRef = useRef(chatId !== null && isStreamingTail(initialRows ?? [])); - const strippedRowRef = useRef( - stripRef.current ? (initialRows ?? [])[initialRows!.length - 1] : null, + const anchorRef = useRef<{ id: string; stepsPersisted: number } | null>( + (() => { + if (chatId === null || !isStreamingTail(initialRows ?? [])) return null; + const rows = initialRows ?? []; + const tail = rows[rows.length - 1]; + return { id: tail.id, stepsPersisted: stepsPersistedOf(tail) }; + })(), ); // Effect-owned backoff timers (not lifecycle flags): the reconnect ladder and the // stalled inactivity cap. Cleared by the cancelReconnect effect / the cap effect. const reconnectTimerRef = useRef | null>(null); const idleCapTimerRef = useRef | null>(null); + // #491 tail-only: seed EVERY persisted row unchanged (no strip). The streaming + // tail holds steps 0..N-1; the run-stream registry's tail (steps >= N) is APPENDED + // to it by the SDK continuation (readUIMessageStream({ message })), so it must be + // present in the store for the attach to continue the RIGHT message. const initialMessages = useMemo( - () => - seedRows( - initialRows ?? [], - stripRef.current && autonomousRunsEnabled === true, - ).map(rowToUiMessage), + () => (initialRows ?? []).map(rowToUiMessage), [initialRows], ); @@ -335,21 +343,16 @@ export default function ChatThread({ (eff: RunEffect, epoch: number) => { switch (eff.type) { case "resumeStream": { - // The attach GET. Stamp the outcome's generation (I1). A reconnect - // attempt filters the pinned live row from the store first (the mount - // seed already stripped it), so the live replay's text-start rebuilds it - // without duplicating parts (#430). + // The attach GET. Stamp the outcome's generation (I1). #491 tail-only: the + // store already holds EXACTLY the persisted steps 0..N-1 (the mount seed IS + // persist; a reconnect was re-seeded from persist BEFORE FINISH_DISCONNECT + // scheduled it — see the onFinish disconnect handler), so there is nothing + // to filter here: the SDK continues that seeded message, appending the tail + // (steps >= N) without duplicating the pre-drop partial step. pendingAttachEpochRef.current = epoch; // The resumed stream's onFinish is stamped with THIS attach generation // (F1), so a superseded attempt's late finish is dropped. turnEpochRef.current = epoch; - if (machineRef.current.phase.name === "reconnecting") { - const anchor = strippedRowRef.current; - if (anchor) - setMessagesRef.current?.((prev) => - prev.filter((m) => m.id !== anchor.id), - ); - } void resumeStreamRef.current?.(); break; } @@ -464,18 +467,23 @@ export default function ChatThread({ new DefaultChatTransport({ api: "/api/ai-chat/stream", credentials: "include", - prepareReconnectToStreamRequest: () => ({ - // Build the attach URL from the REAL chat id. ?expect=live&anchor= - // only when a streaming tail was stripped: expect=live opts into a - // finished-retained replay (safe only because the row is stripped and the - // replay rebuilds it), and the anchor pins the replay to OUR run — a - // mismatching (newer) run 204s into the restore+poll path instead. - api: `/api/ai-chat/runs/${chatIdRef.current}/stream${ - stripRef.current - ? `?expect=live&anchor=${strippedRowRef.current!.id}` - : "" - }`, - }), + prepareReconnectToStreamRequest: () => { + // #491 tail-only attach URL. When there is an anchor (a streaming/active + // tail to resume) build `?anchor=&n=`: the + // server returns the TAIL — a synthetic `start` frame + frames for steps + // >= n, then live — which the SDK continuation appends to the seeded row. + // The server 204s (-> restore-noop + poll) when it cannot cover the + // frontier (overflow/rotation gap) or the anchor mismatches (a newer run). + // No anchor (a user tail / pre-first-frame break) => no params. + const anchor = anchorRef.current; + return { + api: `/api/ai-chat/runs/${chatIdRef.current}/stream${ + anchor + ? `?anchor=${anchor.id}&n=${anchor.stepsPersisted}` + : "" + }`, + }; + }, fetch: async (input: RequestInfo | URL, init: RequestInit = {}) => { if ((init.method ?? "GET") !== "GET") { // Send path (POST). #488 commit 5: NO client 409 retry ladder anymore @@ -562,8 +570,9 @@ export default function ChatThread({ // Attach GET outcome -> FSM event. The epoch guard replaces BOTH the one-shot // 204 guard (noStreamHandledRef) and the unmount gate: a stale/superseded or - // post-DISPOSE outcome is dropped (I1). For a NONE outcome the attachStrategy - // recovery (restore the stripped row + invalidate for a fresh poll) runs first. + // post-DISPOSE outcome is dropped (I1). #491 tail-only: on a NONE outcome there is + // NOTHING to restore — the anchor row was never stripped from the view (the seed + // keeps it) — so we only invalidate for a fresh poll + dispatch the FSM event. const handleAttachOutcome = useCallback( (ep: number, wasReconnecting: boolean, live: boolean) => { if (ep !== epochRef.current) return; // stale generation — drop @@ -575,10 +584,6 @@ export default function ChatThread({ ); return; } - if (strippedRowRef.current) - setMessagesRef.current?.((prev) => - mergeById(prev, rowToUiMessage(strippedRowRef.current!)), - ); queryClient.invalidateQueries({ queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current), }); @@ -661,56 +666,31 @@ export default function ChatThread({ // keeps executing server-side — must win; only a NON-disconnect error (a // provider 500, `{ isError:true, isDisconnect:false }`) is terminal. if (isDisconnect) { - if (wasObserver) { - // A resumed/attached OBSERVER stream dropped. Recover via the degraded - // poll (restore the stripped row only when there is no visible content; - // never clobber a fuller on-screen tail, invariant 9). The FSM decides - // reconnect-vs-poll from liveFollow (a live-follow drop reconnects again, - // #488 commit 3; a mount-resume drop polls). - if (mountedRef.current) { - const hasVisible = msgHasVisible; - if (!hasVisible && strippedRowRef.current) - setMessages((prev) => - mergeById(prev, rowToUiMessage(strippedRowRef.current!)), - ); - queryClient.invalidateQueries({ - queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current), - }); - dispatch({ - type: "FINISH_DISCONNECT", - hasVisibleContent: hasVisible, - epoch: stampEpoch, - }); - } + if (!mountedRef.current) { setStopNotice(null); return; } - // A LOCAL live turn dropped. #488 commit 2: recover by the RUN-FACT, not by - // the presence of an assistant message — a setup-phase break (before the - // first frame) still leaves a detached run writing to pages. In autonomous - // mode a run is active for the whole turn, so seed the run-fact from the - // start-metadata runId when known, else a sentinel (the attach GET goes by - // chatId, not runId). Pin the assistant row as the strip/anchor when present. - if (autonomousRunsEnabled === true && mountedRef.current) { - const hasAnchor = - message?.role === "assistant" && typeof message.id === "string"; - if (hasAnchor) { - strippedRowRef.current = { - id: message.id, - role: "assistant", - content: "", - status: "streaming", - createdAt: new Date().toISOString(), - metadata: { parts: message.parts }, - }; - stripRef.current = true; - } else { - strippedRowRef.current = null; - stripRef.current = false; - } + // No detached run to recover (legacy, non-autonomous): a plain disconnect — + // terminal notice, no reconnect. (An observer only exists in autonomous mode, + // so this is always a local turn.) + if (autonomousRunsEnabled !== true) { dispatch({ - type: "RUN_FACT", - runFact: { runId: extractRunId(message) ?? "pending" }, + type: "FINISH_DISCONNECT", + hasVisibleContent: false, + epoch: stampEpoch, + }); + setStopNotice("disconnect"); + return; + } + // A mount-resume OBSERVER (one-shot resume, NOT live-follow) drop falls to + // the degraded POLL, which merges by id — it does NOT attach, so there is + // nothing to re-seed. #491 tail-only: the anchor row was never removed from + // the view (the seed keeps it; the continuation only APPENDED), so nothing to + // restore either. The FSM routes this to `polling` (ownership observer, + // !liveFollow). + if (wasObserver && !machineRef.current.ctx.liveFollow) { + queryClient.invalidateQueries({ + queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current), }); dispatch({ type: "FINISH_DISCONNECT", @@ -718,14 +698,66 @@ export default function ChatThread({ epoch: stampEpoch, }); setStopNotice(null); - } else { + return; + } + // We will (re-)ENTER THE RECONNECT LADDER (an attach): a LOCAL live turn's + // first drop, OR a live-follow observer's SUBSEQUENT drop (#488 commit 3). + // #488 commit 2: recover by the RUN-FACT, not by the presence of an assistant + // message — a setup-phase break still leaves a detached run writing to pages. + // + // #491 tail-only (THE crux): the live store holds a PARTIAL step that is AHEAD + // of the persisted boundary; tail-applying the reconnect's step frames over it + // would DUPLICATE that partial step. So entering reconnecting is ALWAYS via a + // RE-SEED FROM PERSIST — never the live store. Fetch the authoritative + // persisted assistant row (`getRun` returns the projected `message`), replace + // the live partial by id (mergeById -> the store now holds EXACTLY steps + // 0..N-1), and set the anchor to `{ id, n = stepsPersisted }`. Only AFTER the + // re-seed is applied do we enter the ladder (FINISH_DISCONNECT schedules the + // backoff) — so the attach can never tail-apply over the live partial. + const cid = chatIdRef.current; + // The live-message runId is the run-fact source (the attach GET keys on + // chatId, so a sentinel still recovers a setup-phase break). + const runId = extractRunId(message ?? undefined) ?? "pending"; + const enterReconnect = (fact: string): void => { + if (!mountedRef.current) return; + dispatch({ type: "RUN_FACT", runFact: { runId: fact } }); dispatch({ type: "FINISH_DISCONNECT", - hasVisibleContent: false, + hasVisibleContent: msgHasVisible, epoch: stampEpoch, }); - setStopNotice("disconnect"); + }; + if (cid) { + void getRun(cid) + .then((res) => { + if (!mountedRef.current) return; + const persisted = res.message; + if (persisted && persisted.role === "assistant") { + anchorRef.current = { + id: persisted.id, + stepsPersisted: stepsPersistedOf(persisted), + }; + // Replace the live partial with the persisted row IN PLACE by id — + // the re-seed from persist. + setMessages((prev) => mergeById(prev, rowToUiMessage(persisted))); + } else { + // No persisted assistant row (pre-first-frame break): no anchor; the + // reconnect attaches with no params and replays from start. + anchorRef.current = null; + } + enterReconnect(res.run?.id ?? runId); + }) + .catch(() => { + // Persist read failed: keep whatever anchor we had and still enter the + // ladder so a live detached run is not stranded (the attach 204s -> + // degraded poll if it cannot cover). + enterReconnect(runId); + }); + } else { + anchorRef.current = null; + enterReconnect(runId); } + setStopNotice(null); return; } // A NON-disconnect stream error (a provider 500 etc.) -> terminal error banner. @@ -746,11 +778,10 @@ export default function ChatThread({ if (mountedRef.current) { const hasVisible = msgHasVisible; if (!hasVisible) { - // Starved replay: restore the stripped row + poll to the real terminal. - if (strippedRowRef.current) - setMessages((prev) => - mergeById(prev, rowToUiMessage(strippedRowRef.current!)), - ); + // Starved replay (the tail carried no new steps). #491 tail-only: the + // seeded steps 0..N-1 are still on screen (the SDK continuation never + // wiped them — `start` does not reset parts), so there is nothing to + // restore; just poll to the real terminal. queryClient.invalidateQueries({ queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current), }); @@ -863,12 +894,12 @@ export default function ChatThread({ const tail = rows[rows.length - 1]; if (!tail || tail.role !== "assistant") return; setMessages((prev) => mergeById(prev, rowToUiMessage(tail))); - // Anchor-mismatch coherence: a restored stripped row A that a DIFFERENT run's - // row B has replaced as the tail would linger as an orphan — settle A from - // fresh history so no phantom row survives. - const stripped = strippedRowRef.current; - if (stripped && stripped.id !== tail.id) { - const historical = rows.find((r) => r.id === stripped.id); + // Anchor-mismatch coherence: if a DIFFERENT run's row B has replaced our anchor + // row A as the tail, A would linger as an orphan — reconcile A by id from FRESH + // PERSISTED history (not the pinned live row) so no phantom row survives. + const anchor = anchorRef.current; + if (anchor && anchor.id !== tail.id) { + const historical = rows.find((r) => r.id === anchor.id); if (historical) setMessages((prev) => mergeById(prev, rowToUiMessage(historical))); } diff --git a/apps/client/src/features/ai-chat/utils/resume-helpers.test.ts b/apps/client/src/features/ai-chat/utils/resume-helpers.test.ts index 9278d004..42c0d1b6 100644 --- a/apps/client/src/features/ai-chat/utils/resume-helpers.test.ts +++ b/apps/client/src/features/ai-chat/utils/resume-helpers.test.ts @@ -4,7 +4,8 @@ import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.t import { isStreamingTail, isSettledAssistantTail, - seedRows, + stepsPersistedOf, + mergeDeltaRowsIntoPages, mergeById, } from "./resume-helpers.ts"; @@ -12,8 +13,18 @@ function row( id: string, role: string, status?: string, + stepsPersisted?: number, ): IAiChatMessageRow { - return { id, role, content: "", status, createdAt: "2026-01-01T00:00:00Z" }; + return { + id, + role, + content: "", + status, + createdAt: "2026-01-01T00:00:00Z", + ...(stepsPersisted !== undefined + ? { metadata: { stepsPersisted } } + : {}), + }; } function makeMsg(id: string, text: string): UIMessage { @@ -65,23 +76,92 @@ describe("isSettledAssistantTail", () => { }); }); -describe("seedRows", () => { - const rows = [row("u1", "user"), row("a1", "assistant", "streaming")]; - - it("returns the rows unchanged when not stripping", () => { - expect(seedRows(rows, false)).toBe(rows); +describe("stepsPersistedOf", () => { + it("reads metadata.stepsPersisted", () => { + expect(stepsPersistedOf(row("a1", "assistant", "streaming", 3))).toBe(3); + expect(stepsPersistedOf(row("a1", "assistant", "streaming", 0))).toBe(0); }); - it("drops the last row when stripping", () => { - const seeded = seedRows(rows, true); - expect(seeded).toHaveLength(1); - expect(seeded[0].id).toBe("u1"); + it("defaults to 0 for a pre-#491 row (absent), null/undefined, or a bad value", () => { + expect(stepsPersistedOf(row("a1", "assistant", "streaming"))).toBe(0); + expect(stepsPersistedOf(null)).toBe(0); + expect(stepsPersistedOf(undefined)).toBe(0); + expect( + stepsPersistedOf({ + id: "a1", + role: "assistant", + content: "", + createdAt: "x", + metadata: { stepsPersisted: -2 }, + }), + ).toBe(0); }); - it("returns an empty list when stripping a single-row list", () => { - expect(seedRows([row("a1", "assistant", "streaming")], true)).toHaveLength( - 0, - ); + it("floors a non-integer count", () => { + expect( + stepsPersistedOf({ + id: "a1", + role: "assistant", + content: "", + createdAt: "x", + metadata: { stepsPersisted: 2.9 }, + }), + ).toBe(2); + }); +}); + +describe("mergeDeltaRowsIntoPages", () => { + const pages = () => [ + { items: [row("u1", "user"), row("a1", "assistant", "streaming", 1)], meta: {} }, + ]; + + it("returns the pages unchanged for an empty delta", () => { + const p = pages(); + expect(mergeDeltaRowsIntoPages(p, [])).toBe(p); + }); + + it("appends a genuinely new row to the last page in chronological order", () => { + const merged = mergeDeltaRowsIntoPages(pages(), [row("a2", "assistant", "streaming", 0)]); + expect(merged[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]); + }); + + it("replaces a grown row in place (per-step growth), never appends a duplicate", () => { + const merged = mergeDeltaRowsIntoPages(pages(), [ + row("a1", "assistant", "streaming", 2), + ]); + expect(merged[0].items.map((i) => i.id)).toEqual(["u1", "a1"]); + // the in-place replacement carries the grown step frontier. + expect(stepsPersistedOf(merged[0].items[1])).toBe(2); + }); + + it("does not mutate the input pages", () => { + const input = pages(); + const before = input[0].items.slice(); + mergeDeltaRowsIntoPages(input, [row("a2", "assistant", "streaming", 0)]); + expect(input[0].items).toEqual(before); // untouched + }); + + // #491 CONTRACT: the delta overlap window re-delivers the same rows, so merging + // MUST be idempotent — applying a delta twice equals applying it once (no growth, + // no reorder). A regression re-introduces duplicate assistant bubbles per poll. + it("is idempotent: applying the same delta twice equals once", () => { + const delta = [ + row("a1", "assistant", "streaming", 2), // grown existing row + row("a2", "assistant", "streaming", 0), // new row + ]; + const once = mergeDeltaRowsIntoPages(pages(), delta); + const twice = mergeDeltaRowsIntoPages(once, delta); + const thrice = mergeDeltaRowsIntoPages(twice, delta); + expect(once[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]); + expect(twice[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]); + expect(twice).toEqual(once); + expect(thrice).toEqual(once); + }); + + it("seeds a first page when the cache is empty", () => { + const merged = mergeDeltaRowsIntoPages([], [row("u1", "user")]); + expect(merged).toHaveLength(1); + expect(merged[0].items.map((i) => i.id)).toEqual(["u1"]); }); }); diff --git a/apps/client/src/features/ai-chat/utils/resume-helpers.ts b/apps/client/src/features/ai-chat/utils/resume-helpers.ts index bddb4d01..8c21f42b 100644 --- a/apps/client/src/features/ai-chat/utils/resume-helpers.ts +++ b/apps/client/src/features/ai-chat/utils/resume-helpers.ts @@ -11,9 +11,10 @@ import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.t /** * A STREAMING tail: the last persisted row is an assistant row still marked - * `status === 'streaming'`. Such a tail is stripped from the seed and rebuilt by - * the replay (`expect=live`), since the SDK's `text-start` always pushes a new - * part and replaying over a seeded in-progress row would duplicate its text. + * `status === 'streaming'`. #491 (tail-only): such a tail is seeded UNCHANGED — + * it carries the persisted steps 0..N-1 — and the run-stream registry's tail + * (frames for steps >= N) is APPENDED to it by the SDK's `readUIMessageStream` + * continuation. Only the presence of this tail decides WHETHER to attach. */ export function isStreamingTail(rows: IAiChatMessageRow[]): boolean { const tail = rows[rows.length - 1]; @@ -32,15 +33,61 @@ export function isSettledAssistantTail(rows: IAiChatMessageRow[]): boolean { } /** - * Seed rows for `useChat`: return the rows unchanged, or without the last row when - * `strip` is set (the streaming tail is stripped so the live replay rebuilds it - * without duplicating parts). + * #491 tail-only anchor: the count of FINISHED steps whose parts are persisted in + * THIS assistant row (`metadata.stepsPersisted`), written atomically with `parts` + * server-side. The resume client reads it as its persisted step frontier N — the + * tail-only attach asks the run-stream registry for the frames of step N onward + * (the seed already carries steps 0..N-1). Absent on pre-#491 rows => 0. */ -export function seedRows( +export function stepsPersistedOf( + row: IAiChatMessageRow | null | undefined, +): number { + const n = row?.metadata?.stepsPersisted; + return typeof n === "number" && n >= 0 ? Math.floor(n) : 0; +} + +/** One page of the messages infinite-query cache (`{ items, meta }`). */ +export interface IMessagePage { + items: IAiChatMessageRow[]; + meta: unknown; +} + +/** + * #491 delta-poll merge: upsert the delta poll's `rows` into the messages + * infinite-query page structure IDEMPOTENTLY by id. The delta endpoint's overlap + * window GUARANTEES occasional REPEATS, so this MUST converge: a row already + * present is REPLACED IN PLACE (per-step growth of an in-progress row), a new row + * is APPENDED to the last page in chronological order (the server returns delta + * rows oldest-first). Applying the same delta twice equals applying it once. Never + * mutates the input pages (returns fresh page objects with cloned item arrays). + */ +export function mergeDeltaRowsIntoPages( + pages: IMessagePage[], rows: IAiChatMessageRow[], - strip: boolean, -): IAiChatMessageRow[] { - return strip ? rows.slice(0, -1) : rows; +): IMessagePage[] { + if (rows.length === 0) return pages; + const next: IMessagePage[] = pages.map((p) => ({ + ...p, + items: p.items.slice(), + })); + const locate = (id: string): [number, number] | null => { + for (let pi = 0; pi < next.length; pi++) { + const ii = next[pi].items.findIndex((it) => it.id === id); + if (ii !== -1) return [pi, ii]; + } + return null; + }; + for (const row of rows) { + const at = locate(row.id); + if (at) { + next[at[0]].items[at[1]] = row; // replace in place — idempotent by id + } else if (next.length > 0) { + next[next.length - 1].items.push(row); // append chronologically + } else { + next.push({ items: [row], meta: undefined }); + } + } + return next; } /** -- 2.52.0 From a0ecf21cb5a061d2edd8132c10642d63fac59954 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 22:50:09 +0300 Subject: [PATCH 6/7] =?UTF-8?q?fix(ai-chat):=20=D1=80=D0=B5-=D1=80=D0=B5?= =?UTF-8?q?=D0=B2=D1=8C=D1=8E=20#491=20=E2=80=94=20=D0=B4=D1=83=D0=B1?= =?UTF-8?q?=D0=BB=D1=8C=20=D0=BD=D0=B0=20getRun-fail=20+=20epoch=20RUN=5FF?= =?UTF-8?q?ACT=20+=20doc=20(#491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ре-ревью нашло регрессию класса #137/#161 на пути отказа getRun при tail-only re-seed: - НАХОДКА 1 (MEDIUM/HIGH): в onFinish-disconnect ветка .catch (отказ getRun) на локальном дропе входила в reconnect-ladder БЕЗ ре-сида и БЕЗ фильтра «живой» частичной строки. anchorRef оставался устаревшим (mount-инициализатор), живая частичная строка со шагами N..M-1 — последней; через ~1с реконнект строил ?anchor=&n=N_mount, и при живом ране с покрытием от N_mount (flaky-сеть: SSE и getRun упали, сеть поднялась за 1с) сервер отдавал кадры ≥N_mount → SDK дописывал их к строке, где они УЖЕ есть → дублирование (клиентского дедупа реплея против parts нет). Фикс: восстановлена структурная гарантия удалённого resumeStream- фильтра — на отказе getRun (и на no-persisted-row, и на no-cid) живая частичная строка удаляется из стора по id + anchorRef=null → реконнект реплеит со start в ЧИСТЫЙ стор (полная пересборка) либо 204→poll. Нет пути, где attach tail-applies на строку с уже присутствующими шагами. Тест: getRun-reject на локальном дисконнекте → живая строка отфильтрована + URL без параметров (mutation-verify: без фикса тест краснеет — фильтр не срабатывает). - НАХОДКА 2 (LOW): RUN_FACT в enterReconnect теперь epoch-штампуется (epoch: stampEpoch), как везде (postRun): getRun-rtt расширяет окно onFinish→dispatch, конкурентный SEND_LOCAL во время rtt теперь дропает устаревший RUN_FACT по I1, а не перетирает runFact.runId нового хода. - НАХОДКА 3 (LOW, doc): run-fsm.spec.md обновлён — stripRef/strippedRowRef → anchorRef {id, stepsPersisted}, tail-only + re-seed-from-persist. FSM run-fsm.ts не тронут; инварианты #488 (epoch/honor-in-stopping/ownership-reset/ disconnect-first/render-gate) сохранены. Клиент ai-chat vitest 399 зелёный, tsc 0 ai-chat-ошибок; сервер delta(6, реально исполняется)/registry/step-marker/attach + integration attach — зелёное. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ai-chat/components/chat-thread.test.tsx | 37 ++++++++++++++++ .../ai-chat/components/chat-thread.tsx | 44 +++++++++++++++---- .../features/ai-chat/state/run-fsm.spec.md | 10 +++-- 3 files changed, 78 insertions(+), 13 deletions(-) diff --git a/apps/client/src/features/ai-chat/components/chat-thread.test.tsx b/apps/client/src/features/ai-chat/components/chat-thread.test.tsx index 5fbcebaa..47e047cd 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.test.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.test.tsx @@ -1089,6 +1089,43 @@ describe("ChatThread — live reconnect + stalled", () => { ); }); + it("#491 regression (#137/#161 dup): getRun REJECT on a live disconnect drops the live partial + nulls the anchor", async () => { + // The re-seed source (getRun) FAILS — a flaky-network blip (SSE + getRun both + // fail, network recovers in ~1s). The OLD .catch just re-entered the ladder with + // NO re-seed and NO filter, so the reconnect could tail-apply the registry's + // frames onto the live partial that ALREADY has those steps -> duplicated text. + renderLive(); + h.state.getRun.mockReset(); + h.state.getRun.mockRejectedValue(new Error("network")); + await disconnect(); // live partial = liveMsg (id "a2") + expect(h.state.getRun).toHaveBeenCalledWith("c1"); + // THE GUARANTEE: on the getRun failure the live partial (a2) is FILTERED from the + // store, so the reconnect can never tail-apply already-present steps onto it. + // MUTATION-VERIFY: revert the .catch fix (enterReconnect only, no filter) and no + // setMessages call removes a2 -> this reddens. + const removedLivePartial = ( + h.state.setMessages as unknown as { + mock: { calls: [unknown][] }; + } + ).mock.calls.some(([updater]) => { + if (typeof updater !== "function") return false; + const out = (updater as (p: { id: string }[]) => { id: string }[])([ + { id: "a2" }, + { id: "u1" }, + ]); + return !out.some((m) => m.id === "a2"); + }); + expect(removedLivePartial).toBe(true); + expect(screen.getByText(/reconnecting/i)).toBeTruthy(); + advanceToAttempt(1); + expect(h.state.resumeStream).toHaveBeenCalledTimes(1); + // Anchor was nulled -> replay-from-start (no params) / 204 -> poll; never a stale + // ?anchor=&n= over the live partial. + expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe( + "/api/ai-chat/runs/c1/stream", + ); + }); + it("#488 (browser QA): the reconnect banner is SHOWN, not masked by the residual useChat error", async () => { // The drop sets useChat `error` (real SDK), and the terminal errorView describes // it ("Lost connection to the server"). The FSM phase-gate must let the diff --git a/apps/client/src/features/ai-chat/components/chat-thread.tsx b/apps/client/src/features/ai-chat/components/chat-thread.tsx index 01e703ce..6eb7fdaa 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.tsx @@ -720,13 +720,33 @@ export default function ChatThread({ const runId = extractRunId(message ?? undefined) ?? "pending"; const enterReconnect = (fact: string): void => { if (!mountedRef.current) return; - dispatch({ type: "RUN_FACT", runFact: { runId: fact } }); + // Epoch-stamp the run-fact too (I1): the getRun rtt widens the + // onFinish->dispatch window, so a concurrent SEND_LOCAL during it must be + // able to drop this stale RUN_FACT (else it clobbers the new turn's + // runFact.runId). Consistent with the postRun RUN_FACT stamp. + dispatch({ type: "RUN_FACT", runFact: { runId: fact }, epoch: stampEpoch }); dispatch({ type: "FINISH_DISCONNECT", hasVisibleContent: msgHasVisible, epoch: stampEpoch, }); }; + // Restore the STRUCTURAL guarantee that the live partial is never the + // tail-apply base: drop the live partial from the store by id and null the + // anchor, so the reconnect replays from step 0 into a CLEAN store (a full + // rebuild) or, past any rotation, 204s -> degraded poll. Used on BOTH the + // no-persisted-row and getRun-FAILURE paths — after this there is no path + // where the attach tail-applies frames onto a row that already has them + // (the #137/#161 duplication class). + const dropLivePartialAndReplayFromStart = (): void => { + if (message?.role === "assistant" && typeof message.id === "string") { + const liveId = message.id; + setMessagesRef.current?.((prev) => + prev.filter((m) => m.id !== liveId), + ); + } + anchorRef.current = null; + }; if (cid) { void getRun(cid) .then((res) => { @@ -738,23 +758,29 @@ export default function ChatThread({ stepsPersisted: stepsPersistedOf(persisted), }; // Replace the live partial with the persisted row IN PLACE by id — - // the re-seed from persist. + // the re-seed from persist. The attach's tail (steps >= N) then + // appends to a store holding EXACTLY steps 0..N-1: no duplication. setMessages((prev) => mergeById(prev, rowToUiMessage(persisted))); } else { - // No persisted assistant row (pre-first-frame break): no anchor; the - // reconnect attaches with no params and replays from start. - anchorRef.current = null; + // No persisted assistant row (pre-first-frame break): drop the live + // partial + replay from start (no anchor/n) so nothing is duplicated. + dropLivePartialAndReplayFromStart(); } enterReconnect(res.run?.id ?? runId); }) .catch(() => { - // Persist read failed: keep whatever anchor we had and still enter the - // ladder so a live detached run is not stranded (the attach 204s -> - // degraded poll if it cannot cover). + if (!mountedRef.current) return; + // Persist read FAILED: we cannot re-seed from fresh persist, and a + // stale mount-time anchor over the live partial would tail-apply + // already-present steps -> duplication (a flaky-network blip: + // SSE + getRun both fail, network recovers in ~1s, the registry still + // covers from the mount frontier). Restore the removed-filter guarantee + // instead: drop the live partial + replay from start / 204 -> poll. + dropLivePartialAndReplayFromStart(); enterReconnect(runId); }); } else { - anchorRef.current = null; + dropLivePartialAndReplayFromStart(); enterReconnect(runId); } setStopNotice(null); diff --git a/apps/client/src/features/ai-chat/state/run-fsm.spec.md b/apps/client/src/features/ai-chat/state/run-fsm.spec.md index 4f70f744..c40f3fbf 100644 --- a/apps/client/src/features/ai-chat/state/run-fsm.spec.md +++ b/apps/client/src/features/ai-chat/state/run-fsm.spec.md @@ -121,8 +121,7 @@ holds. **Pending column: empty.** | 11 | `stopPendingRef` | **FSM phase `stopping`** | the deferred stop fires from the chat-id adoption effect while `stopping` | | 12 | `mountedRef` | **retained (React liveness)** | orthogonal to run-lifecycle; gates imperative onFinish side-effects post-unmount. Epoch (I1) handles stale COMMAND-outcomes; DISPOSE bumps it | | 13 | `attemptResumeRef` | **FSM `ATTACH_START` + run-fact** | mount arms attach ONLY on a confirmed active run (commit 4b: streaming-tail status, or POST /run for a user tail) | -| 14 | `stripRef` | **data** (attachStrategy) | strip+replay detail; the `resumeStream` effect reads it | -| 15 | `strippedRowRef` | **data** (attachStrategy) | the anchor row | +| 14–15 | `anchorRef {id, stepsPersisted}` | **data** (attachStrategy) | #491 tail-only: replaced `stripRef`/`strippedRowRef`. The PERSISTED assistant row that pins the run (server invariant 6) + its step frontier N; feeds `?anchor=&n=`. No strip — the seed keeps every row; entering reconnecting re-seeds from persist | | 16 | `attachAbortRef` | **effect-owned controller** | aborted by the `abortAttach` effect in cleanup (I5) | | 17–25 | `chatIdRef`, `openPageRef`, `getEditorSelectionRef`, `roleIdRef`, `stableIdRef`, `queuedRef`, `sendMessageRef`, `statusRef`, `lastForwardedChatIdRef` | **data** (identity/send mirrors) | unchanged — not lifecycle flags | | NEW | `pendingSupersedeRef` | **data** (send-plumbing) | the runId injected into the next `POST /stream {supersede}`; the single replacement for the 3 DELETED one-shots (#8/#9/#10) — net −2 refs | @@ -178,6 +177,9 @@ Pessimism rule: a stale-but-positive fact PERMITS entering recovery (attach); th /run) are effect-owned and aborted in cleanup (`abortAttach` on `DISPOSE`), not render-phase refs. A client abort of an already-sent POST does not cancel the server action, so disarming on unmount is safe. -- **attachStrategy** (strip+replay today) is behind the `resumeStream` effect; the - resume-stack iteration (#491) swaps it to tail-only WITHOUT touching the FSM. +- **attachStrategy** is behind the `resumeStream` effect; #491 swapped it to + tail-only (`?anchor=&n=`, `anchorRef` data) WITHOUT touching the FSM. Entering + reconnecting always re-seeds from persist; on a getRun failure the live partial + is dropped + replay-from-start so it is never the tail-apply base (no #137/#161 + duplication). - **Queue** stays a data structure; flush/interrupt decisions are transitions. -- 2.52.0 From 3550bfa41107d7d89cf46dd0d5494a4e103366fa Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sun, 12 Jul 2026 00:48:26 +0300 Subject: [PATCH 7/7] =?UTF-8?q?fix(ai-chat):=20=D1=80=D0=B5-=D1=80=D0=B5?= =?UTF-8?q?=D0=B2=D1=8C=D1=8E=20#518=20=E2=80=94=20CI-=D1=84=D0=B8=D0=B4?= =?UTF-8?q?=D0=B5=D0=BB=D0=B8=D1=82=D0=B8=20delta,=20attach-=D0=B4=D1=83?= =?UTF-8?q?=D0=B1=D0=BB=D1=8C,=20=D1=87=D0=B8=D1=81=D1=82=D0=BA=D0=B8=20(#?= =?UTF-8?q?491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ребейз на обновлённый feat/490 (PR #510, Option-A): стек #491 перенесён с 6e872f2d на 6e42752a; run-fsm.ts/chat-thread.tsx приехали из develop-версии (W1/S4: RUN_ALREADY_ACTIVE несёт activeRunId + supersede-CAS адопция) — мои #491-правки легли сверху без конфликтов, W1 НЕ откачен (run-fsm.ts == develop байт-в-байт; RUN_ALREADY_ACTIVE диспатчит activeRunId). 1. [CI-fidelity] Дельта-курсор спек падал в CI unit-лейне: `.spec.ts` дефолтил на НЕмигрированный docmost → 5/6 ERROR relation does not exist (skip-гард ловит только connection-fail). Переименован в `*.int-spec.ts` (исключается из unit- regex `.spec.ts$`, гоняется в test:int, чей global-setup мигрирует docmost_test) + DSN по умолчанию → docmost_test. Теперь 6/6 реально исполняются в CI-верном окружении; overlap-мутация роняет RACE-тесты (не вакуумен). 2. [regression #137/#161] attach: отсутствие `n` схлопывалось в frontier 0 → finished-неротированный ран (coverageFloor 0) отдавал ВЕСЬ tail вместо 204; парамслесс/легаси-вкладка допишет полный replay → дубль. Различаем ОТСУТСТВИЕ `n` (null — не tail-aware) от `n=0` (tail-aware): контроллер шлёт null при missing/invalid; registry.attach(n: number|null) 204-ит finished-ран при n===null (старый `finished && !expectLive` гейт), n=0 по-прежнему отдаёт хвост. Тесты (registry + controller) + mutation-verify: нейтрализация гейта роняет их. 4. [conventions] Ring-кап env-var переименован RUN_STREAM_MAX_BUFFER_BYTES → AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES (префикс как у сиблингов) + запись в .env.example (дефолт 4MB, 0/invalid→дефолт, subscriber-cap=2×). 5. [docs] run-fsm.spec.md: item4 переписан («реализовано в #491 — дельта несёт run:{id,status}|null; клиент run-поле ещё не потребляет») + добавлена строка перехода POLL_IDLE_CAP stopping→idle (Review #4, редьюсер это делает). 6. [simplification] Удалена мёртвая цепочка reconstructRunParts / reconstructPartsFromRow (ноль прод-вызовов) + опц. messageRepo-инъекция в AiChatRunService + спек-блоки; вернётся с первым реальным вызывателем. Маркер metadata.stepsPersisted (реально используемый) сохранён. DROP-пункты ревьюера (осиротевший import, DELTA_POLL_MAX_ROWS) не трогаю. Прогон: server tsc 0, ai-chat unit 202, delta-int 6/6 (int-lane), int attach 6/6, client vitest 403, tsc client 0 ai-chat, mcp 834/0. FSM-инварианты #488 сохранены. Co-Authored-By: Claude Opus 4.8 (1M context) --- .env.example | 10 +++ .../features/ai-chat/state/run-fsm.spec.md | 9 ++- .../core/ai-chat/ai-chat-run.service.spec.ts | 73 ------------------- .../src/core/ai-chat/ai-chat-run.service.ts | 34 --------- .../ai-chat-stream-registry.service.ts | 37 +++++++--- .../ai-chat/ai-chat-stream-registry.spec.ts | 29 +++++++- .../ai-chat/ai-chat.controller.attach.spec.ts | 35 ++++++++- .../src/core/ai-chat/ai-chat.controller.ts | 14 +++- .../src/core/ai-chat/ai-chat.service.ts | 28 ------- .../core/ai-chat/ai-chat.step-marker.spec.ts | 43 +---------- ...=> ai-chat-message.repo.delta.int-spec.ts} | 12 ++- 11 files changed, 126 insertions(+), 198 deletions(-) rename apps/server/src/database/repos/ai-chat/{ai-chat-message.repo.delta.spec.ts => ai-chat-message.repo.delta.int-spec.ts} (92%) diff --git a/.env.example b/.env.example index 6d7dec43..51fc4d4d 100644 --- a/.env.example +++ b/.env.example @@ -302,6 +302,16 @@ MCP_DOCMOST_PASSWORD= # enabled for a workspace, and the same single-instance constraint applies (the # registry is process-local). # AI_CHAT_RESUMABLE_STREAM=false +# +# Per-run replay ring cap (#491), in BYTES, for the resumable-stream registry +# above. The registry buffers the run's recent SSE tail so a reopened tab can +# attach and continue from the step it already persisted; the ring is bounded and +# rotates on every confirmed step-persist. This caps the un-persisted tail between +# rotations — an overflow evicts the oldest frames and a late attach falls back to +# 204 -> degraded poll, so correctness never depends on the size. Default 4194304 +# (4MB); a 0/invalid value falls back to the default. The per-subscriber backpressure +# cap is derived as 2x this value. Only meaningful with AI_CHAT_RESUMABLE_STREAM on. +# AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES=4194304 # --- Run lifecycle tunables (#487) --- # These govern the universal run machinery (every turn is now a first-class run, diff --git a/apps/client/src/features/ai-chat/state/run-fsm.spec.md b/apps/client/src/features/ai-chat/state/run-fsm.spec.md index c40f3fbf..86ae3d73 100644 --- a/apps/client/src/features/ai-chat/state/run-fsm.spec.md +++ b/apps/client/src/features/ai-chat/state/run-fsm.spec.md @@ -48,6 +48,7 @@ Legend: **†** = command-transition (bumps `epoch`, I1). Effects in `[…]`. | `RETRY` (manual, stalled banner) | stalled | polling(attach-none) **†** | `[armPoll]` | | `POLL_TERMINAL` (settled tail merged) | polling, reconnecting, stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (I4) | | `POLL_IDLE_CAP` (inactivity cap) | polling, reconnecting | stalled | `[disarmPoll, cancelReconnect]` (commit 4a — no more silent) | +| `POLL_IDLE_CAP` (inactivity cap) | stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (Review #4: a Stop-armed poll with no SDK/terminal backstop gets a bounded exit — NOT `stalled`, Stop was already pressed so nothing to retry) | | `RUN_FACT{null}` (POST /run → null/terminal, 204) | reconnecting/attaching/polling/stopping | idle | `[cancelReconnect, disarmPoll]`, runFact←null (I3 fresh-negative gate) | | `RUN_FACT{runId}` | any | (same) | runFact←runId (pessimism toward an attempt) | | `STOP_REQUESTED` (user Stop) | streaming, reconnecting, polling | stopping **†** | `[stopRun, abortAttach, cancelReconnect, armPoll]` (poll drives the terminal — I4 exit by data) | @@ -150,8 +151,12 @@ message. Sources, in the order they update `ctx.runFact`: 3. **Attach outcomes:** `ATTACH_LIVE` (2xx) confirms active; a 204 on a non-stripped path is an authoritative NEGATIVE fact → the runtime dispatches `RUN_FACT{null}`, which cancels recovery (I3 fresh-negative gate). -4. **Poll (future resume-stack iteration #491):** the delta will carry the run field; - until then the poll drives to a terminal ROW, dispatched as `POLL_TERMINAL`. +4. **Poll (#491, implemented):** the degraded poll now hits the delta endpoint + (`POST /ai-chat/messages/delta`), which ALREADY carries the run fact + (`run: {id, status} | null`) alongside the changed rows. The client does NOT yet + consume that run field — it still drives to a terminal ROW (merged by id), + dispatched as `POLL_TERMINAL` — so the run field rides the wire for a future + client that settles straight off it. Pessimism rule: a stale-but-positive fact PERMITS entering recovery (attach); the 204 then cuts it. A fresh negative fact gates recovery OUT immediately. 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 41640cb0..ade4563c 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,76 +771,3 @@ 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 37ab3827..1530c3b5 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,9 +1,6 @@ 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'; @@ -215,10 +212,6 @@ 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, ) {} /** @@ -791,33 +784,6 @@ 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-stream-registry.service.ts b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts index 743bad30..e338d01a 100644 --- a/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts @@ -88,14 +88,14 @@ export const RUN_STREAM_RETAIN_FINISHED_MS = 30_000; * its persisted frontier come from the seed, not the ring). The ring stays bounded * because it rotates on every confirmed persist; this cap is only the ceiling for * the un-persisted tail between rotations. Env-tunable via - * RUN_STREAM_MAX_BUFFER_BYTES (bytes); a 0/invalid value falls back to this. + * AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES (bytes); a 0/invalid value falls back to this. */ -export const RUN_STREAM_MAX_BUFFER_BYTES = 4 * 1024 * 1024; +export const AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES = 4 * 1024 * 1024; // 2× the ring cap: a just-written full-tail burst alone can never trip the // per-subscriber cap (see controller); only a genuinely stalled socket can. This // derivative relationship is preserved even when the ring cap is env-overridden. -export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * RUN_STREAM_MAX_BUFFER_BYTES; +export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES; /** * A finish-step boundary frame is exactly `data: {"type":"finish-step"...}\n\n` @@ -110,12 +110,12 @@ const FINISH_STEP_FRAME_PREFIX = 'data: {"type":"finish-step"'; /** Resolve the ring cap from the environment, falling back to the default. */ function resolveMaxBufferBytes(): number { - const raw = process.env.RUN_STREAM_MAX_BUFFER_BYTES; - if (!raw) return RUN_STREAM_MAX_BUFFER_BYTES; + const raw = process.env.AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES; + if (!raw) return AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES; const parsed = Number(raw); return Number.isFinite(parsed) && parsed > 0 ? Math.floor(parsed) - : RUN_STREAM_MAX_BUFFER_BYTES; + : AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES; } export interface RunStreamCallbacks { @@ -321,7 +321,11 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { async attach( chatId: string, anchor: string | undefined, - n: number, + // The client's persisted step frontier. `null` = a NOT-tail-aware client (no + // `n` query param) — a legacy/parameterless tab that expects the old + // "finished -> 204 -> poll" contract; distinct from `0` (a tail-aware client + // with nothing persisted yet). + n: number | null, cb: RunStreamCallbacks, ): Promise { const entry = this.entries.get(chatId); @@ -329,14 +333,25 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { // Invariant 6: cross-run replay is forbidden. Before bind, assistantMessageId // is undefined and mismatches any anchor -> 204 -> client restore+poll path. if (anchor && entry.assistantMessageId !== anchor) return null; + // #491 regression guard (#137/#161 dup): a NOT-tail-aware client (no `n`) + // resuming a FINISHED run must 204 and poll — the old `finished && !expectLive` + // gate. Without this, a missing `n` collapsing to frontier 0 would serve the + // WHOLE tail of a finished, NON-rotated run (coverageFloor 0), and a + // parameterless client that never stripped its transcript would APPEND that + // full replay onto the steps it already shows -> duplicated text. A tail-aware + // client (n present, incl. n=0) still gets the tail past its frontier. + if (entry.finished && n === null) return null; // A finished entry with NOTHING in the ring (aborted before the first frame, // or fully overflowed) has no tail to deliver -> 204 -> the client polls. if (entry.finished && entry.frames.length === 0) return null; + // A LIVE run with no `n` (legacy parameterless) replays from step 0 (the old + // behavior); a tail-aware client resumes from its frontier. + const frontier = n ?? 0; const floor = this.coverageFloor(entry); - if (floor > n) { + if (floor > frontier) { this.logger.warn( `run-stream attach gap for run=${entry.runId}: coverageFloor=${floor} ` + - `> client n=${n} -> 204 (client refetches + re-attaches)`, + `> client frontier=${frontier} -> 204 (client refetches + re-attaches)`, ); return null; } @@ -345,7 +360,7 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { const sliceTail = (): string[] => { const out: string[] = [startFrame]; for (let i = 0; i < entry.frames.length; i++) { - if (entry.stamps[i] >= n) out.push(entry.frames[i]); + if (entry.stamps[i] >= frontier) out.push(entry.frames[i]); } return out; }; @@ -368,7 +383,7 @@ export class AiChatStreamRegistryService implements OnModuleDestroy { pendingBytes: 0, overflowed: false, pendingEnd: false, - minStamp: n, + minStamp: frontier, }; // Register + snapshot in the SAME synchronous block (invariant 4). No await // separates them, so a concurrently ingested frame cannot be lost/duplicated. diff --git a/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts b/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts index 06df6d75..9388b2b6 100644 --- a/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts @@ -1,6 +1,6 @@ import { AiChatStreamRegistryService, - RUN_STREAM_MAX_BUFFER_BYTES, + AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES, RUN_STREAM_RETAIN_FINISHED_MS, RunStreamCallbacks, } from './ai-chat-stream-registry.service'; @@ -377,7 +377,7 @@ describe('AiChatStreamRegistryService step-aligned retention (#491)', () => { const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); // Step 0: a fat step that blows past the cap with NO persist confirmation. - const big = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 2)); + const big = 'x'.repeat(Math.floor(AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES / 2)); src.push(textDelta('t0', big)); // 0 src.push(textDelta('t0', big)); // 0 src.push(textDelta('t0', big)); // 0 -> overflow evicts stamp-0 frames @@ -420,7 +420,7 @@ describe('AiChatStreamRegistryService step-aligned retention (#491)', () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); registry.bind(CHAT, 'run-1', 'assist-1', src.stream); - const big = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 2)); + const big = 'x'.repeat(Math.floor(AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES / 2)); src.push(textDelta('t0', big)); // 0 src.push(textDelta('t0', big)); // 0 src.push(finishStep()); // 0 (still stamp 0) @@ -460,6 +460,29 @@ describe('AiChatStreamRegistryService step-aligned retention (#491)', () => { expect(entryOf().subscribers.size).toBe(0); }); + it('#491 regression (#137/#161 dup): a PARAMETERLESS attach (n=null) to a finished NON-rotated run -> 204, but n=0 still gets the tail', async () => { + // A finished, non-rotated run: frames present, coverageFloor 0. A missing `n` + // (null — a legacy/parameterless tab that never stripped its transcript) must + // 204 -> poll, NOT receive the whole tail it would append (duplicate). A + // tail-aware client (n=0 present) still resumes. + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push(textDelta('t0', 'a')); // 0 + src.push(finishStep()); // 0 + src.push(finish()); // 1 + src.close(); + await flush(); + // NOT rotated (no confirmPersistedStep) -> stamps[0]=0, coverageFloor=0. + // MUTATION-VERIFY: revert the `finished && n === null -> null` gate (default n + // to 0) and the parameterless attach below serves the full tail instead of 204. + expect(await registry.attach(CHAT, 'assist-1', null, collector().cb)).toBeNull(); + // A tail-aware client at frontier 0 IS served (the distinction: null != 0). + const tailAware = await registry.attach(CHAT, 'assist-1', 0, collector().cb); + expect(tailAware).not.toBeNull(); + expect(tailAware!.finished).toBe(true); + }); + it('confirmPersistedStep is monotonic and identity-checked', async () => { registry.open(CHAT, 'run-1'); const src = makePushStream(); diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts b/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts index 05bc19b8..50666847 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts @@ -182,7 +182,38 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { ); }); - it('floors n to 0 when the query is absent/invalid', async () => { + it('#491: an ABSENT/invalid n passes null (not 0) so a finished run 204s (not-tail-aware)', async () => { + // Distinguishing a MISSING `n` from `n=0` is the #137/#161 dup guard: a + // parameterless/legacy tab must be handed null (-> the registry 204s a finished + // run) rather than frontier 0 (which would serve a finished non-rotated run's + // whole tail). MUTATION-VERIFY: revert to `Number(n) || 0` and this asserts 0. + const { controller, streamRegistry } = makeController({ + chat: owned, + attachment: null, + }); + for (const bad of [undefined, '', 'abc']) { + streamRegistry.attach.mockClear(); + const { res } = makeRawRes(); + const { req } = makeReq(); + await controller.attachRunStream( + 'c1', + undefined, + bad, + req, + res, + user, + workspace, + ); + expect(streamRegistry.attach).toHaveBeenCalledWith( + 'c1', + undefined, + null, + expect.anything(), + ); + } + }); + + it('#491: a PRESENT n=0 passes 0 (tail-aware, distinct from absent)', async () => { const { controller, streamRegistry } = makeController({ chat: owned, attachment: null, @@ -192,7 +223,7 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => { await controller.attachRunStream( 'c1', undefined, - undefined, + '0', req, res, user, diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.ts b/apps/server/src/core/ai-chat/ai-chat.controller.ts index b97ad6ff..2b8aa6ab 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.ts @@ -365,9 +365,17 @@ export class AiChatController { @AuthWorkspace() workspace: Workspace, ): Promise { await this.assertOwnedChat(chatId, user, workspace); // same gate as getRun - // The client's persisted step frontier. A missing/invalid value floors to 0 - // ("give me everything") which, past any rotation, safely 204s. - const frontier = Number.isFinite(Number(n)) ? Math.max(0, Number(n)) : 0; + // The client's persisted step frontier. #491: distinguish a MISSING/invalid `n` + // (null — a NOT-tail-aware, legacy/parameterless tab expecting the old + // "finished -> 204 -> poll" contract) from `n=0` (a tail-aware client with + // nothing persisted yet). Passing 0 for a missing `n` would serve a finished, + // non-rotated run's WHOLE tail and a parameterless client would append it onto + // the steps it already shows -> #137/#161 duplicate. null makes the registry + // 204 such a finished run (see attach); a tail-aware n=0 still resumes. + const frontier: number | null = + n === undefined || n === '' || !Number.isFinite(Number(n)) + ? null + : Math.max(0, Number(n)); // The per-subscriber backpressure cap tracks the (env-tunable) ring cap. const subscriberCap = this.streamRegistry?.subscriberMaxBufferedBytes ?? 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 e58788b0..9922ba5a 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -2747,34 +2747,6 @@ 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 + 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 index 98c53e4b..a4398a3d 100644 --- 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 @@ -1,15 +1,11 @@ -import { - flushAssistant, - reconstructPartsFromRow, -} from './ai-chat.service'; -import type { AiChatMessage } from '@docmost/db/types/entity.types'; +import { flushAssistant } from './ai-chat.service'; /** * #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. + * every flush shape. */ // A finished step carrying one line of text and one tool call/result. @@ -67,38 +63,3 @@ describe('flushAssistant step marker (#491)', () => { 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); - }); -}); diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.int-spec.ts similarity index 92% rename from apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts rename to apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.int-spec.ts index 94ccddeb..063b2d15 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.spec.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.delta.int-spec.ts @@ -26,6 +26,15 @@ import { AiChatRunRepo } from './ai-chat-run.repo'; * 4. the overlap GUARANTEES repeats across close polls — the contract behind the * client's idempotent merge (mergeById). * + * INTEGRATION lane (`*.int-spec.ts`): runs under `test:int`, whose global-setup + * DROPS + RE-CREATES + MIGRATES `docmost_test`, so the real `ai_chat_messages` / + * `ai_chat_runs` tables EXIST here. (It was previously a `.spec.ts` defaulting to + * the UNmigrated dev `docmost`; in the CI unit lane — where `WAL_TEST_DATABASE_URL` + * is unset and only `test:int` migrates — that meant 5/6 ERROR + * `relation "ai_chat_messages" does not exist`, silently voiding coverage of the + * risky cursor/overlap logic. Renaming to `.int-spec.ts` + defaulting the DSN to + * `docmost_test` fixes the CI fidelity.) + * * FK triggers are bypassed (`session_replication_role = replica`) so synthetic * chat/workspace ids need no parent fixtures; a single pooled connection (max 1) * keeps that session setting for every query. SKIPS cleanly when the DB is @@ -33,7 +42,8 @@ import { AiChatRunRepo } from './ai-chat-run.repo'; */ const CONN = process.env.WAL_TEST_DATABASE_URL ?? - 'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost'; + process.env.TEST_DATABASE_URL ?? + 'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost_test'; let db: Kysely; let sqlClient: ReturnType; -- 2.52.0