fix(ai-chat): settle detached runs on pre-stream failures + review fixes (#184)

CRITICAL: any failure between a successful beginRun and streamText's terminal
callbacks taking ownership (the bare awaits: user-message insert, history load,
convertToModelMessages, settings resolve; the buildSystemPrompt/forUser block;
and synchronous streamText wiring) left ai_chat_runs stuck 'running' forever
(sweepRunning only runs at startup), which then 409'd every future turn in the
chat and made the observer tab poll forever. Wrap the body of stream() after
beginRun in a safety-net try/catch that settles the run to 'error' (via
onSettled) before rethrowing, and make finalizeRun idempotent (active.delete is
the once-guard) so a settle here and a settle from a streamText callback collapse
to a single terminal write.

Also from review comment 2519:
- correct three client comments that falsely claimed /ai-chat/run is "flag-gated
  server-side and would 403" — it is owner-gated only; with the feature off the
  chat simply has no runs so the endpoint returns { run: null }
  (ai-chat-window.tsx, ai-chat-service.ts, ai-chat-query.ts).
- remove the dead UpdatableAiChatRun type (zero usages; the repo update uses an
  inline Partial<...>).
- add controller specs for POST /ai-chat/run and /ai-chat/stop (owner-gating,
  run:null when no run, run+message, stop by runId and by chatId).
- add tests: an exception after beginRun settles the run to 'error' and drops the
  in-memory entry (next turn is not 409'd); finalizeRun is idempotent.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

apps/client/src/features/ai-chat/components/ai-chat-window.tsx

@@ -166,8 +166,10 @@ export default function AiChatWindow() {
     useAiChatMessagesQuery(activeChatId ?? undefined);
 
   // #184 reconnect-and-live-follow. Whether detached agent runs are enabled for
-  // this workspace; the reconnect endpoint is flag-gated server-side, so we must
-  // not poll it when the feature is off.
+  // this workspace. The reconnect endpoint itself is NOT flag-gated server-side
+  // (it is only owner-gated and returns `{ run: null }` when the chat has no
+  // run); but when the feature is off no runs are ever created, so polling it
+  // would always come back empty — we gate it off here to avoid pointless polls.
   const workspace = useAtomValue(workspaceAtom);
   const autonomousRunsEnabled =
     workspace?.settings?.ai?.autonomousRuns === true;

apps/client/src/features/ai-chat/queries/ai-chat-query.ts

@@ -60,11 +60,12 @@ export const AI_CHAT_RUN_RQ_KEY = (chatId: string) => ["ai-chat-run", chatId];
 export function useAiChatsQuery() {
   const query = useInfiniteQuery({
     queryKey: AI_CHATS_RQ_KEY,
-    queryFn: ({ pageParam }) =>
-      getAiChats({ cursor: pageParam, limit: 50 }),
+    queryFn: ({ pageParam }) => getAiChats({ cursor: pageParam, limit: 50 }),
     initialPageParam: undefined as string | undefined,
     getNextPageParam: (lastPage) =>
-      lastPage.meta.hasNextPage ? (lastPage.meta.nextCursor ?? undefined) : undefined,
+      lastPage.meta.hasNextPage
+        ? (lastPage.meta.nextCursor ?? undefined)
+        : undefined,
   });
 
   const data = useMemo<IPagination<IAiChat> | undefined>(() => {
@@ -94,7 +95,9 @@ export function useAiChatMessagesQuery(chatId: string | undefined) {
       getAiChatMessages({ chatId: chatId as string, cursor: pageParam }),
     initialPageParam: undefined as string | undefined,
     getNextPageParam: (lastPage) =>
-      lastPage.meta.hasNextPage ? (lastPage.meta.nextCursor ?? undefined) : undefined,
+      lastPage.meta.hasNextPage
+        ? (lastPage.meta.nextCursor ?? undefined)
+        : undefined,
     enabled: !!chatId,
   });
 
@@ -144,9 +147,10 @@ export function useAiChatMessagesQuery(chatId: string | undefined) {
  * is thus naturally bounded by the run terminating; no separate timeout cap.
  *
  * `enabled` gates the whole thing: callers pass `false` when the autonomous-runs
- * feature is off (the endpoint is flag-gated server-side and would 403) OR when
- * THIS tab is the one actively streaming the run (the live SSE owns the view, so
- * we must not also poll/merge). The global `retry: false` means a 403/anything
+ * feature is off (the endpoint is NOT flag-gated server-side, but with the feature
+ * off the chat has no runs, so polling would only ever return `{ run: null }`) OR
+ * when THIS tab is the one actively streaming the run (the live SSE owns the view,
+ * so we must not also poll/merge). The global `retry: false` means a failed fetch
  * leaves `data` undefined, so refetchInterval(undefined run) returns false — a
  * failed fetch can never spin a tight loop.
  */
@@ -311,11 +315,14 @@ export function useImportAiRolesFromCatalogMutation() {
     mutationFn: (payload) => importAiRolesFromCatalog(payload),
     onSuccess: (result) => {
       notifications.show({
-        message: t("Imported {{created}}, renamed {{renamed}}, skipped {{skipped}}", {
-          created: result.created,
-          renamed: result.renamed,
-          skipped: result.skipped,
-        }),
+        message: t(
+          "Imported {{created}}, renamed {{renamed}}, skipped {{skipped}}",
+          {
+            created: result.created,
+            renamed: result.renamed,
+            skipped: result.skipped,
+          },
+        ),
       });
       // Surface partial failures (e.g. unique-name races) as a red warning.
       if (result.errors.length > 0) {

apps/client/src/features/ai-chat/services/ai-chat-service.ts

@@ -49,7 +49,9 @@ export async function getAiChatMessages(
  * partial output while the run is in-flight, the final output once it finished).
  * The DB is the source of truth, so this works for an in-flight run (the browser
  * dropped, the run kept going) and a finished one alike; `{ run: null }` when the
- * chat has never had a run. Owner-gated and flag-gated server-side.
+ * chat has never had a run. Owner-gated server-side (the requesting user must own
+ * the chat); it is NOT flag-gated — when the feature is off the chat simply has no
+ * runs, so the endpoint returns `{ run: null }`.
  */
 export async function getAiChatRun(
   chatId: string,

apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts

@@ -12,10 +12,13 @@ function uniqueViolation(constraintName: string): Error & {
   code: string;
   constraint_name: string;
 } {
-  return Object.assign(new Error('duplicate key value violates unique constraint'), {
-    code: '23505',
-    constraint_name: constraintName,
-  });
+  return Object.assign(
+    new Error('duplicate key value violates unique constraint'),
+    {
+      code: '23505',
+      constraint_name: constraintName,
+    },
+  );
 }
 
 /**
@@ -260,6 +263,33 @@ describe('AiChatRunService run lifecycle', () => {
     );
   });
 
+  it('finalizeRun is IDEMPOTENT: a second settle no-ops (single terminal write)', async () => {
+    // The #184 review fix: AiChatService.stream wraps the turn in a safety-net
+    // catch that settles a failed turn AND streamText's terminal callback may
+    // also settle — both routes call finalizeRun. Only the FIRST may write the
+    // terminal row; the second must no-op so a late settle can never clobber the
+    // real terminal status or double-write the row.
+    const repo = makeRepo();
+    const svc = new AiChatRunService(repo as never);
+    await svc.beginRun({
+      chatId: 'chat-1',
+      workspaceId: 'ws-1',
+      userId: 'user-1',
+    });
+
+    await svc.finalizeRun('run-1', 'ws-1', 'error', 'first');
+    expect(svc.isLocallyActive('run-1')).toBe(false);
+    // A second settle (e.g. a streamText callback firing after the catch) no-ops.
+    await svc.finalizeRun('run-1', 'ws-1', 'completed', undefined);
+
+    expect(repo.update).toHaveBeenCalledTimes(1);
+    expect(repo.update).toHaveBeenCalledWith(
+      'run-1',
+      'ws-1',
+      expect.objectContaining({ status: 'failed', error: 'first' }),
+    );
+  });
+
   it('recordStep / linkAssistantMessage are best-effort: a repo failure is swallowed', async () => {
     const repo = makeRepo({
       update: jest.fn(async () => {

apps/server/src/core/ai-chat/ai-chat-run.service.ts

@@ -201,10 +201,19 @@ export class AiChatRunService implements OnModuleInit {
 
   /**
    * Finalize a run to its terminal status (succeeded / failed / aborted),
-   * stamping finishedAt + any error, and DROP its in-memory entry. Idempotent
-   * and best-effort: the at-most-once turn terminal callbacks call it, but a
-   * double call (or a call after the row was swept) merely re-writes the same
-   * terminal row.
+   * stamping finishedAt + any error, and DROP its in-memory entry. Best-effort.
+   *
+   * IDEMPOTENT (#184 review): the terminal write happens AT MOST ONCE per run.
+   * `this.active.delete(runId)` returns false when the run was already settled
+   * (its in-memory entry already dropped); in that case we no-op. This collapses
+   * a legitimate double-settle to a single write: AiChatService.stream wraps the
+   * turn in a safety-net catch that settles the run to 'error' on any failure
+   * BEFORE streamText's terminal callbacks own the lifecycle — and on the rare
+   * path where streamText DID attach (so a callback also settles) the two would
+   * otherwise both call onSettled. The first caller wins and writes the terminal
+   * row; the second returns early, so a late settle can never clobber the real
+   * terminal status or double-write. beginRun always registers the entry before
+   * the turn can settle, so a legitimate first finalize always finds it.
    */
   async finalizeRun(
     runId: string,
@@ -212,7 +221,7 @@ export class AiChatRunService implements OnModuleInit {
     turnStatus: TurnTerminalStatus,
     error?: string,
   ): Promise<void> {
-    this.active.delete(runId);
+    if (!this.active.delete(runId)) return;
     try {
       await this.runRepo.update(runId, workspaceId, {
         status: mapTurnStatusToRun(turnStatus),
@@ -258,10 +267,7 @@ export class AiChatRunService implements OnModuleInit {
 
   /** Fetch a run by id (workspace-scoped). Used to resolve + ownership-check an
    *  explicit stop targeting a runId. */
-  getRun(
-    runId: string,
-    workspaceId: string,
-  ): Promise<AiChatRun | undefined> {
+  getRun(runId: string, workspaceId: string): Promise<AiChatRun | undefined> {
     return this.runRepo.findById(runId, workspaceId);
   }
 

apps/server/src/core/ai-chat/ai-chat.controller.run.spec.ts

@@ -0,0 +1,163 @@
+import { BadRequestException, ForbiddenException } from '@nestjs/common';
+import { AiChatController } from './ai-chat.controller';
+import type { User, Workspace } from '@docmost/db/types/entity.types';
+
+/**
+ * Wiring spec for the #184 run-reconnect / run-stop endpoints
+ * (`POST /ai-chat/run` and `POST /ai-chat/stop`). Both are OWNER-gated via
+ * assertOwnedChat (the requesting user must own the chat) and NOT flag-gated.
+ * Exercised with hand-rolled mocks — no Nest graph, no DB. The controller's
+ * constructor order is (aiChatService, aiChatRunService, aiChatRepo,
+ * aiChatMessageRepo, aiTranscription).
+ */
+describe('AiChatController run endpoints (#184)', () => {
+  const user = { id: 'u1' } as User;
+  const workspace = { id: 'ws1' } as Workspace;
+
+  function makeController(opts: {
+    chat?: unknown; // what aiChatRepo.findById returns (owner-gate)
+    run?: unknown; // getLatestForChat / getRun result
+    activeRun?: unknown; // getActiveForChat result
+    message?: unknown; // aiChatMessageRepo.findById result
+    stopped?: boolean; // requestStop result
+  }) {
+    const aiChatRunService = {
+      getLatestForChat: jest.fn().mockResolvedValue(opts.run),
+      getRun: jest.fn().mockResolvedValue(opts.run),
+      getActiveForChat: jest.fn().mockResolvedValue(opts.activeRun),
+      requestStop: jest.fn().mockResolvedValue(opts.stopped ?? false),
+    };
+    const aiChatRepo = {
+      findById: jest.fn().mockResolvedValue(opts.chat),
+    };
+    const aiChatMessageRepo = {
+      findById: jest.fn().mockResolvedValue(opts.message),
+    };
+    const controller = new AiChatController(
+      {} as never, // aiChatService
+      aiChatRunService as never,
+      aiChatRepo as never,
+      aiChatMessageRepo as never,
+      {} as never, // aiTranscription
+    );
+    return { controller, aiChatRunService, aiChatRepo, aiChatMessageRepo };
+  }
+
+  describe('POST /ai-chat/run (getRun)', () => {
+    it('owner-gates: a chat the user does not own throws ForbiddenException', async () => {
+      const { controller, aiChatRunService } = makeController({
+        chat: { id: 'c1', creatorId: 'someone-else' },
+      });
+      await expect(
+        controller.getRun({ chatId: 'c1' }, user, workspace),
+      ).rejects.toBeInstanceOf(ForbiddenException);
+      // It must NOT reach the run lookup once the owner-gate fails.
+      expect(aiChatRunService.getLatestForChat).not.toHaveBeenCalled();
+    });
+
+    it('returns { run: null, message: null } when the chat has never had a run', async () => {
+      const { controller, aiChatRunService } = makeController({
+        chat: { id: 'c1', creatorId: 'u1' },
+        run: undefined,
+      });
+      const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
+      expect(res).toEqual({ run: null, message: null });
+      expect(aiChatRunService.getLatestForChat).toHaveBeenCalledWith(
+        'c1',
+        'ws1',
+      );
+    });
+
+    it('returns the run and its projected assistant message', async () => {
+      const run = { id: 'run-1', chatId: 'c1', assistantMessageId: 'm1' };
+      const message = { id: 'm1', role: 'assistant' };
+      const { controller, aiChatMessageRepo } = makeController({
+        chat: { id: 'c1', creatorId: 'u1' },
+        run,
+        message,
+      });
+      const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
+      expect(res).toEqual({ run, message });
+      expect(aiChatMessageRepo.findById).toHaveBeenCalledWith('m1', 'ws1');
+    });
+
+    it('returns message: null when the run has no linked assistant message', async () => {
+      const run = { id: 'run-1', chatId: 'c1', assistantMessageId: null };
+      const { controller, aiChatMessageRepo } = makeController({
+        chat: { id: 'c1', creatorId: 'u1' },
+        run,
+      });
+      const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
+      expect(res).toEqual({ run, message: null });
+      expect(aiChatMessageRepo.findById).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('POST /ai-chat/stop (stopRun)', () => {
+    it('throws BadRequestException when neither runId nor chatId is given', async () => {
+      const { controller } = makeController({});
+      await expect(
+        controller.stopRun({}, user, workspace),
+      ).rejects.toBeInstanceOf(BadRequestException);
+    });
+
+    it('stops by runId: owner-gates via the run’s chat, then requests the stop', async () => {
+      const { controller, aiChatRunService, aiChatRepo } = makeController({
+        run: { id: 'run-1', chatId: 'c1' },
+        chat: { id: 'c1', creatorId: 'u1' },
+        stopped: true,
+      });
+      const res = await controller.stopRun({ runId: 'run-1' }, user, workspace);
+      expect(res).toEqual({ stopped: true });
+      expect(aiChatRunService.getRun).toHaveBeenCalledWith('run-1', 'ws1');
+      expect(aiChatRepo.findById).toHaveBeenCalledWith('c1', 'ws1');
+      expect(aiChatRunService.requestStop).toHaveBeenCalledWith('run-1', 'ws1');
+    });
+
+    it('stops by runId: a foreign run’s chat throws ForbiddenException (no stop)', async () => {
+      const { controller, aiChatRunService } = makeController({
+        run: { id: 'run-1', chatId: 'c1' },
+        chat: { id: 'c1', creatorId: 'someone-else' },
+      });
+      await expect(
+        controller.stopRun({ runId: 'run-1' }, user, workspace),
+      ).rejects.toBeInstanceOf(ForbiddenException);
+      expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
+    });
+
+    it('stops by runId: an unknown run reports { stopped: false }', async () => {
+      const { controller, aiChatRunService } = makeController({
+        run: undefined,
+      });
+      const res = await controller.stopRun({ runId: 'gone' }, user, workspace);
+      expect(res).toEqual({ stopped: false });
+      expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
+    });
+
+    it('stops by chatId: owner-gates, resolves the active run, requests the stop', async () => {
+      const { controller, aiChatRunService, aiChatRepo } = makeController({
+        chat: { id: 'c1', creatorId: 'u1' },
+        activeRun: { id: 'run-9' },
+        stopped: true,
+      });
+      const res = await controller.stopRun({ chatId: 'c1' }, user, workspace);
+      expect(res).toEqual({ stopped: true });
+      expect(aiChatRepo.findById).toHaveBeenCalledWith('c1', 'ws1');
+      expect(aiChatRunService.getActiveForChat).toHaveBeenCalledWith(
+        'c1',
+        'ws1',
+      );
+      expect(aiChatRunService.requestStop).toHaveBeenCalledWith('run-9', 'ws1');
+    });
+
+    it('stops by chatId: reports { stopped: false } when no run is active', async () => {
+      const { controller, aiChatRunService } = makeController({
+        chat: { id: 'c1', creatorId: 'u1' },
+        activeRun: undefined,
+      });
+      const res = await controller.stopRun({ chatId: 'c1' }, user, workspace);
+      expect(res).toEqual({ stopped: false });
+      expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
+    });
+  });
+});

apps/server/src/core/ai-chat/ai-chat.service.lifecycle.spec.ts

@@ -1,5 +1,7 @@
 import { Logger } from '@nestjs/common';
-import { AiChatService } from './ai-chat.service';
+import { AiChatService, AiChatRunHooks } from './ai-chat.service';
+import { AiChatRunService } from './ai-chat-run.service';
+import type { User, Workspace } from '@docmost/db/types/entity.types';
 
 /**
  * Lifecycle unit tests for AiChatService.onModuleInit (#183 crash-recovery
@@ -59,3 +61,97 @@ describe('AiChatService.onModuleInit (startup sweep)', () => {
     expect(String(warnSpy.mock.calls[0][0])).toContain('db unavailable');
   });
 });
+
+/**
+ * #184 CRITICAL run-lifecycle safety net (review fix). A transient failure
+ * AFTER a successful beginRun but BEFORE streamText's terminal callbacks own the
+ * lifecycle must STILL settle the run — otherwise the run row is stuck 'running'
+ * forever (sweepRunning only runs at startup) and the partial unique index + the
+ * controller pre-check 409 every future turn in that chat until a restart. Here
+ * we model the very first bare await after beginRun (the user-message insert)
+ * throwing, wiring the run hooks to a REAL AiChatRunService (mock repo) exactly
+ * as the controller does, and assert the run is settled to 'error' and its
+ * in-memory entry dropped (so a follow-up turn would NOT be 409'd).
+ */
+describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
+  const user = { id: 'u1' } as User;
+  const workspace = { id: 'ws1' } as Workspace;
+
+  afterEach(() => jest.restoreAllMocks());
+
+  it('an exception after beginRun settles the run to error and drops the in-memory entry', async () => {
+    jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined);
+
+    // Real run service over a mock repo, so finalizeRun's in-memory bookkeeping
+    // (active.delete) is exercised for real.
+    const runRepo = {
+      insert: jest.fn().mockResolvedValue({ id: 'run-1', status: 'running' }),
+      update: jest.fn().mockResolvedValue({ id: 'run-1' }),
+    };
+    const runService = new AiChatRunService(runRepo as never);
+
+    // The user-message insert (the first bare await after beginRun) throws.
+    const aiChatMessageRepo = {
+      insert: jest.fn().mockRejectedValue(new Error('insert boom')),
+    };
+    const aiChatRepo = {
+      // Existing chat -> chatId stays, no new-chat insert path.
+      findById: jest.fn().mockResolvedValue({ id: 'chat-1', creatorId: 'u1' }),
+    };
+
+    const service = new AiChatService(
+      {} as never, // ai
+      aiChatRepo as never,
+      aiChatMessageRepo as never,
+      {} as never, // aiSettings
+      {} as never, // tools
+      {} as never, // mcpClients
+      {} as never, // aiAgentRoleRepo
+      {} as never, // pageRepo
+      {} as never, // pageAccess
+    );
+
+    const runHooks: AiChatRunHooks = {
+      begin: (chatId) =>
+        runService.beginRun({
+          chatId,
+          workspaceId: workspace.id,
+          userId: user.id,
+          trigger: 'user',
+        }),
+      onSettled: (runId, status, error) =>
+        runService.finalizeRun(runId, workspace.id, status, error),
+    };
+
+    await expect(
+      service.stream({
+        user,
+        workspace,
+        sessionId: 'sess',
+        body: {
+          chatId: 'chat-1',
+          messages: [
+            { id: 'm', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
+          ],
+        },
+        res: {} as never,
+        signal: new AbortController().signal,
+        model: {} as never,
+        role: null,
+        runHooks,
+      }),
+    ).rejects.toThrow('insert boom');
+
+    // The run was begun...
+    expect(runRepo.insert).toHaveBeenCalledTimes(1);
+    // ...then settled to a terminal FAILED status by the safety net...
+    expect(runRepo.update).toHaveBeenCalledTimes(1);
+    expect(runRepo.update).toHaveBeenCalledWith(
+      'run-1',
+      'ws1',
+      expect.objectContaining({ status: 'failed' }),
+    );
+    // ...and the in-memory entry is gone, so a follow-up turn is NOT 409'd.
+    expect(runService.isLocallyActive('run-1')).toBe(false);
+  });
+});

apps/server/src/database/types/entity.types.ts

@@ -56,16 +56,12 @@ export type UpdatableAiChat = Updateable<Omit<AiChats, 'id'>>;
 // full-text search. It is omitted from the public type so it never leaks
 // into HTTP responses or the chat history fed to the language model.
 export type AiChatMessage = Omit<Selectable<AiChatMessages>, 'tsv'>;
-export type InsertableAiChatMessage = Omit<
-  Insertable<AiChatMessages>,
-  'tsv'
->;
+export type InsertableAiChatMessage = Omit<Insertable<AiChatMessages>, 'tsv'>;
 
 // AI Chat Run (#184 phase 1): the agent run as a first-class lifecycle object,
 // detached from the HTTP request / browser window.
 export type AiChatRun = Selectable<AiChatRuns>;
 export type InsertableAiChatRun = Insertable<AiChatRuns>;
-export type UpdatableAiChatRun = Updateable<Omit<AiChatRuns, 'id'>>;
 
 // AI Provider Credentials
 // SECURITY (D9/§8.1): holds encrypted per-workspace provider API keys.
@@ -211,11 +207,14 @@ export type UpdatableFavorite = Updateable<Omit<Favorites, 'id'>>;
 // Page Transclusion
 export type PageTransclusion = Selectable<PageTransclusions>;
 export type InsertablePageTransclusion = Insertable<PageTransclusions>;
-export type UpdatablePageTransclusion = Updateable<Omit<PageTransclusions, 'id'>>;
+export type UpdatablePageTransclusion = Updateable<
+  Omit<PageTransclusions, 'id'>
+>;
 
 // Page Transclusion Reference
 export type PageTransclusionReference = Selectable<PageTransclusionReferences>;
-export type InsertablePageTransclusionReference = Insertable<PageTransclusionReferences>;
+export type InsertablePageTransclusionReference =
+  Insertable<PageTransclusionReferences>;
 export type UpdatablePageTransclusionReference = Updateable<
   Omit<PageTransclusionReferences, 'id'>
 >;
@@ -285,7 +284,9 @@ export type UpdatablePagePermission = Updateable<Omit<_PagePermissions, 'id'>>;
 // Page Verification
 export type PageVerification = Selectable<_PageVerifications>;
 export type InsertablePageVerification = Insertable<_PageVerifications>;
-export type UpdatablePageVerification = Updateable<Omit<_PageVerifications, 'id'>>;
+export type UpdatablePageVerification = Updateable<
+  Omit<_PageVerifications, 'id'>
+>;
 
 // Page Verifier
 export type PageVerifier = Selectable<_PageVerifiers>;