feat(ai-chat): resumable SSE run-stream registry — server, dormant (#381 PR 1)

PR 1 of 2 (#381, фаза 1.5 #184): серверный реестр SSE-стримов агентских ранов, чтобы любая вкладка могла подключиться к живому рану с реплеем кадров + живым хвостом. «Спящий» — весь провод за флагом AI_CHAT_RESUMABLE_STREAM (off по умолчанию); клиент (PR 2) ещё не написан. - ai-chat-stream-registry.service.ts: in-memory реестр (open/bind/abortEntry/ attach). attach — снапшот+подписка в ОДНОМ синхронном блоке (инвариант 4: нет await между `subscribers.add` и `frames.slice()`), paused-подписчик, overflow, retention с identity-guard (инвариант 2), open поверх live entry даёт ровно один onEnd (инвариант 3), anchor против кросс-ранового реплея (инвариант 6). - ai-chat.controller.ts: begin-хук open(chatId, runId) + GET-attach эндпоинт (403 чужой чат; 204 нет-entry/finished/anchor-мисматч; cleanup до первой записи + recheck req.raw.destroyed; cap→destroy). - ai-chat.service.ts: tee SSE-кадров в реестр (consumeSseStream + generateMessageId, гейт на runId && flag) + abortEntry из внешнего catch. - environment.service.ts: флаг isAiChatResumableStreamEnabled(). Флаг OFF ⇒ байт-в-байт legacy И #184-фаза-1 (нет start.messageId, нет tee). Инжектируемые провайдеры НЕ @Optional() → поломка вайринга роняет старт, а не тихо выключает фичу. Тесты: registry unit (16), controller.attach (9), service pipe-options (4, вкл. flag-off-with-runId негатив), int-spec ai-chat-attach (6, реальный MockLanguageModelV3). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-06 07:22:20 +03:00
parent 3085ec1b50
commit 52ee3c1f3e
9 changed files with 1912 additions and 4 deletions
@@ -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> | boolean,
+  { timeoutMs = 15_000, stepMs = 25 } = {},
+): Promise<void> {
+  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<void>;
+}> {
+  return new Promise((resolve) => {
+    const server = http.createServer((_req, res) => {
+      resolve({
+        res,
+        cleanup: () =>
+          new Promise<void>((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<any>;
+  const stream = new ReadableStream<any>({
+    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<any>;
+  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<void> }> {
+    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<string> {
+    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<any>({
+          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();
+    }
+  });
+});