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 new file mode 100644 index 00000000..b6a5fc47 --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.service.ts @@ -0,0 +1,294 @@ +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. + * + * 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. + */ + +/** 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). */ +export const RUN_STREAM_MAX_BUFFER_BYTES = 4 * 1024 * 1024; + +// 2x the replay cap: a just-written 4MB replay burst alone can never trip the +// per-subscriber cap (see controller); only a genuinely stalled socket can. +export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * RUN_STREAM_MAX_BUFFER_BYTES; + +export interface RunStreamCallbacks { + onFrame: (frame: string) => void; + onEnd: () => void; +} + +export interface RunStreamAttachment { + replay: string[]; + finished: boolean; + start(): void; // drain pending frames (order preserved) and go live + unsubscribe(): void; // safe to call at any point, idempotent +} + +interface Subscriber extends RunStreamCallbacks { + started: boolean; + pending: string[]; + pendingEnd: boolean; +} + +interface Entry { + runId: string; + // 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; + frames: string[]; + bytes: number; + overflowed: boolean; + finished: boolean; + subscribers: Set; + retainTimer?: NodeJS.Timeout; +} + +@Injectable() +export class AiChatStreamRegistryService implements OnModuleDestroy { + private readonly logger = new Logger(AiChatStreamRegistryService.name); + private readonly entries = new Map(); // key: chatId + + /** + * Register a fresh entry at the START of a run (before any frame), so a tab + * that attaches in the begin->seed window finds an entry to wait on. If an + * entry already exists for this chat (a previous, possibly still-live run whose + * tee loop is draining), it is terminated MIRRORING the done-path (invariant 3) + * so its subscribers are released and its retention timer is cleared; a late + * `done` from that old tee then fires against the closed-over old reference and, + * thanks to identity checks, never touches this new entry. + */ + open(chatId: string, runId: string): void { + const existing = this.entries.get(chatId); + if (existing) { + if (existing.retainTimer) { + clearTimeout(existing.retainTimer); + existing.retainTimer = undefined; + } + // Started subscribers get exactly one onEnd() and are removed; paused ones + // are marked pendingEnd (their start() will end them). finished=true guards + // any later done from the old tee loop from double-notifying. + this.terminateSubscribers(existing); + } + this.entries.set(chatId, { + runId, + frames: [], + bytes: 0, + overflowed: false, + finished: false, + subscribers: new Set(), + }); + } + + /** + * Tee a run's SSE frame stream into its entry (called from consumeSseStream). + * No-op with a warning when there is no entry or the entry belongs to a + * different run (invariant 1). The reader loop is fire-and-forget: the tee + * branch outlives the client socket by design. + */ + bind( + chatId: string, + runId: string, + assistantMessageId: string | undefined, + stream: ReadableStream, + ): void { + const entry = this.entries.get(chatId); + if (!entry || entry.runId !== runId) { + // Invariant 1: only the matching run may mutate the entry. + this.logger.warn( + `bind: no matching run-stream entry for chat=${chatId} run=${runId}`, + ); + return; + } + entry.assistantMessageId = assistantMessageId; + const reader = stream.getReader(); + const pump = async (): Promise => { + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + this.ingestFrame(entry, value); + } + this.finalizeEntry(chatId, entry); + } catch { + // A read error is a terminal event too — release subscribers. + this.finalizeEntry(chatId, entry); + } + }; + void pump(); + } + + /** + * 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 + * on runId (invariant 1); the shared terminal path is idempotent. + */ + abortEntry(chatId: string, runId: string): void { + const entry = this.entries.get(chatId); + if (!entry || entry.runId !== runId) return; + this.finalizeEntry(chatId, entry); + } + + /** + * 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. + * + * 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. + */ + async attach( + chatId: string, + expectLive: boolean, + anchor: string | undefined, + cb: RunStreamCallbacks, + ): Promise { + const entry = this.entries.get(chatId); + if (!entry || entry.overflowed) 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) { + // Replay-only: the run is done, no subscriber is registered. + return { + replay: entry.frames.slice(), + finished: true, + start: () => undefined, + unsubscribe: () => undefined, + }; + } + + const sub: Subscriber = { + onFrame: cb.onFrame, + onEnd: cb.onEnd, + started: false, + pending: [], + pendingEnd: false, + }; + entry.subscribers.add(sub); + // Snapshot in the SAME synchronous block as the registration (invariant 4). + const replay = entry.frames.slice(); + return { + replay, + finished: false, + start: () => { + // Deliver frames buffered while paused, in order, then go live. + for (const frame of sub.pending) { + try { + sub.onFrame(frame); + } catch { + entry.subscribers.delete(sub); + return; + } + } + sub.pending = []; + sub.started = true; + if (sub.pendingEnd) { + try { + sub.onEnd(); + } catch { + // The socket is gone; nothing to end. + } + entry.subscribers.delete(sub); + } + }, + unsubscribe: () => { + entry.subscribers.delete(sub); + }, + }; + } + + onModuleDestroy(): void { + for (const entry of this.entries.values()) { + if (entry.retainTimer) clearTimeout(entry.retainTimer); + } + this.entries.clear(); + } + + /** Buffer + fan-out a single frame. See invariant/overflow semantics inline. */ + private ingestFrame(entry: Entry, frame: string): void { + entry.bytes += Buffer.byteLength(frame); + if (!entry.overflowed) { + 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`, + ); + } + } + for (const sub of entry.subscribers) { + if (sub.started) { + try { + sub.onFrame(frame); + } catch { + entry.subscribers.delete(sub); + } + } else { + sub.pending.push(frame); + } + } + } + + /** + * Shared terminal path for done / read-error / external-abort. Idempotent: a + * second call (already finished) is a no-op, so an open()-replaced or + * abort-then-done entry is never double-armed or double-ended. + */ + private finalizeEntry(chatId: string, entry: Entry): void { + if (entry.finished) return; + this.terminateSubscribers(entry); + const timer = setTimeout(() => { + // Invariant 2: only delete OUR entry (a replacement may already own the key). + if (this.entries.get(chatId) === entry) this.entries.delete(chatId); + }, RUN_STREAM_RETAIN_FINISHED_MS); + timer.unref?.(); + entry.retainTimer = timer; + } + + /** + * Mark the entry finished and release its subscribers, mirroring the done-path: + * started subscribers get exactly one onEnd() and are removed; paused ones are + * flagged pendingEnd so their start() ends them. Deleting the current element + * during Set iteration is safe. + */ + private terminateSubscribers(entry: Entry): void { + entry.finished = true; + for (const sub of entry.subscribers) { + if (sub.started) { + try { + sub.onEnd(); + } catch { + // The socket is gone; nothing to end. + } + entry.subscribers.delete(sub); + } else { + sub.pendingEnd = true; + } + } + } +} 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 new file mode 100644 index 00000000..e2171986 --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat-stream-registry.spec.ts @@ -0,0 +1,357 @@ +import { + AiChatStreamRegistryService, + RUN_STREAM_MAX_BUFFER_BYTES, + RUN_STREAM_RETAIN_FINISHED_MS, + 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. + */ + +// A ReadableStream whose frames the test pushes explicitly, plus close/error. +function makePushStream(): { + stream: ReadableStream; + push: (f: string) => void; + close: () => void; + error: (e?: unknown) => void; +} { + let controller!: ReadableStreamDefaultController; + const stream = new ReadableStream({ + start(c) { + controller = c; + }, + }); + return { + stream, + push: (f) => controller.enqueue(f), + close: () => controller.close(), + error: (e) => controller.error(e ?? new Error('read error')), + }; +} + +// Let the fire-and-forget pump drain queued frames (reader.read() resolves on a +// macrotask boundary for an already-enqueued value). +const flush = (): Promise => new Promise((r) => setTimeout(r, 0)); + +function collector(): { + cb: RunStreamCallbacks; + frames: string[]; + ended: () => number; +} { + const frames: string[] = []; + let ends = 0; + return { + frames, + ended: () => ends, + cb: { + onFrame: (f) => frames.push(f), + onEnd: () => { + ends += 1; + }, + }, + }; +} + +describe('AiChatStreamRegistryService', () => { + const CHAT = 'chat-1'; + let registry: AiChatStreamRegistryService; + + beforeEach(() => { + registry = new AiChatStreamRegistryService(); + jest.spyOn((registry as any).logger, 'warn').mockImplementation(() => {}); + }); + + afterEach(() => { + registry.onModuleDestroy(); + }); + + it('replays frames in arrival order (live attach)', 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.push('c'); + await flush(); + + const c = collector(); + const att = await registry.attach(CHAT, false, undefined, c.cb); + expect(att).not.toBeNull(); + expect(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 () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + src.push('a'); + src.push('b'); + await flush(); + + const c = collector(); + const att = (await registry.attach(CHAT, false, undefined, c.cb))!; + expect(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 () => { + 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(); + // 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']); + 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 + expect(c.frames).toEqual(['b', 'c']); + src.push('d'); // now live + 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'); + const c = collector(); + const att = (await registry.attach(CHAT, false, undefined, c.cb))!; + // Terminate the run while the subscriber is still paused. + 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 () => { + 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(); + + 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(); + }); + + it('matching anchor with expect=live attaches', 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, true, 'assist-1', c.cb); + expect(att).not.toBeNull(); + expect(att!.replay).toEqual(['a']); + }); + + it('overflow: attach returns null, but the LIVE subscriber keeps receiving (incl. the crossing frame)', 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 c = collector(); + const att = (await registry.attach(CHAT, false, undefined, c.cb))!; + att.start(); + + const oneMb = 'x'.repeat(1024 * 1024); + // 5 x 1MB = 5MB > 4MB cap; the 5th frame is the one that crosses. + for (let i = 0; i < 5; i++) src.push(oneMb + 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(oneMb + 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('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 + + 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 + const still = (registry as any).entries.get(CHAT); + expect(still).toBe(newEntry); + expect(still.runId).toBe('run-2'); + }); + + it('bind with a foreign runId is a no-op (invariant 1)', async () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'WRONG-run', 'assist-x', src.stream); + 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(); + }); + + it('abortEntry with a foreign runId is a no-op (invariant 1)', async () => { + registry.open(CHAT, 'run-1'); + registry.abortEntry(CHAT, 'WRONG-run'); + 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 () => { + registry.open(CHAT, 'run-1'); + const src = makePushStream(); + registry.bind(CHAT, 'run-1', 'assist-1', src.stream); + + const bad = collector(); + const badAtt = (await registry.attach(CHAT, false, undefined, { + onFrame: () => { + throw new Error('boom'); + }, + onEnd: bad.cb.onEnd, + }))!; + badAtt.start(); + + 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 + 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']); + }); +}); + +/** + * Retention + replace timer behavior. Fake timers, and entries are finalized via + * the synchronous abortEntry() path so no stream pump / microtask juggling is + * needed. + */ +describe('AiChatStreamRegistryService retention timers', () => { + const CHAT = 'chat-r'; + let registry: AiChatStreamRegistryService; + + beforeEach(() => { + jest.useFakeTimers(); + registry = new AiChatStreamRegistryService(); + jest.spyOn((registry as any).logger, 'warn').mockImplementation(() => {}); + }); + + afterEach(() => { + registry.onModuleDestroy(); + jest.useRealTimers(); + }); + + it('a finished entry is removed after the retention window', () => { + registry.open(CHAT, 'run-1'); + registry.abortEntry(CHAT, 'run-1'); // finalize -> retention armed + expect((registry as any).entries.get(CHAT)).toBeDefined(); + jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1); + expect((registry as any).entries.get(CHAT)).toBeUndefined(); + }); + + 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. + 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 + const clearSpy = jest.spyOn(global, 'clearTimeout'); + registry.open(CHAT, 'run-2'); // must clear run-1's retain timer + expect(clearSpy).toHaveBeenCalled(); + jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1); + const entry = (registry as any).entries.get(CHAT); + expect(entry).toBeDefined(); + expect(entry.runId).toBe('run-2'); + }); +}); 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 new file mode 100644 index 00000000..dda95ff2 --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat.controller.attach.spec.ts @@ -0,0 +1,324 @@ +import { ForbiddenException } from '@nestjs/common'; +import { AiChatController } from './ai-chat.controller'; +import type { + RunStreamAttachment, + RunStreamCallbacks, +} from './ai-chat-stream-registry.service'; +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 + * (`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, + * aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo, + * streamRegistry, environment). + */ +describe('AiChatController attach endpoint (#184 phase 1.5)', () => { + const user = { id: 'u1' } as User; + const workspace = { id: 'ws1' } as Workspace; + + function makeRawRes() { + const raw: any = { + writableEnded: false, + writableLength: 0, + destroyed: false, + written: [] as string[], + head: null as any, + write: jest.fn((f: string) => { + raw.written.push(f); + return true; + }), + writeHead: jest.fn((code: number, headers: any) => { + raw.head = { code, headers }; + return raw; + }), + flushHeaders: jest.fn(), + end: jest.fn(() => { + raw.writableEnded = true; + }), + destroy: jest.fn(() => { + raw.destroyed = true; + }), + on: jest.fn(), + once: jest.fn(), + }; + const res: any = { + raw, + status: jest.fn(() => res), + send: jest.fn(), + hijack: jest.fn(), + }; + return { res, raw }; + } + + function makeReq(destroyed = false) { + const handlers: Record void> = {}; + const raw: any = { + destroyed, + once: jest.fn((ev: string, fn: () => void) => { + handlers[ev] = fn; + }), + }; + return { req: { raw } as any, raw, fireClose: () => handlers['close']?.() }; + } + + function makeAttachment( + over: Partial = {}, + ): RunStreamAttachment { + return { + replay: [], + finished: false, + start: jest.fn(), + unsubscribe: jest.fn(), + ...over, + }; + } + + function makeController(opts: { + chat?: unknown; + attachment?: RunStreamAttachment | null; + }) { + const aiChatRepo = { findById: jest.fn().mockResolvedValue(opts.chat) }; + let capturedCb: RunStreamCallbacks | undefined; + const streamRegistry = { + attach: jest.fn( + ( + _chatId: string, + _live: boolean, + _anchor: string | undefined, + cb: RunStreamCallbacks, + ) => { + capturedCb = cb; + return Promise.resolve( + opts.attachment === undefined ? makeAttachment() : opts.attachment, + ); + }, + ), + }; + const environment = { isAiChatResumableStreamEnabled: () => true }; + const controller = new AiChatController( + {} as never, // aiChatService + {} as never, // aiChatRunService + aiChatRepo as never, + {} as never, // aiChatMessageRepo + {} as never, // aiTranscription + {} as never, // pageRepo + streamRegistry as never, + environment as never, + ); + return { + controller, + aiChatRepo, + streamRegistry, + getCb: () => capturedCb!, + }; + } + + const owned = { id: 'c1', creatorId: 'u1' }; + + it('owner-gates: a foreign chat throws ForbiddenException and never attaches', async () => { + const { controller, streamRegistry } = makeController({ + chat: { id: 'c1', creatorId: 'someone-else' }, + }); + const { res } = makeRawRes(); + const { req } = makeReq(); + await expect( + controller.attachRunStream( + 'c1', + undefined, + undefined, + req, + res, + user, + workspace, + ), + ).rejects.toBeInstanceOf(ForbiddenException); + expect(streamRegistry.attach).not.toHaveBeenCalled(); + }); + + it('answers 204 when the registry has nothing to resume (no entry / finished / anchor-mismatch)', async () => { + const { controller } = makeController({ chat: owned, attachment: null }); + const { res } = makeRawRes(); + const { req } = makeReq(); + await controller.attachRunStream( + 'c1', + undefined, + undefined, + req, + res, + user, + workspace, + ); + expect(res.status).toHaveBeenCalledWith(204); + expect(res.send).toHaveBeenCalled(); + expect(res.hijack).not.toHaveBeenCalled(); + }); + + it('threads expect=live and anchor through to the registry', async () => { + const { controller, streamRegistry } = makeController({ + chat: owned, + attachment: null, + }); + const { res } = makeRawRes(); + const { req } = makeReq(); + await controller.attachRunStream( + 'c1', + 'live', + 'anchor-1', + req, + res, + user, + workspace, + ); + expect(streamRegistry.attach).toHaveBeenCalledWith( + 'c1', + true, + 'anchor-1', + expect.anything(), + ); + }); + + it('passes expect=false when the query is absent', async () => { + const { controller, streamRegistry } = makeController({ + chat: owned, + attachment: null, + }); + const { res } = makeRawRes(); + const { req } = makeReq(); + await controller.attachRunStream( + 'c1', + undefined, + undefined, + req, + res, + user, + workspace, + ); + expect(streamRegistry.attach).toHaveBeenCalledWith( + 'c1', + false, + undefined, + expect.anything(), + ); + }); + + it('hijacks, writes headers + replay, registers a close cleanup, then goes live', async () => { + const start = jest.fn(); + const attachment = makeAttachment({ + replay: ['f1', 'f2'], + finished: false, + start, + }); + const { controller } = makeController({ chat: owned, attachment }); + const { res, raw } = makeRawRes(); + const { req } = makeReq(); + await controller.attachRunStream( + 'c1', + undefined, + undefined, + req, + res, + user, + workspace, + ); + expect(res.hijack).toHaveBeenCalled(); + expect(raw.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ 'content-type': 'text/event-stream' }), + ); + expect(raw.written).toEqual(['f1', 'f2']); // replay + expect(start).toHaveBeenCalled(); // go live after replay + expect(req.raw.once).toHaveBeenCalledWith('close', expect.any(Function)); + }); + + it('finished replay ends the response immediately without going live', async () => { + const start = jest.fn(); + const attachment = makeAttachment({ + replay: ['f1'], + finished: true, + start, + }); + const { controller } = makeController({ chat: owned, attachment }); + const { res, raw } = makeRawRes(); + const { req } = makeReq(); + await controller.attachRunStream( + 'c1', + 'live', + 'a1', + req, + res, + user, + workspace, + ); + expect(raw.written).toEqual(['f1']); + expect(raw.end).toHaveBeenCalled(); + expect(start).not.toHaveBeenCalled(); // finished -> returns before start() + }); + + it('a close during the awaits (req.raw.destroyed) unsubscribes and writes nothing', async () => { + const attachment = makeAttachment({ replay: ['f1'] }); + const { controller } = makeController({ chat: owned, attachment }); + const { res, raw } = makeRawRes(); + const { req } = makeReq(true); // destroyed already at registration time + await controller.attachRunStream( + 'c1', + undefined, + undefined, + req, + res, + user, + workspace, + ); + expect(attachment.unsubscribe).toHaveBeenCalled(); + expect(raw.writeHead).not.toHaveBeenCalled(); + expect(raw.written).toEqual([]); + }); + + it('the registered close handler unsubscribes the attachment', async () => { + const attachment = makeAttachment({ replay: [] }); + const { controller } = makeController({ chat: owned, attachment }); + const { res } = makeRawRes(); + const { req, fireClose } = makeReq(); + await controller.attachRunStream( + 'c1', + undefined, + undefined, + req, + res, + user, + workspace, + ); + expect(attachment.unsubscribe).not.toHaveBeenCalled(); + fireClose(); // socket closed + expect(attachment.unsubscribe).toHaveBeenCalled(); + }); + + it('onFrame destroys the socket when the buffered length exceeds the cap', async () => { + const attachment = makeAttachment({ replay: [] }); + const { controller, getCb } = makeController({ chat: owned, attachment }); + const { res, raw } = makeRawRes(); + const { req } = makeReq(); + await controller.attachRunStream( + 'c1', + undefined, + undefined, + req, + res, + user, + workspace, + ); + const cb = getCb(); + // Normal live frame writes. + raw.writableLength = 0; + cb.onFrame('live-1'); + expect(raw.written).toContain('live-1'); + // A stalled socket over the cap is destroyed instead of buffering. + raw.writableLength = SUBSCRIBER_MAX_BUFFERED_BYTES + 1; + raw.write.mockClear(); + cb.onFrame('too-much'); + expect(raw.destroy).toHaveBeenCalled(); + expect(raw.write).not.toHaveBeenCalled(); + }); +}); 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 18b509cc..8cc7b03d 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.ts @@ -4,11 +4,15 @@ import { ConflictException, Controller, ForbiddenException, + Get, HttpCode, HttpException, HttpStatus, Logger, + Param, + ParseUUIDPipe, Post, + Query, Req, Res, ServiceUnavailableException, @@ -54,6 +58,12 @@ import { } from './dto/ai-chat.dto'; import { describeProviderError } from '../../integrations/ai/ai-error.util'; import { buildChatMarkdown } from './chat-markdown.util'; +import { + AiChatStreamRegistryService, + SUBSCRIBER_MAX_BUFFERED_BYTES, +} from './ai-chat-stream-registry.service'; +import { startSseHeartbeat } from './sse-resilience'; +import { EnvironmentService } from '../../integrations/environment/environment.service'; /** * Per-user AI chat API (§6.1). Routes are POST to match this codebase's @@ -72,6 +82,11 @@ export class AiChatController { private readonly aiChatMessageRepo: AiChatMessageRepo, private readonly aiTranscription: AiTranscriptionService, private readonly pageRepo: PageRepo, + // #184 phase 1.5. OPTIONAL so existing positional constructions (controller + // specs) compile unchanged; Nest always injects the real providers in + // production. Only touched on the resumable-stream (flag-on) path. + private readonly streamRegistry?: AiChatStreamRegistryService, + private readonly environment?: EnvironmentService, ) {} /** List the requesting user's chats in this workspace (paginated). */ @@ -233,6 +248,102 @@ export class AiChatController { return { stopped }; } + /** + * 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. + * + * `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. + */ + @SkipTransform() + @UseGuards(JwtAuthGuard, UserThrottlerGuard) + @Throttle({ [AI_CHAT_THROTTLER]: { limit: 60, ttl: 60000 } }) + @Get('runs/:chatId/stream') + async attachRunStream( + @Param('chatId', new ParseUUIDPipe()) chatId: string, + @Query('expect') expect: string | undefined, + @Query('anchor') anchor: 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 + 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(); + } + }, + 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; + } + res.hijack(); + // Cleanup BEFORE any write (invariant 5): a torn-down socket must not orphan + // a paused subscriber whose pending queue would buffer the whole run. + req.raw.once('close', () => { + attachment.unsubscribe(); + stopHeartbeat(); + }); + // A close emitted DURING the awaits above was missed by the listener — check. + // (Healthy pending GETs have req.raw.destroyed === false, so no false + // positives; returning without end() is fine — the socket is gone.) + if (req.raw.destroyed) { + attachment.unsubscribe(); + return; + } + res.raw.on('error', () => undefined); + try { + res.raw.writeHead(200, { + 'content-type': 'text/event-stream', + 'cache-control': 'no-cache', + 'x-vercel-ai-ui-message-stream': 'v1', + 'x-accel-buffering': 'no', + // deliberately NO Connection/Keep-Alive (hop-by-hop; Safari/HTTP2) + }); + res.raw.flushHeaders?.(); + for (const frame of attachment.replay) res.raw.write(frame); + if (attachment.finished) { + res.raw.end(); + return; + } + stopHeartbeat = startSseHeartbeat(res.raw, 15_000); + attachment.start(); // drain pending accumulated during replay, go live + } catch { + attachment.unsubscribe(); + stopHeartbeat(); + res.raw.destroy(); + } + } + /** Rename a chat. */ @HttpCode(HttpStatus.OK) @Post('rename') @@ -344,13 +455,25 @@ export class AiChatController { // its progress, and settle its terminal status — see AiChatRunService. const runHooks: AiChatRunHooks | undefined = autonomousRuns ? { - begin: (chatId) => - this.aiChatRunService.beginRun({ + begin: async (chatId) => { + const handle = await this.aiChatRunService.beginRun({ chatId, workspaceId: workspace.id, userId: user.id, trigger: 'user', - }), + }); + // #184 phase 1.5: register the run-stream entry at BEGIN (before any + // frame) so a tab that attaches in the begin->seed window finds an + // entry to wait on. Gated on AI_CHAT_RESUMABLE_STREAM: with the flag + // off nothing is registered and attach always 204s. + if ( + handle?.runId && + this.environment?.isAiChatResumableStreamEnabled?.() + ) { + this.streamRegistry?.open(chatId, handle.runId); + } + return handle; + }, onAssistantSeeded: (runId, messageId) => this.aiChatRunService.linkAssistantMessage( runId, diff --git a/apps/server/src/core/ai-chat/ai-chat.module.ts b/apps/server/src/core/ai-chat/ai-chat.module.ts index c0e8aa02..0a634b78 100644 --- a/apps/server/src/core/ai-chat/ai-chat.module.ts +++ b/apps/server/src/core/ai-chat/ai-chat.module.ts @@ -4,6 +4,7 @@ import { TokenModule } from '../auth/token.module'; import { AiChatController } from './ai-chat.controller'; import { AiChatService } from './ai-chat.service'; import { AiChatRunService } from './ai-chat-run.service'; +import { AiChatStreamRegistryService } from './ai-chat-stream-registry.service'; import { AiTranscriptionService } from './ai-transcription.service'; import { AiChatToolsService } from './tools/ai-chat-tools.service'; import { EmbeddingModule } from './embedding/embedding.module'; @@ -44,6 +45,7 @@ import { PublicShareChatToolsService } from './tools/public-share-chat-tools.ser providers: [ AiChatService, AiChatRunService, + AiChatStreamRegistryService, AiTranscriptionService, AiChatToolsService, PublicShareChatService, diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index 5869d7a7..19e82cee 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -1,4 +1,13 @@ -import { ForbiddenException } from '@nestjs/common'; +import { ForbiddenException, Logger } from '@nestjs/common'; +// Mock ONLY streamText so a driven stream() call can capture the pipe-options +// object (consumeSseStream / generateMessageId). Everything else in the AI SDK +// stays REAL (requireActual), so the pure-helper suites in this file are +// unaffected — none of them call stream()/streamText. +jest.mock('ai', () => ({ + ...jest.requireActual('ai'), + streamText: jest.fn(), +})); +import { streamText } from 'ai'; import { AiChatService, compactToolOutput, @@ -1059,3 +1068,181 @@ describe('isInterruptResume', () => { expect(isInterruptResume(withPrev(null), true)).toBe(false); }); }); + +/** + * #184 phase 1.5 — the run-wrapped pipe options (unit). Drives stream() to the + * pipe call with streamText mocked, capturing the options object, and asserts: + * - flag OFF while a runId IS present -> the LEGACY option shape (no + * consumeSseStream, no generateMessageId), and the registry is never touched. + * This is the exact dormancy guarantee this PR rests on. + * - flag ON + runId -> consumeSseStream tees into the registry and + * generateMessageId returns the seeded assistant DB row id. + * - flag ON but no runHooks (runId undefined) -> legacy (the runId gate). + * - flag ON + runId -> the outer catch releases the entry via abortEntry. + */ +describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', () => { + const streamTextMock = streamText as unknown as jest.Mock; + let pipeMock: jest.Mock; + + beforeEach(() => { + streamTextMock.mockReset(); + pipeMock = jest.fn(); + streamTextMock.mockReturnValue({ + consumeStream: jest.fn(), + pipeUIMessageStreamToResponse: pipeMock, + }); + // Silence the service's diagnostic logging for a clean test run. + jest.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined as never); + jest + .spyOn(Logger.prototype, 'error') + .mockImplementation(() => undefined as never); + jest + .spyOn(Logger.prototype, 'warn') + .mockImplementation(() => undefined as never); + }); + + afterEach(() => jest.restoreAllMocks()); + + // A raw-response stub sufficient for the post-streamText wiring. + function makeRes() { + return { + raw: { + writeHead: jest.fn(), + write: jest.fn(), + once: jest.fn(), + on: jest.fn(), + flushHeaders: jest.fn(), + writableEnded: false, + destroyed: false, + }, + }; + } + + // Wire only the deps reached on the way to the pipe call, plus a spy registry. + function makeService(opts: { resumable: boolean }) { + const aiChatRepo = { + findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })), + insert: jest.fn(), + }; + const aiChatMessageRepo = { + // Both the user insert and the assistant seed return the same row id. + insert: jest.fn(async () => ({ id: 'msg-1' })), + findAllByChat: jest.fn(async () => []), + update: jest.fn(async () => ({ id: 'msg-1' })), + }; + const aiSettings = { resolve: jest.fn(async () => ({})) }; + const tools = { forUser: jest.fn(async () => ({})) }; + const mcpClients = { + toolsFor: jest.fn(async () => ({ + tools: {}, + clients: [], + outcomes: [], + instructions: [], + })), + }; + const streamRegistry = { + open: jest.fn(), + bind: jest.fn(), + abortEntry: jest.fn(), + }; + const svc = new AiChatService( + {} as never, // ai (model is injected) + aiChatRepo as never, + aiChatMessageRepo as never, + {} as never, // aiChatPageSnapshotRepo + aiSettings as never, + tools as never, + mcpClients as never, + {} as never, // aiAgentRoleRepo + {} as never, // pageRepo (openPage undefined -> never touched) + {} as never, // pageAccess + { + isAiChatDeferredToolsEnabled: () => false, + isAiChatResumableStreamEnabled: () => opts.resumable, + } as never, + streamRegistry as never, + ); + return { svc, streamRegistry }; + } + + const body = { + chatId: 'chat-1', + messages: [ + { id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] }, + ], + }; + + const makeRunHooks = () => ({ + begin: jest.fn(async () => ({ + runId: 'run-1', + signal: new AbortController().signal, + })), + onAssistantSeeded: jest.fn(), + onStep: jest.fn(), + onSettled: jest.fn(), + }); + + async function drive(svc: AiChatService, hooks: unknown): Promise { + await svc.stream({ + user: { id: 'u1' } as never, + workspace: { id: 'ws-1' } as never, + sessionId: 's1', + body: body as never, + res: makeRes() as never, + signal: new AbortController().signal, + model: {} as never, + role: null, + runHooks: hooks as never, + }); + } + + it('flag OFF + runId present: LEGACY option shape (no consumeSseStream / generateMessageId); registry untouched', async () => { + const { svc, streamRegistry } = makeService({ resumable: false }); + await drive(svc, makeRunHooks()); + expect(pipeMock).toHaveBeenCalledTimes(1); + const options = pipeMock.mock.calls[0][1]; + // The dormancy guarantee: a live run with the flag off tees NOTHING and does + // not stamp a message id — byte-for-byte the pre-1.5 wire. + expect(options.consumeSseStream).toBeUndefined(); + expect(options.generateMessageId).toBeUndefined(); + expect(streamRegistry.bind).not.toHaveBeenCalled(); + expect(streamRegistry.abortEntry).not.toHaveBeenCalled(); + }); + + it('flag ON + runId: consumeSseStream tees into the registry; generateMessageId returns the seeded row id', async () => { + const { svc, streamRegistry } = makeService({ resumable: true }); + await drive(svc, makeRunHooks()); + const options = pipeMock.mock.calls[0][1]; + expect(typeof options.consumeSseStream).toBe('function'); + expect(typeof options.generateMessageId).toBe('function'); + // generateMessageId stamps the seeded assistant DB row id. + expect(options.generateMessageId()).toBe('msg-1'); + // consumeSseStream binds the tee: (chatId, runId, assistantId, stream). + const fakeStream = {} as ReadableStream; + options.consumeSseStream({ stream: fakeStream }); + expect(streamRegistry.bind).toHaveBeenCalledWith( + 'chat-1', + 'run-1', + 'msg-1', + fakeStream, + ); + }); + + it('flag ON but NO runHooks (runId undefined): pipe options stay legacy (the runId gate)', async () => { + const { svc, streamRegistry } = makeService({ resumable: true }); + await drive(svc, undefined); + const options = pipeMock.mock.calls[0][1]; + expect(options.consumeSseStream).toBeUndefined(); + expect(options.generateMessageId).toBeUndefined(); + expect(streamRegistry.bind).not.toHaveBeenCalled(); + }); + + it('flag ON + runId: the outer catch calls abortEntry when the stream throws', async () => { + const { svc, streamRegistry } = makeService({ resumable: true }); + streamTextMock.mockImplementation(() => { + throw new Error('boom'); + }); + await expect(drive(svc, makeRunHooks())).rejects.toThrow('boom'); + expect(streamRegistry.abortEntry).toHaveBeenCalledWith('chat-1', 'run-1'); + }); +}); 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 5b3e47e1..d8e027b9 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -32,6 +32,7 @@ import { import { AiChatToolsService } from './tools/ai-chat-tools.service'; import { McpClientsService } from './external-mcp/mcp-clients.service'; import { EnvironmentService } from '../../integrations/environment/environment.service'; +import { AiChatStreamRegistryService } from './ai-chat-stream-registry.service'; import { buildSystemPrompt } from './ai-chat.prompt'; import { CORE_TOOL_KEYS, @@ -276,6 +277,10 @@ export class AiChatService implements OnModuleInit { // Reads the AI_CHAT_DEFERRED_TOOLS toggle (#332). Injected last so existing // positional constructor callers (tests) only append one stub. private readonly environment: EnvironmentService, + // #184 phase 1.5 run-stream registry. OPTIONAL so existing positional + // constructions (int-specs) compile unchanged; Nest always injects the real + // provider in production. Only ever touched on the run-wrapped + flag-on path. + private readonly streamRegistry?: AiChatStreamRegistryService, ) {} /** @@ -1174,6 +1179,35 @@ export class AiChatService implements OnModuleInit { // as the cumulative authoritative usage so the client never jumps DOWN. let cumulativeStepUsage: ChatStreamUsage | undefined; result.pipeUIMessageStreamToResponse(res.raw, { + // #184 phase 1.5: run-wrapped mode only — the legacy path (flag off) stays + // byte-for-byte identical, including the absence of start.messageId. Both + // fields are gated on `runId` (present only for a durable run) AND the + // AI_CHAT_RESUMABLE_STREAM flag; the seed `assistantId` is unconditional, + // so gating on `assistantId` alone would change the legacy wire. + ...(runId && this.environment?.isAiChatResumableStreamEnabled?.() + ? { + // Tee the SSE frames into the run-stream registry so late tabs can + // attach (replay + live tail). + consumeSseStream: ({ + stream, + }: { + stream: ReadableStream; + }) => + this.streamRegistry?.bind( + chatId, + runId!, + assistantId, + stream, + ), + // Stamp the persisted assistant row's DB id onto the streamed + // message so every tab renders the SAME id as the DB row (id-based + // reconciliation). Seeding is best-effort: when it failed, let the + // client generate the id. + ...(assistantId + ? { generateMessageId: () => assistantId } + : {}), + } + : {}), headers: { 'X-Accel-Buffering': 'no' }, // Surface the authoritative chatId on the streamed assistant UI message so // the client adopts the REAL id of the row we created, instead of guessing @@ -1239,6 +1273,12 @@ export class AiChatService implements OnModuleInit { // finalizeRun (onSettled) is idempotent — a settle here and a settle from a // streamText callback collapse to a single terminal write. if (runId) { + // #184 phase 1.5: a failure here means the tee `done` will never arrive, + // so release the registry entry's subscribers explicitly — otherwise an + // attached tab hangs forever. Same flag gate as the tee wiring above. + if (this.environment?.isAiChatResumableStreamEnabled?.()) { + this.streamRegistry?.abortEntry(chatId, runId); + } await runHooks?.onSettled?.( runId, 'error', diff --git a/apps/server/src/integrations/environment/environment.service.ts b/apps/server/src/integrations/environment/environment.service.ts index 6cad98bf..08ff8dd3 100644 --- a/apps/server/src/integrations/environment/environment.service.ts +++ b/apps/server/src/integrations/environment/environment.service.ts @@ -292,6 +292,23 @@ export class EnvironmentService { return enabled === 'true'; } + /** + * Resumable SSE transport for durable agent runs (#184 phase 1.5). When + * enabled, a run tees its SSE frames into the in-memory run-stream registry so + * a late/reloaded tab can attach (replay + live tail) via + * `GET /ai-chat/runs/:chatId/stream`. Defaults to DISABLED: PR 1 ships the + * server code dormant — with the flag off, `open`/`bind`/`generateMessageId` + * are never called and attach always answers 204, so the legacy and #184 + * phase-1 wire paths stay byte-for-byte identical. Set + * AI_CHAT_RESUMABLE_STREAM=true to activate it (paired with the PR 2 client). + */ + isAiChatResumableStreamEnabled(): boolean { + const enabled = this.configService + .get('AI_CHAT_RESUMABLE_STREAM', 'false') + .toLowerCase(); + return enabled === 'true'; + } + getPostHogHost(): string { return this.configService.get('POSTHOG_HOST'); } diff --git a/apps/server/test/integration/ai-chat-attach.int-spec.ts b/apps/server/test/integration/ai-chat-attach.int-spec.ts new file mode 100644 index 00000000..7faad55f --- /dev/null +++ b/apps/server/test/integration/ai-chat-attach.int-spec.ts @@ -0,0 +1,564 @@ +import * as http from 'node:http'; +import { Kysely } from 'kysely'; +import { + MockLanguageModelV3, + convertArrayToReadableStream, +} from 'ai/test'; +import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo'; +import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; +import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo'; +import { AiChatService } from 'src/core/ai-chat/ai-chat.service'; +import { AiChatRunService } from 'src/core/ai-chat/ai-chat-run.service'; +import { + AiChatStreamRegistryService, + RunStreamCallbacks, +} from 'src/core/ai-chat/ai-chat-stream-registry.service'; +import { + getTestDb, + destroyTestDb, + createWorkspace, + createUser, + createChat, +} from './db'; + +/** + * #184 phase 1.5 — the resumable transport end to end against REAL Postgres, + * the REAL `streamText` (seeded via MockLanguageModelV3) and a REAL Node + * ServerResponse, driving the REAL `AiChatService.stream` run-wrapped path with a + * REAL `AiChatStreamRegistryService`. The run-hooks mirror the controller: they + * begin a durable run and `open()` the registry entry at begin, and the service + * 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. + */ + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +async function waitFor( + cond: () => Promise | boolean, + { timeoutMs = 15_000, stepMs = 25 } = {}, +): Promise { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (await cond()) return; + await sleep(stepMs); + } + throw new Error('waitFor: condition not met within timeout'); +} + +// A real Node ServerResponse wired to a live socket (as in the stream int-spec). +function makeRealResponse(): Promise<{ + res: http.ServerResponse; + cleanup: () => Promise; +}> { + return new Promise((resolve) => { + const server = http.createServer((_req, res) => { + resolve({ + res, + cleanup: () => + new Promise((done) => { + try { + if (!res.writableEnded) res.end(); + } catch { + /* socket already gone */ + } + server.close(() => done()); + }), + }); + }); + server.listen(0, () => { + const port = (server.address() as any).port; + const creq = http.request({ port, method: 'GET' }, (cres) => { + cres.resume(); + }); + creq.on('error', () => undefined); + creq.end(); + }); + }); +} + +// A full, successful single-step turn. +function successStream() { + return convertArrayToReadableStream([ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: 't1' }, + { type: 'text-delta', id: 't1', delta: 'Hello' }, + { type: 'text-delta', id: 't1', delta: ' there' }, + { type: 'text-end', id: 't1' }, + { + type: 'finish', + finishReason: 'stop', + usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 }, + }, + ] as any); +} + +// A stream the test feeds chunk by chunk (to attach mid-flight / drive a stop). +function makeControlledStream() { + let controller!: ReadableStreamDefaultController; + const stream = new ReadableStream({ + start(c) { + controller = c; + }, + }); + return { + stream, + emit: (chunk: any) => controller.enqueue(chunk), + close: () => controller.close(), + }; +} + +// Collect replay + live frames from an attachment. +function liveSink(): { + cb: RunStreamCallbacks; + frames: string[]; + ended: () => boolean; +} { + const frames: string[] = []; + let ended = false; + return { + frames, + ended: () => ended, + cb: { + onFrame: (f) => frames.push(f), + onEnd: () => { + ended = true; + }, + }, + }; +} + +// The SSE `start` frame carries the message id; pull it out of a `data: {...}`. +function parseStartMessageId(frames: string[]): string | 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; + } catch { + /* not this frame */ + } + } + return undefined; +} + +describe('AiChatService run-stream attach [integration]', () => { + let db: Kysely; + let aiChatRepo: AiChatRepo; + let msgRepo: AiChatMessageRepo; + let runRepo: AiChatRunRepo; + let workspaceId: string; + let userId: string; + + const mcpClients = { + toolsFor: async () => ({ + tools: {}, + clients: [], + outcomes: [], + instructions: [], + }), + }; + + // Build the service with the run-stream registry wired and the resumable flag + // ON (the property under test). Deferred tools OFF (irrelevant here). + function buildService(registry: AiChatStreamRegistryService): AiChatService { + return new AiChatService( + { getChatModel: async () => null } as any, + aiChatRepo, + msgRepo, + {} as any, + { resolve: async () => null } as any, + { forUser: async () => ({}) } as any, + mcpClients as any, + {} as any, + {} as any, + {} as any, + { + isAiChatDeferredToolsEnabled: () => false, + isAiChatResumableStreamEnabled: () => true, + } as any, + registry, + ); + } + + // Run-hooks mirroring the controller: begin the durable run AND open() the + // registry entry at begin. Captures the runId so a test can stop it. + function makeRunHooks( + runService: AiChatRunService, + registry: AiChatStreamRegistryService, + box: { runId?: string }, + ) { + return { + begin: async (chatId: string) => { + const handle = await runService.beginRun({ + chatId, + workspaceId, + userId, + trigger: 'user', + }); + box.runId = handle.runId; + registry.open(chatId, handle.runId); + return handle; + }, + onAssistantSeeded: (runId: string, messageId: string) => + runService.linkAssistantMessage(runId, workspaceId, messageId), + onStep: (runId: string, n: number) => + void runService.recordStep(runId, workspaceId, n), + onSettled: (runId: string, status: any, error?: string) => + runService.finalizeRun(runId, workspaceId, status, error), + }; + } + + function userUiMessage(text: string) { + return { + id: `u-${Math.random()}`, + role: 'user', + parts: [{ type: 'text', text }], + }; + } + + async function startRun(opts: { + registry: AiChatStreamRegistryService; + runService?: AiChatRunService; + model: MockLanguageModelV3; + chatId: string; + body: any; + box?: { runId?: string }; + }): Promise<{ res: http.ServerResponse; cleanup: () => Promise }> { + const service = buildService(opts.registry); + const { res, cleanup } = await makeRealResponse(); + const runHooks = opts.runService + ? makeRunHooks(opts.runService, opts.registry, opts.box ?? {}) + : undefined; + await service.stream({ + user: { id: userId, workspaceId } as any, + workspace: { id: workspaceId, name: 'WS' } as any, + sessionId: 'sess-1', + body: opts.body, + res: { raw: res } as any, + signal: new AbortController().signal, + model: opts.model as any, + role: null, + runHooks, + } as any); + return { res, cleanup }; + } + + async function assistantRowId(chatId: string): Promise { + const rows = await msgRepo.findAllByChat(chatId, workspaceId); + const row = rows.find((r: any) => r.role === 'assistant'); + return row!.id as string; + } + + beforeAll(async () => { + db = getTestDb(); + aiChatRepo = new AiChatRepo(db as any); + msgRepo = new AiChatMessageRepo(db as any); + runRepo = new AiChatRunRepo(db as any); + workspaceId = (await createWorkspace(db)).id; + userId = (await createUser(db, workspaceId)).id; + }); + + afterAll(async () => { + await destroyTestDb(); + }); + + it('run-wrapped: replay is the full frame sequence incl [DONE], start.messageId == the seeded DB row id', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + const registry = new AiChatStreamRegistryService(); + const runService = new AiChatRunService(runRepo, { + isCloud: () => false, + } as never); + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: successStream() }), + } as any); + + const box: { runId?: string } = {}; + const { cleanup } = await startRun({ + registry, + runService, + model, + chatId, + body: { chatId, messages: [userUiMessage('Hi')] }, + box, + }); + try { + // Wait for the assistant row to settle (terminal callbacks run async). + await waitFor(async () => { + const rows = await msgRepo.findAllByChat(chatId, workspaceId); + return rows.some( + (r: any) => + r.role === 'assistant' && + ['completed', 'error', 'aborted'].includes(r.status), + ); + }); + const rowId = await assistantRowId(chatId); + + // Finished-run replay with expect=live + the correct anchor. + const sink = liveSink(); + const att = await registry.attach(chatId, true, rowId, 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'); + expect(att!.replay.some((f) => f.includes('[DONE]'))).toBe(true); + } finally { + registry.onModuleDestroy(); + await cleanup(); + } + }); + + it('anchor mismatch with expect=live returns null (invariant 6)', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + const registry = new AiChatStreamRegistryService(); + const runService = new AiChatRunService(runRepo, { + isCloud: () => false, + } as never); + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: successStream() }), + } as any); + + const { cleanup } = await startRun({ + registry, + runService, + model, + chatId, + body: { chatId, messages: [userUiMessage('Hi')] }, + }); + try { + await waitFor(async () => { + const rows = await msgRepo.findAllByChat(chatId, workspaceId); + return rows.some( + (r: any) => r.role === 'assistant' && r.status === 'completed', + ); + }); + 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), + ).toBeNull(); + } finally { + registry.onModuleDestroy(); + await cleanup(); + } + }); + + it('an attach opened BEFORE the first frame follows the live stream from frame 0', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + const registry = new AiChatStreamRegistryService(); + const runService = new AiChatRunService(runRepo, { + isCloud: () => false, + } as never); + const controlled = makeControlledStream(); + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: controlled.stream }), + } as any); + + const { cleanup } = await startRun({ + registry, + runService, + model, + chatId, + body: { chatId, messages: [userUiMessage('Slow please')] }, + }); + 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 + att.start(); // go live (drains nothing, then follows) + + // Now emit the whole turn. + controlled.emit({ type: 'stream-start', warnings: [] }); + controlled.emit({ type: 'text-start', id: 't1' }); + controlled.emit({ type: 'text-delta', id: 't1', delta: 'Zero' }); + controlled.emit({ type: 'text-end', id: 't1' }); + controlled.emit({ + type: 'finish', + finishReason: 'stop', + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + }); + controlled.close(); + + await waitFor(() => sink.frames.some((f) => f.includes('[DONE]'))); + // The subscriber saw the stream from the very first frame (`start`) through + // the terminal marker, with the streamed text present. + expect(sink.frames.some((f) => f.includes('"type":"start"'))).toBe(true); + expect(sink.frames.join('')).toContain('Zero'); + expect(sink.frames[sink.frames.length - 1]).toContain('[DONE]'); + expect(sink.ended()).toBe(true); + } finally { + registry.onModuleDestroy(); + await cleanup(); + } + }); + + it('requestStop surfaces {"type":"abort"} + [DONE] + end to the attached subscriber', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + const registry = new AiChatStreamRegistryService(); + const runService = new AiChatRunService(runRepo, { + isCloud: () => false, + } as never); + // An abort-AWARE model: it streams some partial output, then errors the model + // stream with an AbortError when the run signal aborts — exactly as a real + // provider network stream is torn down on abort (a plain in-memory stream + // would just stall, so streamText would never observe the stop). + const model = new MockLanguageModelV3({ + doStream: async ({ abortSignal }: any) => { + const stream = new ReadableStream({ + start(controller) { + controller.enqueue({ type: 'stream-start', warnings: [] }); + controller.enqueue({ type: 'text-start', id: 't1' }); + controller.enqueue({ type: 'text-delta', id: 't1', delta: 'partial' }); + abortSignal?.addEventListener('abort', () => { + try { + controller.error( + new DOMException('Aborted', 'AbortError'), + ); + } catch { + /* already errored/closed */ + } + }); + }, + }); + return { stream }; + }, + } as any); + + const box: { runId?: string } = {}; + const { cleanup } = await startRun({ + registry, + runService, + model, + chatId, + body: { chatId, messages: [userUiMessage('Start then stop')] }, + box, + }); + try { + const sink = liveSink(); + const att = (await registry.attach(chatId, false, undefined, sink.cb))!; + att.start(); + + // Give streamText a beat to begin consuming the partial output. + await sleep(250); + // User presses Stop -> the run signal aborts -> the SDK emits an abort chunk. + await runService.requestStop(box.runId!, workspaceId); + + await waitFor(() => sink.ended()); + expect(sink.frames.some((f) => f.includes('"type":"abort"'))).toBe(true); + expect(sink.frames.some((f) => f.includes('[DONE]'))).toBe(true); + expect(sink.ended()).toBe(true); + } finally { + registry.onModuleDestroy(); + await cleanup(); + } + }); + + it('the outer catch calls abortEntry so an open entry is released (finished)', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + const registry = new AiChatStreamRegistryService(); + const runService = new AiChatRunService(runRepo, { + isCloud: () => false, + } as never); + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: successStream() }), + } as any); + + // A msgRepo whose user-row insert throws: the turn fails AFTER begin (the + // entry is already open) but BEFORE the pipe, exercising the outer catch. + const throwingMsgRepo = { + insert: async () => { + throw new Error('db boom'); + }, + findAllByChat: (...a: any[]) => (msgRepo as any).findAllByChat(...a), + update: (...a: any[]) => (msgRepo as any).update(...a), + findById: (...a: any[]) => (msgRepo as any).findById(...a), + }; + const service = new AiChatService( + { getChatModel: async () => null } as any, + aiChatRepo, + throwingMsgRepo as any, + {} as any, + { resolve: async () => null } as any, + { forUser: async () => ({}) } as any, + mcpClients as any, + {} as any, + {} as any, + {} as any, + { + isAiChatDeferredToolsEnabled: () => false, + isAiChatResumableStreamEnabled: () => true, + } as any, + registry, + ); + const { res, cleanup } = await makeRealResponse(); + const box: { runId?: string } = {}; + try { + await expect( + service.stream({ + user: { id: userId, workspaceId } as any, + workspace: { id: workspaceId, name: 'WS' } as any, + sessionId: 'sess-1', + body: { chatId, messages: [userUiMessage('will throw')] }, + res: { raw: res } as any, + signal: new AbortController().signal, + model: model as any, + role: null, + runHooks: makeRunHooks(runService, registry, box), + } as any), + ).rejects.toThrow(); + // The entry opened at begin was terminated by abortEntry (from the catch), + // so it is finished and a plain attach returns null instead of hanging. + const entry = (registry as any).entries.get(chatId); + expect(entry).toBeDefined(); + expect(entry.finished).toBe(true); + const sink = liveSink(); + expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull(); + } finally { + registry.onModuleDestroy(); + await cleanup(); + } + }); + + it('legacy (no run-hooks): the registry is never populated', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + const registry = new AiChatStreamRegistryService(); + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: successStream() }), + } as any); + + // No runHooks -> runId undefined -> the run-wrapped tee is never wired. + const { cleanup } = await startRun({ + registry, + model, + chatId, + body: { chatId, messages: [userUiMessage('Legacy hi')] }, + }); + try { + await waitFor(async () => { + const rows = await msgRepo.findAllByChat(chatId, workspaceId); + return rows.some( + (r: any) => r.role === 'assistant' && r.status === 'completed', + ); + }); + 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(); + } finally { + registry.onModuleDestroy(); + await cleanup(); + } + }); +});