diff --git a/.env.example b/.env.example index e490f7c9..f3acb7ba 100644 --- a/.env.example +++ b/.env.example @@ -288,6 +288,29 @@ MCP_DOCMOST_PASSWORD= # registry is process-local). # AI_CHAT_RESUMABLE_STREAM=false +# --- Run lifecycle tunables (#487) --- +# These govern the universal run machinery (every turn is now a first-class run, +# both modes) and rarely need changing. +# +# How long a server-side SUPERSEDE ("interrupt and send now") waits for the target +# run to settle after issuing Stop before it degrades to a 409 SUPERSEDE_TIMEOUT +# (nothing sent, the composer keeps the user's text). 10s is generous under a +# healthy DB; do NOT raise it to paper over a slow DB — a SUPERSEDE_TIMEOUT is the +# honest signal. Default 10000 (10s). +# AI_CHAT_SUPERSEDE_TIMEOUT_MS=10000 +# +# How often the periodic bidirectional reconcile job runs (heals runs/messages +# left dangling by a crash or a lost terminal write). Default 120000 (2 min). +# AI_CHAT_RECONCILE_INTERVAL_MS=120000 +# +# Wall-clock cap for a SINGLE in-app tool call (a long paginated read, or a content +# write whose collab commit hangs) — the per-call half of the composite abort +# signal every in-app tool is wrapped with (the other half is the turn's Stop). +# The reconcile staleness floor is derived as max(2 x this cap, 15min), so a very +# high value delays stale-run recovery (the server boot-warns above 30min). Default +# 120000 (2 min). +# AI_CHAT_INAPP_TOOL_CALL_CAP_MS=120000 + # --- Anonymous public-share AI assistant --- # Opt-in per workspace (AI settings -> "public share assistant"; off by default). # When enabled, anonymous visitors of a published share can ask an AI about that diff --git a/AGENTS.md b/AGENTS.md index 08756773..dd548743 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -455,7 +455,7 @@ The API server is a Fastify app with a global `/api` prefix (`main.ts` excludes - `core/ai-chat/tools/` — the agent's ~40 read+write tools. Every tool runs under the **calling user's** CASL permissions via a per-user loopback access token (`docmost-client.loader.ts`), so the agent can never exceed what the user could do. Only **reversible** operations are exposed (page history + trash; no permanent delete). Agent edits get an "AI agent" provenance badge in page history (`20260616T130000-agent-provenance` migration). - `core/ai-chat/embedding/` — RAG indexer + a BullMQ consumer on `AI_QUEUE` that embeds pages into `page_embeddings` (vector search), complementing Postgres full-text search. Pages are (re)indexed on edit; `AI_EMBEDDING_TIMEOUT_MS` bounds a hung embeddings endpoint. - `core/ai-chat/external-mcp/` — admins can attach external MCP servers (e.g. Tavily) to give the agent web access. **`ssrf-guard.ts` validates outbound MCP URLs against SSRF** — keep that guard in the path when touching external-MCP connection logic. - - `core/ai-chat/ai-chat-run.service.ts` + `ai_chat_runs` — **detached/autonomous agent runs** (`#184`), behind the per-workspace `settings.ai.autonomousRuns` flag (off by default). When on, a turn becomes a server-side RUN that survives a browser disconnect; only an explicit `POST /ai-chat/stop` ends it, and a client reconnects/live-follows via `POST /ai-chat/run`. **DEPLOY CONSTRAINT — single-instance only in phase 1:** Stop and the AbortController that backs it are process-local, so a Stop only aborts a run executing on the **same** replica that owns it (cross-instance pub/sub stop is phase 2). Do **not** enable `autonomousRuns` on a horizontally-scaled deployment (multiple replicas behind a load balancer, or Docmost cloud `CLOUD=true`) — run a single instance instead. The server logs a startup WARNING when it detects a multi-instance deployment (`CLOUD=true`) so the constraint is visible. The startup sweep settles any run left dangling by a restart. + - `core/ai-chat/ai-chat-run.service.ts` + `ai_chat_runs` — **every agent turn is now a first-class server-side RUN** (`#184`, universalized in `#487`): its lifecycle is tracked in `ai_chat_runs` in **both** modes, and the single-active-run-per-chat concurrency gate is enforced universally (a legacy second tab now gets a clean `409 A_RUN_ALREADY_ACTIVE` instead of a second parallel stream that interleaved history). The per-workspace `settings.ai.autonomousRuns` flag (off by default) **no longer gates whether a turn is a run** — it now controls **only the browser-disconnect semantics**: when ON the run is *detached* (a disconnect leaves it executing server-side; only an explicit `POST /ai-chat/stop` ends it, and a client reconnects/live-follows via `POST /ai-chat/run`); when OFF (legacy) a disconnect ends the turn by stopping its run via the run's stop lever. `#487` also adds a server-side **supersede** CAS ("interrupt and send now") to `POST /ai-chat/stream` (`supersede: { runId }`): it atomically stops the chat's currently-active run and waits for it to settle before the new turn claims the slot, returning `SUPERSEDE_INVALID` / `SUPERSEDE_TARGET_MISMATCH` / `SUPERSEDE_TIMEOUT` on the non-proceed branches. **DEPLOY CONSTRAINT — single-instance only in phase 1:** Stop and the AbortController that backs it are process-local, so a Stop only aborts a run executing on the **same** replica that owns it (cross-instance pub/sub stop is phase 2). Do **not** enable `autonomousRuns` on a horizontally-scaled deployment (multiple replicas behind a load balancer, or Docmost cloud `CLOUD=true`) — run a single instance instead. The server logs a startup WARNING when it detects a multi-instance deployment (`CLOUD=true`) so the constraint is visible. The startup sweep settles any run left dangling by a restart. ### Client structure Vite SPA. Code is organized by feature under `apps/client/src/features/*` (mirrors the server domains: `page`, `space`, `comment`, `ai-chat`, `editor`, …). Conventions: diff --git a/CHANGELOG.md b/CHANGELOG.md index ed4d5a03..30d27413 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -202,6 +202,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 dangling by a restart. Phase 1 is single-instance-only (cross-instance Stop is not yet reliable); the server warns at startup on a horizontally-scaled deployment. (#184) +- **Server-side "interrupt and send now" (supersede) for AI chat.** `POST + /ai-chat/stream` now accepts a `supersede: { runId }` field: when the user sends + a new message while a run is active, the server atomically stops that run and + waits for it to settle before the new turn claims the chat's single run slot, + instead of the send being rejected as concurrent. The compare-and-set surfaces + three codes on its non-proceed branches — `SUPERSEDE_INVALID` (the targeted run + is malformed / belongs to another chat), `SUPERSEDE_TARGET_MISMATCH` (a + different run is now active; carries the current `activeRunId`), and + `SUPERSEDE_TIMEOUT` (the previous run did not stop within the settle window, so + nothing was sent and the composer keeps the text). Tunable via + `AI_CHAT_SUPERSEDE_TIMEOUT_MS` (default 10s). (#487) - **Out-of-band page transfer via an in-RAM blob sandbox (`stash_page`).** A new MCP tool serializes a whole page (its full ProseMirror JSON, with every internal image/file mirrored) into an ephemeral in-RAM blob and returns only @@ -282,6 +293,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- **Every AI-chat turn is now a first-class server-side run, and one run per chat + is enforced in both modes.** The run machinery from `#184` was universalized: a + turn is tracked in `ai_chat_runs` and gated by the single-active-run-per-chat + index regardless of the `settings.ai.autonomousRuns` flag. **Behavior change:** + a second tab (or a double-submit) that starts a turn while one is already active + on the chat is now rejected up front with `409 A_RUN_ALREADY_ACTIVE` (carrying + the `activeRunId`); previously, on the legacy path, it opened a second parallel + stream on the same chat that interleaved history. The `autonomousRuns` flag no + longer controls whether a turn is a run — it now governs **only** the + browser-disconnect semantics (ON = detached/survives a disconnect; OFF = a + disconnect stops the run). (#487) - **Client markdown paste/copy and AI-chat rendering now go through the canonical converter.** Pasting markdown into the editor, "Copy as markdown", the AI title generator, and the AI-chat markdown renderer all now use diff --git a/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts index 4843a979..ade4563c 100644 --- a/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts @@ -43,6 +43,9 @@ function makeRepo(overrides: Record = {}) { workspaceId: v.workspaceId, })), update: jest.fn(async () => ({ id: 'run-1' })), + // #487: terminal finalize now goes through the CONDITIONAL write. Default + // returns a truthy row (the run WAS active -> this call wrote it). + finalizeIfActive: jest.fn(async () => ({ id: 'run-1', status: 'succeeded' })), markStopRequested: jest.fn(async () => ({ id: 'run-1' })), findActiveByChat: jest.fn(async () => undefined), findLatestByChat: jest.fn(async () => undefined), @@ -336,14 +339,12 @@ describe('AiChatRunService run lifecycle', () => { await svc.finalizeRun('run-1', 'ws-1', 'error', 'provider blew up'); expect(svc.isLocallyActive('run-1')).toBe(false); - expect(repo.update).toHaveBeenCalledWith( + // #487: the terminal write is CONDITIONAL (finalizeIfActive); finishedAt is + // stamped inside the repo method, so the service passes just status + error. + expect(repo.finalizeIfActive).toHaveBeenCalledWith( 'run-1', 'ws-1', - expect.objectContaining({ - status: 'failed', - error: 'provider blew up', - finishedAt: expect.any(Date), - }), + expect.objectContaining({ status: 'failed', error: 'provider blew up' }), ); }); @@ -366,8 +367,8 @@ describe('AiChatRunService run lifecycle', () => { // 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( + expect(repo.finalizeIfActive).toHaveBeenCalledTimes(1); + expect(repo.finalizeIfActive).toHaveBeenCalledWith( 'run-1', 'ws-1', expect.objectContaining({ status: 'failed', error: 'first' }), @@ -389,8 +390,8 @@ describe('AiChatRunService run lifecycle', () => { const updateGate = new Promise((res) => { resolveUpdate = res; }); - const update = jest.fn(() => updateGate); - const repo = makeRepo({ update }); + const finalizeIfActive = jest.fn(() => updateGate); + const repo = makeRepo({ finalizeIfActive }); const svc = new AiChatRunService(repo as never, makeEnv() as never); await svc.beginRun({ chatId: 'chat-1', @@ -399,23 +400,23 @@ describe('AiChatRunService run lifecycle', () => { }); // Fire both before the (pending) update resolves. The first synchronously - // claims the entry (active.delete) and awaits update; the second, started in - // the same macrotask, finds the entry already gone and returns at the claim - // WITHOUT ever calling update. + // claims the entry (active.delete) and awaits the write; the second, started + // in the same macrotask, finds the entry already gone and returns at the claim + // WITHOUT ever writing. const p1 = svc.finalizeRun('run-1', 'ws-1', 'completed'); const p2 = svc.finalizeRun('run-1', 'ws-1', 'error', 'safety-net'); // The decisive assertion: exactly one caller reached the terminal UPDATE. - expect(update).toHaveBeenCalledTimes(1); + expect(finalizeIfActive).toHaveBeenCalledTimes(1); // Let the single in-flight update land; both calls resolve cleanly. - resolveUpdate({ id: 'run-1' }); + resolveUpdate({ id: 'run-1', status: 'succeeded' }); await Promise.all([p1, p2]); - expect(update).toHaveBeenCalledTimes(1); + expect(finalizeIfActive).toHaveBeenCalledTimes(1); // The winner is the FIRST caller ('completed' -> 'succeeded'); the late // 'error' settle never wrote, so it could not clobber the real status. - expect(update).toHaveBeenCalledWith( + expect(finalizeIfActive).toHaveBeenCalledWith( 'run-1', 'ws-1', expect.objectContaining({ status: 'succeeded' }), @@ -431,10 +432,10 @@ describe('AiChatRunService run lifecycle', () => { // 409s until a restart. The fix updates FIRST and retries. let calls = 0; const repo = makeRepo({ - update: jest.fn(async () => { + finalizeIfActive: jest.fn(async () => { calls += 1; if (calls === 1) throw new Error('deadlock detected'); - return { id: 'run-1' }; + return { id: 'run-1', status: 'succeeded' }; }), }); jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined); @@ -447,26 +448,29 @@ describe('AiChatRunService run lifecycle', () => { await svc.finalizeRun('run-1', 'ws-1', 'completed'); - // The retry landed the terminal write: the entry is dropped (slot freed) and - // the row carries the real terminal status — NOT stranded at 'running'. + // The retry landed the terminal write: the entry is dropped (slot freed), no + // zombie left, and the row carries the real terminal status. expect(svc.isLocallyActive('run-1')).toBe(false); - expect(repo.update).toHaveBeenCalledTimes(2); - expect(repo.update).toHaveBeenLastCalledWith( + expect(svc.hasZombie('run-1')).toBe(false); + expect(repo.finalizeIfActive).toHaveBeenCalledTimes(2); + expect(repo.finalizeIfActive).toHaveBeenLastCalledWith( 'run-1', 'ws-1', expect.objectContaining({ status: 'succeeded' }), ); }); - it('F6: if the terminal write keeps failing, the entry is RETAINED and a LATER settle completes it (chat not permanently 409d)', async () => { + it('#487 give-up: if the terminal write keeps failing, finalizeRun leaves a ZOMBIE (does NOT restore the entry) and settleZombie re-drives it', async () => { // Worst case: the DB is down for the whole first finalize (all attempts fail). - // The run must NOT be silently lost — the entry stays so a subsequent settle - // (a streamText callback, requestStop -> onAbort, or a future sweep) can retry. + // #487 changes the give-up behaviour: the entry is NOT restored (a restored + // entry is indistinguishable from a live run). Instead a ZOMBIE record holds + // the intended terminal status, and a re-drive (settleZombie — called by the + // reconcile / supersede / opportunistic paths) applies it later. let healthy = false; const repo = makeRepo({ - update: jest.fn(async () => { + finalizeIfActive: jest.fn(async () => { if (!healthy) throw new Error('pool exhausted'); - return { id: 'run-1' }; + return { id: 'run-1', status: 'succeeded' }; }), }); jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined); @@ -480,35 +484,83 @@ describe('AiChatRunService run lifecycle', () => { userId: 'user-1', }); - // First settle: every bounded attempt fails -> entry retained, NOT settled. + // First settle: every bounded attempt fails -> ZOMBIE, entry NOT restored. await svc.finalizeRun('run-1', 'ws-1', 'completed'); - expect(svc.isLocallyActive('run-1')).toBe(true); - // F12: the give-up emits ONE explicit, greppable ERROR (run + chat context) - // so an operator can tell "gave up, run held in memory" from a per-attempt - // blip — distinct from the per-attempt warns. + expect(svc.isLocallyActive('run-1')).toBe(false); // NOT a live entry + expect(svc.hasZombie('run-1')).toBe(true); + expect(svc.zombieRunIds()).toContain('run-1'); + // The give-up emits ONE explicit, greppable ERROR mentioning the zombie. const gaveUp = errorSpy.mock.calls.some( (c) => /NON-TERMINAL/.test(String(c[0])) && + /ZOMBIE/.test(String(c[0])) && /run-1/.test(String(c[0])) && /chat-1/.test(String(c[0])), ); expect(gaveUp).toBe(true); + // The settle notifier resolved as terminalWriteFailed (a subscriber learns the + // slot still needs the intended status applied). + const outcome = await svc.peekSettled('run-1'); + expect(outcome).toEqual({ + status: 'succeeded', + error: null, + terminalWriteFailed: true, + }); - // The DB recovers; a later settle now succeeds and frees the slot. + // The DB recovers; a re-drive settles the zombie via the conditional UPDATE. healthy = true; - await svc.finalizeRun('run-1', 'ws-1', 'completed'); - expect(svc.isLocallyActive('run-1')).toBe(false); - expect(repo.update).toHaveBeenLastCalledWith( + const redriven = await svc.settleZombie('run-1'); + expect(redriven).toBe(true); + expect(svc.hasZombie('run-1')).toBe(false); + expect(repo.finalizeIfActive).toHaveBeenLastCalledWith( 'run-1', 'ws-1', expect.objectContaining({ status: 'succeeded' }), ); - // And it is now idempotent: a further settle no-ops (terminal row already - // written), so a double-settle can never clobber the real status. - const callsBefore = repo.update.mock.calls.length; + // A later finalizeRun is idempotent (row already terminal): it no-ops at the + // once-gate, never re-writing. + const callsBefore = repo.finalizeIfActive.mock.calls.length; await svc.finalizeRun('run-1', 'ws-1', 'error', 'late'); - expect(repo.update).toHaveBeenCalledTimes(callsBefore); + expect(repo.finalizeIfActive).toHaveBeenCalledTimes(callsBefore); + }); + + it('#487 double-settle collapses to a benign no-op (conditional write; notifier resolves once)', async () => { + // A second concurrent settle is stopped at the synchronous active.delete + // claim, so the terminal write runs exactly once and the notifier resolves + // exactly once with the FIRST settler's outcome. + const repo = makeRepo(); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + await svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'u1' }); + + await svc.finalizeRun('run-1', 'ws-1', 'aborted'); + await svc.finalizeRun('run-1', 'ws-1', 'error', 'late'); // no-op + + expect(repo.finalizeIfActive).toHaveBeenCalledTimes(1); + const outcome = await svc.peekSettled('run-1'); + // peekSettled after resolve+delete falls through (notifier dropped, no zombie) + // -> undefined; the FIRST settler already resolved any earlier subscriber. + expect(outcome).toBeUndefined(); + }); + + it('#487 late settledPromise subscriber gets the resolved outcome', async () => { + const repo = makeRepo(); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + await svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'u1' }); + + // Subscribe BEFORE settle: hold the promise reference (as supersede does). + const early = svc.peekSettled('run-1'); + expect(early).toBeDefined(); + + await svc.finalizeRun('run-1', 'ws-1', 'completed'); + + // The reference grabbed before settle resolves with the written outcome, even + // though the notifier was dropped from the map on resolve (bounded). + await expect(early).resolves.toEqual({ + status: 'succeeded', + error: null, + terminalWriteFailed: false, + }); }); it('recordStep / linkAssistantMessage are best-effort: a repo failure is swallowed', async () => { @@ -525,3 +577,197 @@ describe('AiChatRunService run lifecycle', () => { ).resolves.toBeUndefined(); }); }); + +describe('#487 AiChatRunService.supersede (CAS)', () => { + const chat = 'chat-1'; + const ws = 'ws-1'; + + it('degrade: no active run on the chat -> caller sends a normal turn', async () => { + const repo = makeRepo({ + findById: jest.fn(async () => undefined), + findActiveByChat: jest.fn(async () => undefined), + }); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ kind: 'degrade' }); + }); + + it('invalid: the target run belongs to a DIFFERENT chat -> 400', async () => { + const repo = makeRepo({ + findById: jest.fn(async () => ({ + id: 'run-x', + chatId: 'other-chat', + workspaceId: ws, + })), + }); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ kind: 'invalid' }); + }); + + it('mismatch: a DIFFERENT run is active than the one targeted -> current runId', async () => { + const repo = makeRepo({ + findById: jest.fn(async () => ({ id: 'run-x', chatId: chat, workspaceId: ws })), + findActiveByChat: jest.fn(async () => ({ + id: 'run-live', + chatId: chat, + workspaceId: ws, + status: 'running', + })), + }); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ + kind: 'mismatch', + activeRunId: 'run-live', + }); + }); + + it('ready: the target IS active -> stop it, await its (fast) settle, free the slot', async () => { + // Simulate a live long TOOL (NOT a slow UPDATE): the run stays active until an + // explicit Stop unwinds it; commit-1's race makes that settle land quickly. + // The abort listener stands in for streamText's onAbort -> finalizeRun. + const repo = makeRepo({ + findById: jest.fn(async () => ({ + id: 'run-1', + chatId: chat, + workspaceId: ws, + status: 'aborted', + error: null, + })), + findActiveByChat: jest.fn(async () => ({ + id: 'run-1', + chatId: chat, + workspaceId: ws, + status: 'running', + })), + }); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + const handle = await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' }); + handle.signal.addEventListener('abort', () => { + void svc.finalizeRun('run-1', ws, 'aborted'); + }); + + // supersede: getRun -> getActiveByChat(==target) -> requestStop -> the abort + // listener settles the run -> awaitSettled resolves -> ready. + expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({ + kind: 'ready', + }); + expect(handle.signal.aborted).toBe(true); // Stop reached the run + }); + + it('timeout: the target never settles within W -> 409 SUPERSEDE_TIMEOUT (nothing persisted)', async () => { + const repo = makeRepo({ + findById: jest.fn(async () => ({ id: 'run-1', chatId: chat, workspaceId: ws })), + findActiveByChat: jest.fn(async () => ({ + id: 'run-1', + chatId: chat, + workspaceId: ws, + status: 'running', + })), + }); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' }); + // Do NOT settle the run: a tiny W elapses -> timeout. + const result = await svc.supersede(chat, 'run-1', ws, 30); + expect(result).toEqual({ kind: 'timeout' }); + }); + + it('ready then a DUPLICATE supersede POST degrades (the run is already gone)', async () => { + let active: unknown = { + id: 'run-1', + chatId: chat, + workspaceId: ws, + status: 'running', + }; + const repo = makeRepo({ + findById: jest.fn(async () => ({ + id: 'run-1', + chatId: chat, + workspaceId: ws, + status: 'aborted', + error: null, + })), + findActiveByChat: jest.fn(async () => active), + finalizeIfActive: jest.fn(async () => { + active = undefined; // settling frees the active slot + return { id: 'run-1', status: 'aborted' }; + }), + }); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + const handle = await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' }); + handle.signal.addEventListener('abort', () => { + void svc.finalizeRun('run-1', ws, 'aborted'); + }); + + expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({ + kind: 'ready', + }); + // The duplicate POST for the same target now finds no active run -> degrade. + expect(await svc.supersede(chat, 'run-1', ws)).toEqual({ kind: 'degrade' }); + }); + + it('reconcileStaleRuns: aborts a stale run with NO entry/zombie; NEVER touches a live entry', async () => { + const finalizeIfActive = jest.fn(async () => ({ id: 'x', status: 'aborted' })); + const repo = makeRepo({ + insert: jest.fn(async (v: any) => ({ + id: 'live-1', + status: 'running', + chatId: v.chatId, + workspaceId: v.workspaceId, + })), + finalizeIfActive, + findStaleActive: jest.fn(async () => [ + { id: 'orphan-1', workspaceId: ws, chatId: 'c-orphan' }, + { id: 'live-1', workspaceId: ws, chatId: 'c-live' }, + ]), + }); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + // A LIVE run this replica owns (in the `active` map). + await svc.beginRun({ chatId: 'c-live', workspaceId: ws, userId: 'u1' }); + expect(svc.isLocallyActive('live-1')).toBe(true); + + const aborted = await svc.reconcileStaleRuns(15 * 60 * 1000); + expect(aborted).toBe(1); + // The orphan (no entry) was aborted; the live entry was NEVER passed to the DB. + expect(finalizeIfActive).toHaveBeenCalledTimes(1); + expect(finalizeIfActive).toHaveBeenCalledWith( + 'orphan-1', + ws, + expect.objectContaining({ status: 'aborted' }), + ); + expect(svc.isLocallyActive('live-1')).toBe(true); + }); + + it('gave-up zombie: supersede applies the intended status (settleZombie) then is ready', async () => { + let healthy = false; + let active: unknown = { + id: 'run-1', + chatId: chat, + workspaceId: ws, + status: 'running', + }; + const repo = makeRepo({ + findById: jest.fn(async () => ({ id: 'run-1', chatId: chat, workspaceId: ws })), + findActiveByChat: jest.fn(async () => active), + finalizeIfActive: jest.fn(async () => { + if (!healthy) throw new Error('db down'); + active = undefined; + return { id: 'run-1', status: 'aborted' }; + }), + }); + jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined); + jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined); + const svc = new AiChatRunService(repo as never, makeEnv() as never); + await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' }); + + // The run's terminal write gives up -> zombie (row still 'running'). + await svc.finalizeRun('run-1', ws, 'aborted'); + expect(svc.hasZombie('run-1')).toBe(true); + + // The DB recovers; supersede awaits the (already-resolved, terminalWriteFailed) + // settle, then settleZombie applies the intended status -> ready. + healthy = true; + expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({ + kind: 'ready', + }); + expect(svc.hasZombie('run-1')).toBe(false); + }); +}); diff --git a/apps/server/src/core/ai-chat/ai-chat-run.service.ts b/apps/server/src/core/ai-chat/ai-chat-run.service.ts index b2e0d393..1530c3b5 100644 --- a/apps/server/src/core/ai-chat/ai-chat-run.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat-run.service.ts @@ -34,6 +34,88 @@ export class RunAlreadyActiveError extends Error { export type TurnTerminalStatus = 'completed' | 'error' | 'aborted'; export type RunTerminalStatus = 'succeeded' | 'failed' | 'aborted'; +/** The terminal run statuses — the row is done once it reads one of these. */ +export const RUN_TERMINAL_STATUSES: readonly RunTerminalStatus[] = [ + 'succeeded', + 'failed', + 'aborted', +]; + +/** Whether a persisted run status is terminal (settled). */ +export function isRunTerminal(status: string | null | undefined): boolean { + return ( + status === 'succeeded' || status === 'failed' || status === 'aborted' + ); +} + +/** + * #487: the outcome a run's {@link AiChatRunService.finalizeRun} settled with. + * `terminalWriteFailed` = the terminal write GAVE UP after the bounded retry, so + * the row is still non-terminal ('running') and a ZOMBIE record holds the + * `intended` status for a later re-drive (reconcile / supersede / boot sweep). A + * subscriber (supersede, #487 commit 3) uses this to decide whether the slot is + * genuinely free or must first have the intended status applied. + */ +export interface RunSettleOutcome { + status: RunTerminalStatus; + error: string | null; + terminalWriteFailed: boolean; +} + +/** + * #487: how long a supersede waits for the target run to settle after Stop before + * it degrades to `SUPERSEDE_TIMEOUT`. W=10s is generous under a HEALTHY DB: commit + * 1's race-on-abort makes an in-app tool abort->settle in ms/hundreds of ms, so a + * live run releases its slot well within the window. Under a DB brownout the + * timeout is normal (the write cannot land); W must NOT be raised to paper + * over a slow DB — a SUPERSEDE_TIMEOUT is the honest signal (nothing persisted, + * the composer keeps the user's text). Env-tunable for ops, default 10s. + */ +export const SUPERSEDE_SETTLE_TIMEOUT_MS = (() => { + const raw = Number(process.env.AI_CHAT_SUPERSEDE_TIMEOUT_MS); + return Number.isFinite(raw) && raw > 0 ? raw : 10_000; +})(); + +/** + * #487: the result of the supersede CAS ({@link AiChatRunService.supersede}). + * - `degrade` : no active run on the chat (it ended between click and POST) — + * the caller sends a NORMAL turn (NOT a mismatch); + * - `invalid` : the target runId belongs to a DIFFERENT chat (malformed CAS 400); + * - `mismatch` : a DIFFERENT run is active than the one the client targeted — + * 409 SUPERSEDE_TARGET_MISMATCH carrying the current `activeRunId` + * (the client does NOT auto-retry); + * - `timeout` : the target did not settle within W — 409 SUPERSEDE_TIMEOUT, + * nothing persisted; + * - `ready` : the target was stopped AND settled (or its zombie's intended was + * applied) — the slot is free; the caller may beginRun the new run. + */ +export type SupersedeResult = + | { kind: 'degrade' } + | { kind: 'invalid' } + | { kind: 'mismatch'; activeRunId: string } + | { kind: 'timeout' } + | { kind: 'ready' }; + +/** A one-shot settle notifier (#487): `resolve` is called EXACTLY ONCE. */ +interface Deferred { + promise: Promise; + resolve: (value: T) => void; +} + +/** + * #487: a run whose terminal write GAVE UP (every bounded attempt failed). The + * row is stranded non-terminal ('running'); this record is the ONLY thing that + * distinguishes it from a live run, and carries the `intended` terminal status so + * a re-drive can apply it via the conditional UPDATE. Process-local (phase-1 + * single-process assumption): a restart drops it, and the boot sweep then writes + * 'aborted' over the intended — a documented loss (see finalizeRun). + */ +interface ZombieRun { + workspaceId: string; + chatId: string; + intended: { status: RunTerminalStatus; error: string | null }; +} + export function mapTurnStatusToRun( status: TurnTerminalStatus, ): RunTerminalStatus { @@ -101,6 +183,22 @@ export class AiChatRunService implements OnModuleInit { // uptime — negligible in phase 1's single process. private readonly settled = new Set(); + // #487 runId -> one-shot settle notifier. Kept in a SEPARATE map from `active` + // ON PURPOSE: it must OUTLIVE the `active.delete` claim inside finalizeRun (the + // claim frees the slot the instant finalize starts), so a subscriber can still + // await the outcome after the entry is gone. Created in beginRun, resolved + // EXACTLY ONCE in finalizeRun, then removed (bounded). Absence => this replica + // has no live notifier: a subscriber falls back to the zombie map, then to the + // row (see peekSettled). Process-local (phase-1 single-process assumption). + private readonly settledPromises = new Map>(); + + // #487 runId -> ZOMBIE record: a run whose terminal write gave up (row stranded + // non-terminal). BOUNDED — an entry is added only on give-up and removed on a + // successful re-drive (settleZombie) or when the row is found already terminal; + // a process restart clears it (and the boot sweep settles the stranded row). + // Process-local (phase-1 single-process assumption). + private readonly zombies = new Map(); + // Bounded retry for the terminal write (F6): a single PK UPDATE can fail // transiently under many fire-and-forget writes (pool exhaustion, deadlock, a // brief connection blip). Riding out that blip in-place matters because the @@ -224,6 +322,10 @@ export class AiChatRunService implements OnModuleInit { chatId: args.chatId, workspaceId: args.workspaceId, }); + // #487: arm the one-shot settle notifier BEFORE returning, so a subscriber + // that races in immediately after begin always finds a promise to await. It + // is resolved exactly once when the run settles (or gives up). + this.settledPromises.set(run.id, this.makeDeferred()); return { runId: run.id, signal: controller.signal }; } @@ -263,47 +365,43 @@ export class AiChatRunService implements OnModuleInit { } /** - * Finalize a run to its terminal status (succeeded / failed / aborted), - * stamping finishedAt + any error. Best-effort, but ROBUST against a transient - * terminal-write failure (F6) AND atomically safe against a concurrent settle. + * Finalize a run to its terminal status (succeeded / failed / aborted) via a + * CONDITIONAL UPDATE, stamping finishedAt + any error. Atomically safe against a + * concurrent settle AND robust against a transient terminal-write failure. * * ATOMIC ONCE-CLAIM (the gate must close in ONE synchronous tick): two * finalizeRun calls for the SAME run can race — the documented real path is * AiChatService.stream's safety-net catch settling the turn to 'error' while a * streamText terminal callback (onFinish/onAbort/onError) ALSO settles it. The - * `settled.has` check alone is NOT a gate: it is read BEFORE the awaited UPDATE, - * so two callers can both see `false` and both write the row (last-write-wins - * clobbers the real terminal status, and the bounded retry only widens that - * window). The claim therefore happens via `active.delete`, a SYNCHRONOUS - * check-and-clear with NO await between the gate and the entry removal: the - * second concurrent caller finds the entry already gone and returns in the same - * tick, before any UPDATE. The transition "nobody is finalizing" -> "I am - * finalizing" is thus a single atomic step. + * claim happens via `active.delete`, a SYNCHRONOUS check-and-clear with NO await + * between the gate and the entry removal: the second concurrent caller finds the + * entry already gone and returns in the same tick, before any UPDATE. * - * ORDER MATTERS (F6): once we own the claim, the terminal UPDATE happens FIRST; - * only once it SUCCEEDS do we record the run as settled. If the UPDATE fails on - * every bounded attempt we RESTORE the in-memory entry, leave the run UNsettled, - * and emit an ERROR signal that the row is left non-terminal 'running' (which - * would 409 every future turn in the chat until recovery). An in-process retry - * by a LATER settle is only POSSIBLE, never guaranteed: it needs (a) the entry - * to have been restored at the give-up path AND (b) a fresh settler to arrive - * AFTER that restore. A concurrent settler that arrives DURING the retry window - * — while the entry is deleted for backoff and not yet restored — is consumed at - * the synchronous `active.delete` claim (it finds nothing to delete and returns - * a no-op), so it does NOT become an in-process retrier. The NO-streamText path - * (the turn threw before streamText was wired, so ONLY the safety-net ever - * settles) likewise has no second in-process settler at all. The UNCONDITIONAL - * backstop in every case is the boot sweep on the next restart (phase 1 has no - * periodic in-process sweep); the retained entry is bounded (cleared on restart) - * and harmless meanwhile. + * ALL TERMINAL WRITES ARE CONDITIONAL (#487): `finalizeIfActive` only flips a + * row still in pending|running (mirror of the assistant message's + * `onlyIfStreaming`). So even a settle that DID reach the UPDATE (e.g. a + * reconcile stamp racing an owner finalize) can never clobber a terminal status + * — the loser matches nothing and is a benign no-op. `active.delete` is the + * fast, in-process gate; the conditional WHERE is the authoritative one. * - * IDEMPOTENT on SUCCESS (#184 review): the terminal write happens AT MOST ONCE - * per run. After a successful write the once-gate keys off {@link settled} (the - * terminal row already written) so a settle arriving AFTER the entry was already - * dropped-and-settled returns early; a settle racing the in-flight write is - * stopped earlier still, by the `active.delete` claim. Either way a genuine - * double-settle collapses to a single write and a late settle can never clobber - * the real terminal status or double-write the row. + * ZOMBIE ON GIVE-UP (#487): if every bounded attempt THROWS (the DB is down for + * the whole finalize), we do NOT restore the entry. The row is stranded + * non-terminal ('running'); we record a ZOMBIE `{ terminalWriteFailed, intended + * }` (the ONLY thing distinguishing this dead run from a live one) and resolve + * the settle notifier with `terminalWriteFailed: true`. A restore would make the + * zombie indistinguishable from a live run to every reader; instead a re-drive + * (settleZombie, called by the periodic reconcile / supersede / opportunistic + * paths) applies the intended status later via the same conditional UPDATE. + * + * DOCUMENTED LOSS (#487, single-process phase 1): if the process RESTARTS before + * a zombie is re-driven, the in-memory zombie map is gone and the boot sweep + * (unconditional) writes 'aborted' over the ACTUAL intended status. This is + * unavoidable while the run lifecycle is single-process — there is no durable + * record of `intended`; a cross-process durable intent is deferred to phase 2. + * + * IDEMPOTENT: the settle notifier resolves EXACTLY ONCE; a second settle is + * stopped at `settled.has` or the `active.delete` claim, so a double-settle + * collapses to a single write and can never double-resolve or clobber the row. */ async finalizeRun( runId: string, @@ -314,13 +412,17 @@ export class AiChatRunService implements OnModuleInit { // ---- Atomic once-claim (synchronous; NO await before the gate closes) ---- // Already terminally written -> idempotent no-op. if (this.settled.has(runId)) return; - // Capture the entry BEFORE the delete so a total-failure path can restore it. + // Capture the entry BEFORE the delete for the give-up log context. const entry = this.active.get(runId); // SYNCHRONOUS check-and-clear: the FIRST caller deletes (claims) the entry; // any concurrent SECOND caller finds nothing to delete and returns HERE, in // the same tick, before any await — so it can never reach the UPDATE. if (!this.active.delete(runId)) return; + const status = mapTurnStatusToRun(turnStatus); + const err = error ?? null; + const chatId = entry?.chatId ?? 'unknown'; + let lastError: unknown; for ( let attempt = 1; @@ -328,47 +430,294 @@ export class AiChatRunService implements OnModuleInit { attempt++ ) { try { - await this.runRepo.update(runId, workspaceId, { - status: mapTurnStatusToRun(turnStatus), - finishedAt: new Date(), - error: error ?? null, + const row = await this.runRepo.finalizeIfActive(runId, workspaceId, { + status, + error: err, }); - // Terminal write landed: arm the once-gate. The entry is already gone - // (claimed above); we do NOT restore it. The slot is now free. + // No throw => the row is now terminal (we wrote it, or it was ALREADY + // terminal — another writer won the conditional UPDATE, a benign no-op). this.settled.add(runId); + this.zombies.delete(runId); + // Resolve with the persisted outcome: our status when WE wrote it, else + // the row's real terminal status (re-read on the already-terminal path so + // a subscriber never sees a status we did not actually persist). + const outcome: RunSettleOutcome = row + ? { status, error: err, terminalWriteFailed: false } + : await this.readTerminalOutcome(runId, workspaceId, status, err); + this.resolveSettled(runId, outcome); return; - } catch (err) { - lastError = err; + } catch (err2) { + lastError = err2; this.logger.warn( `Failed to finalize run ${runId} (attempt ${attempt}/${ AiChatRunService.FINALIZE_MAX_ATTEMPTS - }): ${err instanceof Error ? err.message : 'unknown error'}`, + }): ${err2 instanceof Error ? err2.message : 'unknown error'}`, ); if (attempt < AiChatRunService.FINALIZE_MAX_ATTEMPTS) { await this.delay(AiChatRunService.FINALIZE_RETRY_BASE_MS * attempt); } } } - // Every attempt failed: this is a give-up, materially worse than a per-attempt - // blip — the row is left NON-TERMINAL ('running'), so emit ONE explicit, - // greppable ERROR so an operator can tell "survived a blip" from "gave up, run - // held in memory until recovery" (the last warn alone says only "attempt 3/3"). + // Every attempt threw: GIVE UP. The row is stranded non-terminal ('running'). + // Do NOT restore the entry (a restored entry is indistinguishable from a live + // run); leave a ZOMBIE record instead, and resolve the notifier as + // terminalWriteFailed so a subscriber knows the slot still needs the intended + // status applied. One explicit, greppable ERROR so an operator can tell a + // give-up from a per-attempt blip. this.logger.error( - `Run ${runId} (chat ${entry?.chatId ?? 'unknown'}) left NON-TERMINAL ` + - `('running'): terminal write failed after ${ - AiChatRunService.FINALIZE_MAX_ATTEMPTS - } attempts; entry retained in memory, recovery deferred to next settle / ` + - `boot sweep`, + `Run ${runId} (chat ${chatId}) left NON-TERMINAL ('running'): terminal ` + + `write failed after ${AiChatRunService.FINALIZE_MAX_ATTEMPTS} attempts; ` + + `ZOMBIE recorded (intended '${status}'), recovery deferred to reconcile / ` + + `supersede / boot sweep`, lastError, ); - // RESTORE the claimed entry (and leave the run UNsettled) so a LATER settle - // that arrives AFTER this restore MAY retry the terminal write — but that - // in-process retry is NOT guaranteed (a concurrent settler caught in the retry - // window above is consumed at the `active.delete` claim, and the no-streamText - // path has no second settler at all). The UNCONDITIONAL backstop in every case - // is the boot sweep on the next restart; the restored entry is bounded and - // cleared on restart. - if (entry) this.active.set(runId, entry); + this.zombies.set(runId, { + workspaceId, + chatId, + intended: { status, error: err }, + }); + this.resolveSettled(runId, { status, error: err, terminalWriteFailed: true }); + } + + /** + * #487: re-drive a zombie run's intended terminal write (the conditional + * UPDATE). Called by the periodic reconcile (commit 4), an opportunistic + * single-chat reconcile, and supersede (commit 3). On success — the row is now + * terminal (written OR found already terminal) — the zombie is cleared and the + * once-gate armed; on another failure the zombie is kept for a later retry. + * Returns true when the row is now terminal. Best-effort; never throws. + */ + async settleZombie(runId: string): Promise { + const z = this.zombies.get(runId); + if (!z) return false; + try { + await this.runRepo.finalizeIfActive(runId, z.workspaceId, { + status: z.intended.status, + error: z.intended.error, + }); + this.zombies.delete(runId); + this.settled.add(runId); + return true; + } catch (err) { + this.logger.warn( + `Re-drive of zombie run ${runId} (chat ${z.chatId}) failed; will retry ` + + `later: ${err instanceof Error ? err.message : 'unknown error'}`, + ); + return false; + } + } + + /** + * #487 reconcile clause (c): abort runs the DB still shows active (pending| + * running) but that this replica does NOT own — NO live entry AND NO zombie — + * and that have been UNTOUCHED past `staleMs` (from last-progress `updated_at`, + * NOT startedAt, so a legit long marathon is never a candidate). "No entry" is + * the PRIMARY gate: a live entry (an actively-executing run on this replica) is + * NEVER aborted, whatever its age. Returns the number aborted. Best-effort — + * never throws (a periodic-job failure must not crash the process). + */ + async reconcileStaleRuns(staleMs: number): Promise { + let candidates: Array<{ id: string; workspaceId: string; chatId: string }>; + try { + candidates = await this.runRepo.findStaleActive(staleMs); + } catch (err) { + this.logger.warn( + `Reconcile (stale runs) query failed: ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + return 0; + } + let aborted = 0; + for (const c of candidates) { + // PRIMARY gate: never touch a live entry, and never race a zombie we are + // already re-driving (settleZombie owns those). + if (this.active.has(c.id) || this.zombies.has(c.id)) continue; + try { + const row = await this.runRepo.finalizeIfActive(c.id, c.workspaceId, { + status: 'aborted', + error: 'Run aborted by reconcile: no live runner (stale).', + }); + if (row) { + aborted += 1; + this.settled.add(c.id); + } + } catch (err) { + this.logger.warn( + `Reconcile abort of stale run ${c.id} failed: ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + } + } + return aborted; + } + + /** + * #487: the run's settle outcome as seen by THIS replica, or undefined when it + * has no record (the caller then reads the row — the DB is the source of truth). + * A LIVE deferred (still settling, or resolved-but-not-yet-consumed) wins; a + * ZOMBIE synthesizes the give-up outcome. A subscriber (supersede) races this + * against a timeout. + */ + peekSettled(runId: string): Promise | undefined { + const d = this.settledPromises.get(runId); + if (d) return d.promise; + const z = this.zombies.get(runId); + if (z) { + return Promise.resolve({ + status: z.intended.status, + error: z.intended.error, + terminalWriteFailed: true, + }); + } + return undefined; + } + + /** + * #487: await a run's settle outcome, bounded by `timeoutMs`. Returns the + * outcome on settle, or undefined on TIMEOUT (or when this replica has no record + * of the run and its row is not terminal). Uses the LIVE settle notifier / the + * zombie synth when present; else reads the row (the DB is the source of truth + * once the in-memory record is gone). The subscriber (supersede) grabs this + * right after Stop; commit 1's race makes the settle land in ms on a healthy DB. + */ + async awaitSettled( + runId: string, + workspaceId: string, + timeoutMs: number, + ): Promise { + const pending = this.peekSettled(runId); + if (pending) { + let timer: ReturnType | undefined; + const timeout = new Promise((resolve) => { + timer = setTimeout(() => resolve(undefined), timeoutMs); + timer.unref?.(); + }); + try { + return await Promise.race([pending, timeout]); + } finally { + if (timer) clearTimeout(timer); + } + } + // No live notifier and no zombie: read the row (already settled-and-written, + // or unknown here). A terminal row is an outcome; anything else -> undefined. + const row = await this.runRepo.findById(runId, workspaceId); + if (row && isRunTerminal(row.status)) { + return { + status: row.status as RunTerminalStatus, + error: row.error ?? null, + terminalWriteFailed: false, + }; + } + return undefined; + } + + /** + * #487: the SERVER supersede CAS for `POST /stream { supersede: { runId: X } }`. + * Atomically transitions "X is the chat's active run" -> "X is stopped, settled, + * slot free" so the caller can start a replacement run. See {@link + * SupersedeResult} for the branch semantics. + * + * On a `ready` result the caller MUST still go through the normal beginRun gate + * (the partial unique index) — between the slot freeing here and beginRun a + * neighbouring tab's ordinary POST can win the slot (documented SLOT-THEFT: the + * loser then gets a MISMATCH carrying the NEW runId). There is also NO side- + * effect quiescence: an in-flight write of the stopped run may still land AFTER + * the new run starts (commit 1 stops the NEXT call, not one already committing), + * so the caller adds a prompt note to the new run. + */ + async supersede( + chatId: string, + targetRunId: string, + workspaceId: string, + timeoutMs: number = SUPERSEDE_SETTLE_TIMEOUT_MS, + ): Promise { + // Validate the target belongs to THIS chat (a CAS targeting another chat's run + // is malformed -> 400). A missing row is NOT invalid: the run may have ended + // and been pruned; the active-run check below decides degrade vs mismatch. + const target = await this.getRun(targetRunId, workspaceId); + if (target && target.chatId !== chatId) return { kind: 'invalid' }; + + const active = await this.getActiveForChat(chatId, workspaceId); + // No active run: it ended between the client's click and this POST — this is a + // DEGRADE to a normal send, NOT a mismatch (the user's intent still holds). + if (!active) return { kind: 'degrade' }; + // A DIFFERENT run is active than the one the client saw -> mismatch. The + // client does not auto-retry; it surfaces the new runId. + if (active.id !== targetRunId) { + return { kind: 'mismatch', activeRunId: active.id }; + } + + // The target IS active: stop it, then await its settle within W. + await this.requestStop(targetRunId, workspaceId); + const outcome = await this.awaitSettled(targetRunId, workspaceId, timeoutMs); + if (!outcome) return { kind: 'timeout' }; + // Gave up (terminal write failed): apply the intended status via the + // conditional UPDATE so the slot actually frees. If that ALSO fails, the row + // is still stranded -> treat as a timeout (nothing persisted for the new run). + if (outcome.terminalWriteFailed) { + const settled = await this.settleZombie(targetRunId); + if (!settled) return { kind: 'timeout' }; + } + return { kind: 'ready' }; + } + + /** #487 test/diagnostic seam: whether a give-up zombie is held for this run. */ + hasZombie(runId: string): boolean { + return this.zombies.has(runId); + } + + /** #487: every zombie runId held on this replica (reconcile clause a, commit 4). */ + zombieRunIds(): string[] { + return [...this.zombies.keys()]; + } + + /** #487: create a one-shot deferred (resolve captured for a later single call). */ + private makeDeferred(): Deferred { + let resolve!: (value: T) => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { promise, resolve }; + } + + /** #487: resolve a run's settle notifier EXACTLY ONCE, then drop it (bounded). + * A subscriber that already grabbed the promise still resolves; a later one + * falls back to the zombie map / the row (see peekSettled). */ + private resolveSettled(runId: string, outcome: RunSettleOutcome): void { + const d = this.settledPromises.get(runId); + if (!d) return; + this.settledPromises.delete(runId); + d.resolve(outcome); + } + + /** #487: read the persisted terminal outcome when the conditional finalize was a + * no-op (the row was already terminal). Falls back to the intended status when + * the read fails or the row is unexpectedly missing/non-terminal. */ + private async readTerminalOutcome( + runId: string, + workspaceId: string, + fallbackStatus: RunTerminalStatus, + fallbackError: string | null, + ): Promise { + try { + const row = await this.runRepo.findById(runId, workspaceId); + if (row && isRunTerminal(row.status)) { + return { + status: row.status as RunTerminalStatus, + error: row.error ?? null, + terminalWriteFailed: false, + }; + } + } catch { + // Fall through to the intended status — best-effort only. + } + return { + status: fallbackStatus, + error: fallbackError, + terminalWriteFailed: false, + }; } /** Small async backoff between terminal-write retries (F6). Isolated so it is diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.export.spec.ts b/apps/server/src/core/ai-chat/ai-chat.controller.export.spec.ts index bfbc64c0..6afefc61 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.export.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.export.spec.ts @@ -115,7 +115,7 @@ describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', ( // Drive the SAME applyFinalize the service calls (no duplicated logic). async function dispatchFinalize( - repo: { insert: jest.Mock; update: jest.Mock }, + repo: { insert: jest.Mock; finalizeOwner: jest.Mock }, assistantId: string | undefined, flushed: AssistantFlush, ): Promise { @@ -135,21 +135,22 @@ describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', ( expect(planFinalizeAssistant(undefined)).toEqual({ kind: 'insert' }); }); - it('(a) upfront insert succeeded -> finalize UPDATEs the row by id', async () => { - const repo = { insert: jest.fn(), update: jest.fn() }; + it('(a) upfront insert succeeded -> finalize CONDITIONALLY updates the row by id (#487 owner-write)', async () => { + const repo = { insert: jest.fn(), finalizeOwner: jest.fn() }; const flushed = flushAssistant([], 'final answer', 'completed', { finishReason: 'stop', }); await dispatchFinalize(repo, 'a1', flushed); - expect(repo.update).toHaveBeenCalledWith('a1', workspaceId, flushed); + // #487: the owner write is the CONDITIONAL finalizeOwner, not a raw update. + expect(repo.finalizeOwner).toHaveBeenCalledWith('a1', workspaceId, flushed); expect(repo.insert).not.toHaveBeenCalled(); }); it('(b) upfront insert failed -> finalize INSERTs the terminal payload', async () => { - const repo = { insert: jest.fn(), update: jest.fn() }; + const repo = { insert: jest.fn(), finalizeOwner: jest.fn() }; const flushed = flushAssistant([], 'partial', 'error', { error: 'boom' }); await dispatchFinalize(repo, undefined, flushed); - expect(repo.update).not.toHaveBeenCalled(); + expect(repo.finalizeOwner).not.toHaveBeenCalled(); expect(repo.insert).toHaveBeenCalledTimes(1); const arg = repo.insert.mock.calls[0][0]; // The fallback insert carries the terminal content/status/metadata. diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.supersede.spec.ts b/apps/server/src/core/ai-chat/ai-chat.controller.supersede.spec.ts new file mode 100644 index 00000000..ce06d53f --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat.controller.supersede.spec.ts @@ -0,0 +1,279 @@ +import { + BadRequestException, + ConflictException, + ForbiddenException, + HttpException, +} from '@nestjs/common'; +import { AiChatController } from './ai-chat.controller'; +import type { User, Workspace } from '@docmost/db/types/entity.types'; + +/** + * #487 commit 3 — the single concurrency GATE (both modes) + the server supersede + * CAS, at the controller boundary. The gate + CAS run BEFORE res.hijack(), so a + * rejected concurrent start / a CAS branch returns clean JSON (an HttpException + * the controller's post-hijack catch re-serializes). These assert the OBSERVABLE + * HTTP contract against the real controller + a stubbed run service. + */ +describe('#487 AiChatController.stream — gate + supersede', () => { + const user = { id: 'u1' } as User; + + function wsWith(autonomousRuns: boolean): Workspace { + return { + id: 'ws1', + settings: { ai: { chat: true, autonomousRuns } }, + } as unknown as Workspace; + } + + function makeReqRes(body: Record) { + const req = { + raw: { sessionId: 'sess', once: jest.fn(), destroyed: false }, + body, + }; + const res = { + raw: { + writableEnded: false, + headersSent: false, + on: jest.fn(), + once: jest.fn(), + setHeader: jest.fn(), + end: jest.fn(), + statusCode: 200, + flushHeaders: jest.fn(), + }, + hijack: jest.fn(), + status: jest.fn().mockReturnThis(), + send: jest.fn(), + }; + return { req, res }; + } + + function makeController( + runServiceOverrides: Record, + // The chat assertOwnedChat resolves. Default: a chat OWNED by `user` (u1), so + // the ownership gate is transparent to the gate/CAS assertions below. Pass a + // foreign-owner (or undefined) chat to exercise the #487 owner rejection. + chat: { creatorId: string } | undefined = { creatorId: 'u1' }, + ) { + const aiChatService = { + resolveRoleForRequest: jest.fn().mockResolvedValue(null), + getChatModel: jest.fn().mockResolvedValue({}), + stream: jest.fn().mockResolvedValue(undefined), + }; + const aiChatRunService = { + getActiveForChat: jest.fn().mockResolvedValue(undefined), + supersede: jest.fn(), + beginRun: jest.fn().mockResolvedValue({ + runId: 'run-new', + signal: new AbortController().signal, + }), + linkAssistantMessage: jest.fn(), + recordStep: jest.fn(), + finalizeRun: jest.fn(), + requestStop: jest.fn(), + ...runServiceOverrides, + }; + const aiChatRepo = { findById: jest.fn().mockResolvedValue(chat) }; + const controller = new AiChatController( + aiChatService as never, + aiChatRunService as never, + aiChatRepo as never, // aiChatRepo + {} as never, // aiChatMessageRepo + {} as never, // aiTranscription + {} as never, // pageRepo + ); + return { controller, aiChatService, aiChatRunService, aiChatRepo }; + } + + const codeOf = (err: unknown) => + (((err as HttpException).getResponse() as Record) ?? {}) + .code; + + describe('single concurrency gate — BOTH modes reject the second tab with 409', () => { + for (const autonomousRuns of [true, false]) { + it(`rejects a concurrent start with 409 A_RUN_ALREADY_ACTIVE (autonomousRuns=${autonomousRuns})`, async () => { + const { controller, aiChatRunService } = makeController({ + getActiveForChat: jest + .fn() + .mockResolvedValue({ id: 'run-live', chatId: 'c1' }), + }); + const { req, res } = makeReqRes({ chatId: 'c1' }); + let thrown: unknown; + try { + await controller.stream( + req as never, + res as never, + user, + wsWith(autonomousRuns), + ); + } catch (e) { + thrown = e; + } + expect(thrown).toBeInstanceOf(ConflictException); + expect((thrown as HttpException).getStatus()).toBe(409); + expect(codeOf(thrown)).toBe('A_RUN_ALREADY_ACTIVE'); + // Rejected BEFORE committing to the stream (no hijack, no service.stream). + expect(res.hijack).not.toHaveBeenCalled(); + expect(aiChatRunService.getActiveForChat).toHaveBeenCalledWith( + 'c1', + 'ws1', + ); + }); + } + }); + + // #487 [security, F1]: stream() MUST owner-gate an existing chat exactly like its + // six sibling endpoints, BEFORE the supersede CAS. Otherwise a same-workspace + // non-owner could POST a supersede against another user's chat and (a) harvest + // that user's active runId from the 409 SUPERSEDE_TARGET_MISMATCH body, then (b) + // requestStop the foreign run. The gate must reject FIRST — no run lookup, no + // supersede, no stop, no runId leak. + describe('cross-user ownership gate (F1)', () => { + it('a non-owner streaming against someone else\'s chat is rejected (403) with NO runId leak and NO foreign requestStop', async () => { + // A live run exists on the victim's chat. Without the gate the supersede CAS + // would run and (faithful to the run service) return a MISMATCH carrying the + // victim's runId — the exact leak. With the gate it must never be reached. + const getActiveForChat = jest + .fn() + .mockResolvedValue({ id: 'run-victim', chatId: 'c-other' }); + const supersede = jest + .fn() + .mockResolvedValue({ kind: 'mismatch', activeRunId: 'run-victim' }); + const requestStop = jest.fn(); + const { controller, aiChatService } = makeController( + { getActiveForChat, supersede, requestStop }, + { creatorId: 'someone-else' }, // the chat is NOT owned by u1 + ); + const { req, res } = makeReqRes({ + chatId: 'c-other', + supersede: { runId: 'guessed-uuid' }, + }); + let thrown: unknown; + try { + await controller.stream(req as never, res as never, user, wsWith(true)); + } catch (e) { + thrown = e; + } + // Rejected by the ownership gate (403), the SAME shape the neighbors use. + expect(thrown).toBeInstanceOf(ForbiddenException); + expect((thrown as HttpException).getStatus()).toBe(403); + // Crucially NOT a 409 that would carry activeRunId — no runId is leaked. + const payload = JSON.stringify( + (thrown as HttpException).getResponse() ?? {}, + ); + expect(payload).not.toContain('run-victim'); + expect(codeOf(thrown)).not.toBe('SUPERSEDE_TARGET_MISMATCH'); + // The gate short-circuits BEFORE any run machinery runs. + expect(getActiveForChat).not.toHaveBeenCalled(); + expect(supersede).not.toHaveBeenCalled(); + expect(requestStop).not.toHaveBeenCalled(); + expect(aiChatService.stream).not.toHaveBeenCalled(); + expect(res.hijack).not.toHaveBeenCalled(); + }); + }); + + it('supersede MISMATCH -> 409 SUPERSEDE_TARGET_MISMATCH carrying the current runId', async () => { + const { controller } = makeController({ + supersede: jest + .fn() + .mockResolvedValue({ kind: 'mismatch', activeRunId: 'run-other' }), + }); + const { req, res } = makeReqRes({ + chatId: 'c1', + supersede: { runId: 'run-x' }, + }); + let thrown: unknown; + try { + await controller.stream(req as never, res as never, user, wsWith(true)); + } catch (e) { + thrown = e; + } + expect(thrown).toBeInstanceOf(ConflictException); + expect(codeOf(thrown)).toBe('SUPERSEDE_TARGET_MISMATCH'); + expect( + ((thrown as HttpException).getResponse() as Record) + .activeRunId, + ).toBe('run-other'); + expect(res.hijack).not.toHaveBeenCalled(); + }); + + it('supersede TIMEOUT -> 409 SUPERSEDE_TIMEOUT, nothing streamed', async () => { + const { controller } = makeController({ + supersede: jest.fn().mockResolvedValue({ kind: 'timeout' }), + }); + const { req, res } = makeReqRes({ + chatId: 'c1', + supersede: { runId: 'run-x' }, + }); + let thrown: unknown; + try { + await controller.stream(req as never, res as never, user, wsWith(false)); + } catch (e) { + thrown = e; + } + expect(thrown).toBeInstanceOf(ConflictException); + expect(codeOf(thrown)).toBe('SUPERSEDE_TIMEOUT'); + expect(res.hijack).not.toHaveBeenCalled(); + }); + + it('supersede INVALID (target on another chat) -> 400 SUPERSEDE_INVALID', async () => { + const { controller } = makeController({ + supersede: jest.fn().mockResolvedValue({ kind: 'invalid' }), + }); + const { req, res } = makeReqRes({ + chatId: 'c1', + supersede: { runId: 'run-x' }, + }); + let thrown: unknown; + try { + await controller.stream(req as never, res as never, user, wsWith(true)); + } catch (e) { + thrown = e; + } + expect(thrown).toBeInstanceOf(BadRequestException); + expect(codeOf(thrown)).toBe('SUPERSEDE_INVALID'); + }); + + it('supersede without chatId -> 400 SUPERSEDE_INVALID', async () => { + const { controller, aiChatRunService } = makeController({}); + const { req, res } = makeReqRes({ supersede: { runId: 'run-x' } }); + let thrown: unknown; + try { + await controller.stream(req as never, res as never, user, wsWith(true)); + } catch (e) { + thrown = e; + } + expect(thrown).toBeInstanceOf(BadRequestException); + expect(codeOf(thrown)).toBe('SUPERSEDE_INVALID'); + expect(aiChatRunService.supersede).not.toHaveBeenCalled(); + }); + + it('supersede READY -> proceeds to stream with superseded=true', async () => { + const { controller, aiChatService } = makeController({ + supersede: jest.fn().mockResolvedValue({ kind: 'ready' }), + getActiveForChat: jest.fn().mockResolvedValue(undefined), // slot free after CAS + }); + const { req, res } = makeReqRes({ + chatId: 'c1', + supersede: { runId: 'run-x' }, + }); + await controller.stream(req as never, res as never, user, wsWith(true)); + expect(res.hijack).toHaveBeenCalled(); + expect(aiChatService.stream).toHaveBeenCalledTimes(1); + expect(aiChatService.stream.mock.calls[0][0].superseded).toBe(true); + // The run hooks are always present now (both modes). + expect(aiChatService.stream.mock.calls[0][0].runHooks).toBeDefined(); + }); + + it('supersede DEGRADE -> proceeds to a normal send (superseded=false)', async () => { + const { controller, aiChatService } = makeController({ + supersede: jest.fn().mockResolvedValue({ kind: 'degrade' }), + }); + const { req, res } = makeReqRes({ + chatId: 'c1', + supersede: { runId: 'run-x' }, + }); + await controller.stream(req as never, res as never, user, wsWith(false)); + expect(aiChatService.stream).toHaveBeenCalledTimes(1); + expect(aiChatService.stream.mock.calls[0][0].superseded).toBe(false); + }); +}); 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 8cc7b03d..f9b58b98 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.ts @@ -418,6 +418,19 @@ export class AiChatController { const body = (req.body ?? {}) as AiChatStreamBody; + // #487 [security]: gate cross-user access to an EXISTING chat BEFORE anything + // reads its runs. Every sibling endpoint (getRun/stop/history/rename/delete/ + // attachRunStream) owner-checks the chat via assertOwnedChat; stream() must too. + // Without this a same-workspace member who is NOT the chat owner could POST a + // supersede against another user's chat and (a) harvest that user's active runId + // out of the 409 SUPERSEDE_TARGET_MISMATCH body, then (b) requestStop the foreign + // run. Gate on the chatId the client sent, when present — a brand-new chat (no + // chatId) has no prior owner to check. Mirrors /stop's owner check (403 as the + // neighbors do), and runs pre-hijack so it returns clean JSON. + if (body.chatId) { + await this.assertOwnedChat(body.chatId, user, workspace); + } + // Resolve the agent role for this turn BEFORE hijack: existing chats read it // from ai_chats.role_id (authoritative), a new chat from body.roleId. The // role drives both the persona and the optional model override below. @@ -432,12 +445,66 @@ export class AiChatController { // HttpException) instead of breaking mid-stream. const model = await this.aiChatService.getChatModel(workspace.id, role); - // #184: one active run per chat. For an EXISTING chat reject a concurrent - // start with a clean 409 BEFORE hijack (the common double-submit / second-tab - // case), so the user gets JSON, not a mid-stream error. A brand-new chat - // (no chatId) cannot have a prior run, and the DB partial unique index is the - // backstop against any race that slips past this check. - if (autonomousRuns && body.chatId) { + // #487: server-side supersede CAS ("interrupt and send now"). When the client + // asks to replace a live run, atomically STOP it and wait for it to settle + // before this turn claims the slot. Runs BEFORE hijack so every branch returns + // clean JSON (the client keeps the composer text on a 409). See + // AiChatRunService.supersede for the branch semantics. + let superseded = false; + const supersedeRunId = body.supersede?.runId; + if (supersedeRunId) { + if (!body.chatId) { + throw new BadRequestException({ + message: 'supersede requires chatId', + code: 'SUPERSEDE_INVALID', + }); + } + const result = await this.aiChatRunService.supersede( + body.chatId, + supersedeRunId, + workspace.id, + ); + switch (result.kind) { + case 'invalid': + throw new BadRequestException({ + message: 'The run to supersede does not belong to this chat', + code: 'SUPERSEDE_INVALID', + }); + case 'mismatch': + // A DIFFERENT run is active than the one the client targeted. Surface + // the CURRENT runId; the client does NOT auto-retry (a stale CAS). + throw new ConflictException({ + message: 'A different agent run is now active on this chat', + code: 'SUPERSEDE_TARGET_MISMATCH', + activeRunId: result.activeRunId, + }); + case 'timeout': + // The target did not settle within W — nothing was persisted, the + // composer keeps the text. NOT a rollback: the stop is already issued. + throw new ConflictException({ + message: + 'The previous run did not stop in time; nothing was sent — please try again', + code: 'SUPERSEDE_TIMEOUT', + }); + case 'ready': + // The target stopped and settled: the slot is free. Prompt the new run + // that the old run's last operations may still be applying. + superseded = true; + break; + case 'degrade': + // The run already ended between click and POST — send normally. + break; + } + } + + // #487: one active run per chat — ENFORCED IN BOTH MODES now (legacy mode used + // to have NO gate, so two tabs streamed two parallel turns on one chat, which + // interleaved history and crashed convertToModelMessages). Reject a concurrent + // start with a clean pre-hijack 409 (double-submit / second-tab). A brand-new + // chat (no chatId) cannot have a prior run, and the DB partial unique index in + // beginRun is the authoritative backstop for any race that slips past here + // (including a slot stolen between a supersede release and beginRun). + if (body.chatId) { const active = await this.aiChatRunService.getActiveForChat( body.chatId, workspace.id, @@ -446,107 +513,94 @@ export class AiChatController { throw new ConflictException({ message: 'An agent run is already in progress for this chat', code: 'A_RUN_ALREADY_ACTIVE', + activeRunId: active.id, }); } } - // Run-lifecycle hooks (#184), only when the flag is on. They wrap the turn in - // a durable run whose abort is governed by the run (explicit stop), persist - // its progress, and settle its terminal status — see AiChatRunService. - const runHooks: AiChatRunHooks | undefined = autonomousRuns - ? { - 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, - workspace.id, - messageId, - ), - onStep: (runId, stepCount) => - void this.aiChatRunService.recordStep( - runId, - workspace.id, - stepCount, - ), - onSettled: (runId, status, error) => - this.aiChatRunService.finalizeRun( - runId, - workspace.id, - status, - error, - ), + // #487: the turn is ALWAYS a first-class RUN now (both modes). The mode + // difference is only the abort semantics on a browser disconnect (onClose + // below). currentRunId is captured at begin so a legacy disconnect can stop + // the run through its stop lever. + let currentRunId: string | undefined; + const runHooks: AiChatRunHooks = { + begin: async (chatId) => { + const handle = await this.aiChatRunService.beginRun({ + chatId, + workspaceId: workspace.id, + userId: user.id, + trigger: 'user', + }); + currentRunId = handle?.runId; + // #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. + if ( + handle?.runId && + this.environment?.isAiChatResumableStreamEnabled?.() + ) { + this.streamRegistry?.open(chatId, handle.runId); } - : undefined; + return handle; + }, + onAssistantSeeded: (runId, messageId) => + this.aiChatRunService.linkAssistantMessage( + runId, + workspace.id, + messageId, + ), + onStep: (runId, stepCount) => + void this.aiChatRunService.recordStep(runId, workspace.id, stepCount), + onSettled: (runId, status, error) => + this.aiChatRunService.finalizeRun(runId, workspace.id, status, error), + }; - // Abort the agent loop when the client disconnects. `close` also fires on - // normal completion, so only abort when the response has not finished - // writing (a genuine disconnect). `once` fires at most once and self-removes; - // we also drop it on response `finish` so it never lingers after the stream - // completes normally (the AI SDK pipes the response fire-and-forget, so we - // cannot simply remove it once `stream()` returns). + // Handle a client disconnect. `close` also fires on normal completion, so only + // act when the response has not finished writing (a genuine disconnect). `once` + // fires at most once and self-removes; we also drop it on response `finish`. // DIAGNOSTIC (Safari stream-drop investigation) — temporary: wall-clock at // which a Safari disconnect is observed, measured from request receipt. const reqStartedAt = Date.now(); const controller = new AbortController(); const onClose = (): void => { - // A genuine disconnect leaves the response unfinished (unlike a normal - // completion, which also fires `close`). Such a drop — e.g. a reverse - // proxy cutting the SSE mid-answer — is otherwise invisible server-side, - // so log it here. if (!res.raw.writableEnded) { if (autonomousRuns) { - // #184: the turn is a DETACHED run. A disconnect must NOT abort it — - // the run keeps executing and persisting server-side; the client - // reconnects via /ai-chat/run (or re-stops via /ai-chat/stop). Log only. + // #184: a DETACHED run — a disconnect must NOT stop it. The run keeps + // executing and persisting server-side; the client reconnects via + // /ai-chat/run (or re-stops via /ai-chat/stop). Log only. this.logger.log( `AI chat stream: client disconnected; run continues server-side ` + `(elapsed=${Date.now() - reqStartedAt}ms since request received)`, ); } else { + // #487: legacy — a disconnect ENDS the turn, but the turn is now a RUN, + // so stop it through the run's stop lever (requestStop). streamText no + // longer consumes the socket signal (effectiveSignal is the run signal), + // so aborting `controller` would do nothing; requestStop aborts the run. this.logger.warn( - `AI chat stream: client disconnected before completion; aborting turn ` + - `(elapsed=${Date.now() - reqStartedAt}ms since request received)`, + `AI chat stream: client disconnected before completion; stopping the ` + + `run (elapsed=${Date.now() - reqStartedAt}ms since request received)`, ); - controller.abort(); + if (currentRunId) { + void this.aiChatRunService.requestStop(currentRunId, workspace.id); + } } } }; req.raw.once('close', onClose); res.raw.once('finish', () => req.raw.off('close', onClose)); - // #184: in detached mode the turn is NOT aborted on disconnect, so the SDK's - // pipe keeps writing to a socket the client may have dropped — for the rest of - // the (continuing) run. A write to the dead socket can emit an 'error' on the - // raw response; without a listener that surfaces as an unhandled error event. - // Swallow it (the run continues server-side regardless). Legacy mode aborts on - // disconnect, so it does not need this and keeps its exact prior behavior. - if (autonomousRuns) { - res.raw.on('error', (err) => { - this.logger.debug( - `AI chat detached stream: post-disconnect socket error swallowed: ${ - err instanceof Error ? err.message : String(err) - }`, - ); - }); - } + // #184/#487: the run/pipe can outlive the socket in BOTH modes now (autonomous + // keeps going; legacy keeps going until requestStop's abort unwinds the turn). + // The SDK's pipe may then write to a dropped socket and emit an 'error' on the + // raw response — swallow it so it never surfaces as an unhandled error event. + res.raw.on('error', (err) => { + this.logger.debug( + `AI chat stream: post-disconnect socket error swallowed: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + }); // Commit to streaming: hijack so Fastify stops managing the response and // the AI SDK can write the UI-message stream directly to the Node socket. @@ -562,8 +616,10 @@ export class AiChatController { signal: controller.signal, model, role, - // #184: present only when the flag is on; wraps the turn in a durable run. + // #487: the turn is always run-wrapped now (both modes). runHooks, + // #487: warn the new run that a superseded run's last ops may still apply. + superseded, }); } catch (err) { // Any failure AFTER hijack can no longer go through Nest's exception diff --git a/apps/server/src/core/ai-chat/ai-chat.prompt.ts b/apps/server/src/core/ai-chat/ai-chat.prompt.ts index 5107a77d..33c222fd 100644 --- a/apps/server/src/core/ai-chat/ai-chat.prompt.ts +++ b/apps/server/src/core/ai-chat/ai-chat.prompt.ts @@ -101,6 +101,22 @@ const INTERRUPT_NOTE = 'assume your previous response was complete, and do not silently restart the ' + 'partial work — build on it or follow the new instruction.'; +/** + * #487: injected on a turn started by SUPERSEDING a previous run (the user hit + * "interrupt and send now" while a run was live). The previous run was Stopped, + * but there is NO side-effect quiescence — a write it had already committed, or + * one committing at the moment of Stop, may land with a small delay AFTER this new + * run starts. So the model is told its picture of the page/state may be a beat + * stale and to re-read before assuming an edit did or did not apply. + */ +const SUPERSEDE_NOTE = + 'NOTE: A previous agent run in this conversation was just interrupted so this ' + + 'new turn could start. That run was stopped, but any operation it had already ' + + 'begun (e.g. a page edit) may still be applied with a short delay. Do not ' + + 'assume the document/state is exactly as the interrupted run left it — if you ' + + 'need to rely on the current content, RE-READ it with the page tools before ' + + 'acting rather than trusting a cached view.'; + /** * Injected on a turn where the open page was hand-edited by the user (or anyone * else) AFTER the agent's previous response ended (#274). The server takes a @@ -203,6 +219,14 @@ export interface BuildSystemPromptInput { * (partial) answer was cut off by the user's new message. */ interrupted?: boolean; + /** + * #487: true when THIS turn was started by superseding a still-live previous run + * ("interrupt and send now"). Adds SUPERSEDE_NOTE so the model knows the previous + * run's last operations may still be applying and to re-read state it depends on. + * Distinct from `interrupted` (which is about a PARTIAL prior answer in history); + * both can be set together. Self-clears — set only for the superseding turn. + */ + superseded?: boolean; /** * Set only when the open page was edited by the user AFTER the agent's previous * turn ended (#274), confirmed server-side by diffing the current page against @@ -311,6 +335,7 @@ export function buildSystemPrompt({ openedPage, mcpInstructions, interrupted, + superseded, pageChanged, deferredToolsEnabled, toolCatalog, @@ -360,6 +385,13 @@ export function buildSystemPrompt({ context += `\n${INTERRUPT_NOTE}`; } + // Supersede note (#487): present only for a turn that stopped and replaced a + // still-live previous run — warns the model the previous run's last operations + // may still be applying (no side-effect quiescence). + if (superseded) { + context += `\n${SUPERSEDE_NOTE}`; + } + // Per-turn page-change note (#274). Added to the context section (inside the // safety sandwich), present only when the server detected that the open page // was edited by the user since the agent's last turn ended. The diff content is diff --git a/apps/server/src/core/ai-chat/ai-chat.service.lifecycle.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.lifecycle.spec.ts index 66eb5b42..a6f3e75c 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.lifecycle.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.lifecycle.spec.ts @@ -89,6 +89,11 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => { const runRepo = { insert: jest.fn().mockResolvedValue({ id: 'run-1', status: 'running' }), update: jest.fn().mockResolvedValue({ id: 'run-1' }), + // #487: the terminal settle now goes through the CONDITIONAL write. + finalizeIfActive: jest + .fn() + .mockResolvedValue({ id: 'run-1', status: 'failed' }), + findById: jest.fn().mockResolvedValue(undefined), }; const runService = new AiChatRunService(runRepo as never, { isCloud: () => false } as never); @@ -148,9 +153,10 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => { // 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( + // ...then settled to a terminal FAILED status by the safety net (via the + // #487 conditional write)... + expect(runRepo.finalizeIfActive).toHaveBeenCalledTimes(1); + expect(runRepo.finalizeIfActive).toHaveBeenCalledWith( 'run-1', 'ws1', expect.objectContaining({ status: 'failed' }), diff --git a/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts index d167abbe..8b170c11 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts @@ -155,6 +155,8 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => { insert: jest.fn(async () => ({ id: 'msg-1' })), findAllByChat: jest.fn(async () => []), update: jest.fn(async () => ({ id: 'msg-1' })), + finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })), + findStreamingWithTerminalRun: jest.fn(async () => []), }; const aiSettings = { resolve: jest.fn(async () => ({})) }; const tools = { forUser: jest.fn(async () => ({})) }; @@ -332,7 +334,13 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => { usage: {}, steps: [], }); - expect(runHooks.onSettled).toHaveBeenCalledWith('run-1', 'completed'); + // #487: onFinish passes the (undefined) error slot so a message-finalize + // failure could error-mark the run; on the success path it is undefined. + expect(runHooks.onSettled).toHaveBeenCalledWith( + 'run-1', + 'completed', + undefined, + ); }); it('F9: onAbort settles the run "aborted"', async () => { @@ -415,6 +423,8 @@ describe('AiChatService.stream — begin-failure fails the turn (#184 F14 / #486 insert: jest.fn(async () => ({ id: 'msg-1' })), findAllByChat: jest.fn(async () => []), update: jest.fn(async () => ({ id: 'msg-1' })), + finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })), + findStreamingWithTerminalRun: jest.fn(async () => []), }; const aiSettings = { resolve: jest.fn(async () => ({})) }; const tools = { forUser: jest.fn(async () => ({})) }; 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 46e37f61..ca9511ef 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 @@ -1315,8 +1315,12 @@ describe('AiChatService page-change lifecycle (#274)', () => { describe('isInterruptResume', () => { // history tail is the just-inserted user row; [len-2] is the previous turn. const withPrev = ( - prev: { role: string; status?: string | null } | null, - ): Array<{ role: string; status?: string | null }> => + prev: { + role: string; + status?: string | null; + metadata?: unknown; + } | null, + ): Array<{ role: string; status?: string | null; metadata?: unknown }> => prev ? [prev, { role: 'user', status: null }] : [{ role: 'user', status: null }]; @@ -1357,6 +1361,33 @@ describe('isInterruptResume', () => { it('false when there is no preceding turn (only the new user row)', () => { expect(isInterruptResume(withPrev(null), true)).toBe(false); }); + + it('#487 EXCLUDES a reconcile stamp (finalizeFailed) — not a genuine interruption', () => { + // A row a reconcile settled to 'aborted' carries metadata.finalizeFailed. It + // must NOT be treated as an interrupt-resume (that would inject a false + // "you were interrupted" note), even though its status is 'aborted'. + expect( + isInterruptResume( + withPrev({ + role: 'assistant', + status: 'aborted', + metadata: { finalizeFailed: true }, + }), + true, + ), + ).toBe(false); + // A genuine abort (no finalizeFailed) still counts. + expect( + isInterruptResume( + withPrev({ + role: 'assistant', + status: 'aborted', + metadata: { parts: [] }, + }), + true, + ), + ).toBe(true); + }); }); /** @@ -1419,6 +1450,9 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', () insert: jest.fn(async () => ({ id: 'msg-1' })), findAllByChat: jest.fn(async () => []), update: jest.fn(async () => ({ id: 'msg-1' })), + // #487: the terminal owner-write + the opportunistic reconcile query. + finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })), + findStreamingWithTerminalRun: jest.fn(async () => []), }; const aiSettings = { resolve: jest.fn(async () => ({})) }; const tools = { forUser: jest.fn(async () => ({})) }; @@ -1623,6 +1657,19 @@ describe('AiChatService.stream — token-degeneration reaction (#444)', () => { return { id }; }, ), + // #487: the terminal owner-write records into the SAME `updated` recorder so + // assertions on the terminal 'completed'/'error'/'aborted' write still hold. + finalizeOwner: jest.fn( + async ( + id: string, + workspaceId: string, + patch: Record, + ) => { + updated.push({ id, workspaceId, patch }); + return { id }; + }, + ), + findStreamingWithTerminalRun: jest.fn(async () => []), }; const aiSettings = { resolve: jest.fn(async () => ({})) }; const tools = { forUser: jest.fn(async () => ({})) }; @@ -1882,3 +1929,148 @@ describe('AiChatService.stream — token-degeneration reaction (#444)', () => { expect(patch.content).not.toContain(STEP_LIMIT_NO_ANSWER_MARKER); }); }); + +// #487 F3 — the reconcile() / reconcileChat() ORCHESTRATORS. The individual +// clauses are exercised elsewhere; these pin the production orchestration the +// per-clause specs do not: the clause ORDER, the per-clause try/catch ISOLATION +// (one clause throwing must NOT abort the others), and reconcileChat() (which runs +// at the start of every turn and was entirely uncovered). +describe('AiChatService.reconcile / reconcileChat orchestrators (#487 F3)', () => { + let warnSpy: jest.SpyInstance; + beforeEach(() => { + // Silence the intentional clause-failure warnings (kept out of test output). + warnSpy = jest + .spyOn(Logger.prototype, 'warn') + .mockImplementation(() => undefined); + }); + afterEach(() => { + warnSpy.mockRestore(); + }); + + function makeService(opts: { + messageRepo?: Record; + runService?: Record; + }) { + const aiChatMessageRepo = { + findStreamingWithTerminalRun: jest.fn(async () => []), + stampTerminalIfStreaming: jest.fn(async () => undefined), + sweepStreamingWithoutActiveRun: jest.fn(async () => 0), + ...(opts.messageRepo ?? {}), + }; + const aiChatRunService = opts.runService + ? { + zombieRunIds: jest.fn(() => []), + settleZombie: jest.fn(async () => true), + reconcileStaleRuns: jest.fn(async () => 0), + ...opts.runService, + } + : undefined; + const svc = new AiChatService( + {} as never, // ai + {} as never, // aiChatRepo + aiChatMessageRepo as never, + {} as never, // aiChatPageSnapshotRepo + {} as never, // aiSettings + {} as never, // tools + {} as never, // mcpClients + {} as never, // aiAgentRoleRepo + {} as never, // pageRepo + {} as never, // pageAccess + {} as never, // environment + {} as never, // streamRegistry + aiChatRunService as never, // aiChatRunService (#487) + ); + return { svc, aiChatMessageRepo, aiChatRunService }; + } + + it('reconcile() fires all four clauses IN ORDER (a -> b -> c -> d)', async () => { + const order: string[] = []; + const { svc } = makeService({ + messageRepo: { + findStreamingWithTerminalRun: jest.fn(async () => { + order.push('b:find'); + return [ + { messageId: 'm1', workspaceId: 'ws1', runStatus: 'succeeded' }, + ]; + }), + stampTerminalIfStreaming: jest.fn(async () => { + order.push('b:stamp'); + }), + sweepStreamingWithoutActiveRun: jest.fn(async () => { + order.push('d'); + return 0; + }), + }, + runService: { + zombieRunIds: jest.fn(() => ['z1']), + settleZombie: jest.fn(async () => { + order.push('a'); + return true; + }), + reconcileStaleRuns: jest.fn(async () => { + order.push('c'); + return 0; + }), + }, + }); + + await svc.reconcile(); + + expect(order).toEqual(['a', 'b:find', 'b:stamp', 'c', 'd']); + }); + + it('a clause that THROWS does not abort the remaining clauses (per-clause try/catch isolation)', async () => { + const { svc, aiChatMessageRepo, aiChatRunService } = makeService({ + messageRepo: { + // Clause (b) blows up mid-reconcile. + findStreamingWithTerminalRun: jest.fn(async () => { + throw new Error('clause b DB blip'); + }), + }, + runService: { + zombieRunIds: jest.fn(() => ['z1']), + }, + }); + + // reconcile() must SETTLE (the clause-b failure is swallowed), not reject. + await expect(svc.reconcile()).resolves.toBeUndefined(); + + // (a) ran before (b); crucially (c) and (d) STILL ran despite (b) throwing — + // the property a missing try/catch would break. MUTATION-VERIFY: drop clause + // (b)'s try/catch and this reddens (the throw propagates, skipping c + d). + expect(aiChatRunService!.settleZombie).toHaveBeenCalled(); // (a) + expect(aiChatRunService!.reconcileStaleRuns).toHaveBeenCalled(); // (c) + expect( + aiChatMessageRepo.sweepStreamingWithoutActiveRun, + ).toHaveBeenCalled(); // (d) + }); + + it('reconcileChat() settles THIS chat\'s stuck streaming rows by their run status', async () => { + const { svc, aiChatMessageRepo } = makeService({ + messageRepo: { + findStreamingWithTerminalRun: jest.fn(async () => [ + { messageId: 'm1', workspaceId: 'ws1', runStatus: 'failed' }, + { messageId: 'm2', workspaceId: 'ws1', runStatus: 'succeeded' }, + ]), + }, + }); + + await svc.reconcileChat('chat-1', 'ws1'); + + // Scoped to THIS chat and bounded at 50 (the user-facing opportunistic path). + expect( + aiChatMessageRepo.findStreamingWithTerminalRun, + ).toHaveBeenCalledWith(50, { chatId: 'chat-1', workspaceId: 'ws1' }); + // failed-run -> 'error'; every other terminal status -> 'aborted'. + expect(aiChatMessageRepo.stampTerminalIfStreaming).toHaveBeenCalledWith( + 'm1', + 'ws1', + 'error', + ); + expect(aiChatMessageRepo.stampTerminalIfStreaming).toHaveBeenCalledWith( + 'm2', + 'ws1', + 'aborted', + ); + }); +}); 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 14826bd6..02ecc9d1 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -3,6 +3,7 @@ import { ForbiddenException, Injectable, Logger, + OnModuleDestroy, OnModuleInit, ServiceUnavailableException, } from '@nestjs/common'; @@ -42,7 +43,11 @@ import { makeLoadToolsTool, buildExternalToolCatalog, } from './tools/tool-tiers'; -import { RunAlreadyActiveError } from './ai-chat-run.service'; +import { + RunAlreadyActiveError, + AiChatRunService, +} from './ai-chat-run.service'; +import { inAppToolCallCapMs } from './tools/ai-chat-tools.service'; import { computePageChange } from './page-change/page-change.util'; import { sanitizeSelection, @@ -250,15 +255,23 @@ export function cleanGeneratedTitle(text: string): string { * partial output is already in history thanks to the step-granular write path). */ export function isInterruptResume( - history: Array<{ role: string; status?: string | null }>, + history: Array<{ + role: string; + status?: string | null; + metadata?: unknown; + }>, clientInterrupted: boolean | undefined, ): boolean { if (clientInterrupted !== true) return false; const prev = history[history.length - 2]; - return ( - prev?.role === 'assistant' && - (prev.status === 'aborted' || prev.status === 'streaming') - ); + if (prev?.role !== 'assistant') return false; + // #487: a reconcile STAMP (metadata.finalizeFailed) is NOT a genuine user + // interruption — the previous turn's process died and a reconcile settled the + // row as 'aborted'. Treating it as an interrupt-resume would inject a false + // "you were interrupted" note. Exclude any finalizeFailed row. + const meta = prev.metadata as { finalizeFailed?: unknown } | null | undefined; + if (meta && meta.finalizeFailed === true) return false; + return prev.status === 'aborted' || prev.status === 'streaming'; } /** @@ -379,6 +392,14 @@ export interface AiChatStreamBody { // it against persisted history (`isInterruptResume`) before injecting the // interrupt note, so a spoofed/stale flag on an ordinary turn is ignored. interrupted?: boolean; + // #487: server-side supersede CAS. When present, this POST asks the server to + // STOP the run `supersede.runId` (which the client saw as the chat's active run) + // and, once it has settled, start THIS turn in its place. The server validates + // the target against the chat and answers 400 (wrong chat) / 409 + // SUPERSEDE_TARGET_MISMATCH / 409 SUPERSEDE_TIMEOUT, or proceeds normally + // (degrade / ready). Absent => an ordinary send (rejected with 409 + // A_RUN_ALREADY_ACTIVE if a run is already active on the chat). + supersede?: { runId?: string } | null; // useChat sends the full UIMessage list; the last one is the new user turn. messages?: UIMessage[]; } @@ -428,6 +449,11 @@ export interface AiChatStreamArgs { // chat row (existing chat) or the request body (new chat). null => universal // assistant. Carried here so the turn never re-loads it. role: AiAgentRole | null; + // #487: true when this turn was started by SUPERSEDING a still-live previous run + // (the controller ran the supersede CAS to a `ready` result). Adds the + // SUPERSEDE_NOTE to the system prompt (the previous run's last ops may still be + // applying — no side-effect quiescence). Absent on an ordinary send. + superseded?: boolean; } /** @@ -444,7 +470,7 @@ export interface AiChatStreamArgs { * can be rebuilt for `convertToModelMessages`. */ @Injectable() -export class AiChatService implements OnModuleInit { +export class AiChatService implements OnModuleInit, OnModuleDestroy { private readonly logger = new Logger(AiChatService.name); constructor( @@ -465,8 +491,17 @@ export class AiChatService implements OnModuleInit { // 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, + // #487: the run lifecycle service, for the periodic + opportunistic reconcile + // (zombie re-drive + stale-run abort). OPTIONAL so positional test + // constructions compile unchanged; Nest always injects the real singleton, so + // reconcile sees the SAME in-memory active/zombie maps the runner mutates. + private readonly aiChatRunService?: AiChatRunService, ) {} + // #487: periodic reconcile timer (single-process phase 1). Started in + // onModuleInit, cleared in onModuleDestroy. + private reconcileTimer?: ReturnType; + /** * Crash-recovery sweep on server start (#183): any assistant row left in the * 'streaming' state is the relic of a turn whose process died before it @@ -491,6 +526,158 @@ export class AiChatService implements OnModuleInit { }`, ); } + + // #487: start the PERIODIC reconcile (was boot-only). It heals both directions + // of the run<->message lifecycle asymmetry that a boot sweep alone left to the + // NEXT restart. Single-process phase 1: the in-memory active/zombie maps are + // authoritative, so "no live entry" is a safe primary gate. + const staleMs = this.reconcileStalenessMs(); + // boot-warn if the per-call cap is configured so high the derived staleness is + // unusually long (a stale run then lingers longer before reconcile aborts it). + if (staleMs > 30 * 60 * 1000) { + this.logger.warn( + `#487 reconcile staleness is ${Math.round(staleMs / 60000)}min ` + + `(derived from max(2 x per-call cap, 15min)); a per-call cap this high ` + + `delays stale-run recovery. Review AI_CHAT_INAPP_TOOL_CALL_CAP_MS.`, + ); + } + const intervalMs = this.reconcileIntervalMs(); + this.reconcileTimer = setInterval(() => { + void this.reconcile().catch((err) => { + this.logger.warn( + `Periodic reconcile failed: ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + }); + }, intervalMs); + this.reconcileTimer.unref?.(); + } + + /** #487: stop the periodic reconcile timer on shutdown. */ + onModuleDestroy(): void { + if (this.reconcileTimer) { + clearInterval(this.reconcileTimer); + this.reconcileTimer = undefined; + } + } + + /** + * #487: reconcile staleness threshold X — a run/message is only a "no live + * runner" abort candidate once UNTOUCHED past this. Derived as + * max(2 x per-call cap, 15min): 2x the longest legitimate single tool call plus + * a floor, so a marathon turn making steady progress (updatedAt bumped each + * step) is never swept. + */ + private reconcileStalenessMs(): number { + return Math.max(2 * inAppToolCallCapMs(), 15 * 60 * 1000); + } + + /** #487: how often the periodic reconcile runs (env-tunable, default 2min). */ + private reconcileIntervalMs(): number { + const raw = Number(process.env.AI_CHAT_RECONCILE_INTERVAL_MS); + return Number.isFinite(raw) && raw > 0 ? raw : 2 * 60 * 1000; + } + + /** + * #487: the periodic BIDIRECTIONAL reconcile. Runs the clauses IN ORDER; each is + * best-effort (a failure of one never blocks the others). Single-process phase 1 + * — the run service's in-memory maps are authoritative for "live entry". + * + * (a) re-drive ZOMBIE runs (a terminal write that gave up) — apply the intended + * status via the conditional UPDATE; + * (b) message 'streaming' + its RUN terminal -> stamp the message by the run's + * status (succeeded-run + stuck row -> 'aborted'+finalizeFailed, NOT + * 'completed' with empty parts — the final text lived only in the dead + * process's memory, a documented loss); + * (c) run active + NO live entry + NO zombie + stale -> aborted (the run + * service applies the "no entry" primary gate + last-progress staleness); + * (d) message 'streaming' + age>X + NO active run on the chat -> aborted + * (historical-row safety, double-gated). + */ + async reconcile(): Promise { + const staleMs = this.reconcileStalenessMs(); + + // (a) zombie re-drive. + if (this.aiChatRunService) { + for (const runId of this.aiChatRunService.zombieRunIds()) { + try { + await this.aiChatRunService.settleZombie(runId); + } catch (err) { + this.logger.warn( + `Reconcile (a) zombie ${runId} re-drive failed: ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + } + } + } + + // (b) message streaming + run terminal -> stamp message by run status. + try { + const stuck = await this.aiChatMessageRepo.findStreamingWithTerminalRun(); + for (const s of stuck) { + // succeeded-run -> 'aborted' (NOT 'completed'-empty); failed -> 'error'; + // aborted -> 'aborted'. All via the finalizeFailed stamp. + const status = s.runStatus === 'failed' ? 'error' : 'aborted'; + await this.aiChatMessageRepo.stampTerminalIfStreaming( + s.messageId, + s.workspaceId, + status, + ); + } + } catch (err) { + this.logger.warn( + `Reconcile (b) message<-run failed: ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + } + + // (c) stale active run with no live runner -> aborted. + if (this.aiChatRunService) { + try { + await this.aiChatRunService.reconcileStaleRuns(staleMs); + } catch (err) { + this.logger.warn( + `Reconcile (c) stale-run abort failed: ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + } + } + + // (d) historical streaming row, no active run on the chat, stale -> aborted. + try { + await this.aiChatMessageRepo.sweepStreamingWithoutActiveRun(staleMs); + } catch (err) { + this.logger.warn( + `Reconcile (d) historical-row sweep failed: ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + } + } + + /** + * #487: OPPORTUNISTIC single-chat reconcile at the start of a turn (beginRun / + * supersede path), so a user who returns to a chat with a stuck streaming row + * (its run already terminal) sees it settled WITHOUT waiting for the periodic + * job. Best-effort — a failure NEVER fails the turn (swallowed by the caller). + */ + async reconcileChat(chatId: string, workspaceId: string): Promise { + const stuck = await this.aiChatMessageRepo.findStreamingWithTerminalRun(50, { + chatId, + workspaceId, + }); + for (const s of stuck) { + const status = s.runStatus === 'failed' ? 'error' : 'aborted'; + await this.aiChatMessageRepo.stampTerminalIfStreaming( + s.messageId, + s.workspaceId, + status, + ); + } } /** @@ -727,6 +914,7 @@ export class AiChatService implements OnModuleInit { model, role, runHooks, + superseded, }: AiChatStreamArgs): Promise { // Resolve / create the chat. A new chat is created when no valid chatId is // supplied or the supplied one does not belong to this workspace. @@ -835,6 +1023,20 @@ export class AiChatService implements OnModuleInit { } } + // #487: opportunistic single-chat reconcile — settle any streaming row on this + // chat whose run is already terminal BEFORE this turn's history load, so the + // user never waits on the periodic job and the new turn's model history is not + // polluted by a stuck 'streaming' row. Best-effort: it must NEVER fail the turn. + try { + await this.reconcileChat(chatId, workspace.id); + } catch (err) { + this.logger.debug( + `Opportunistic reconcile for chat ${chatId} failed (ignored): ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + } + try { // Extract the incoming user turn (the last user message from useChat). const incoming = lastUserMessage(body.messages); @@ -929,14 +1131,13 @@ export class AiChatService implements OnModuleInit { ); } catch (err) { // An explicit Stop reached the RUN's signal DURING setup: re-throw so the - // outer catch finalizes the run as aborted — never swallow a Stop. Gated on - // `runId`: the re-throw exists ONLY to finalize the run, which exists only - // in autonomous mode. On the legacy path (no runId) `effectiveSignal` is the - // SOCKET signal (it aborts on a client disconnect); re-throwing there would - // change prior behavior and make the controller write JSON to an already- - // closed socket (it only attaches res.raw.on('error') in autonomous mode). - // So legacy keeps its prior behavior — warn + proceed, and streamText then - // observes the aborted socket signal. + // outer catch finalizes the run as aborted — never swallow a Stop. #487: the + // turn is ALWAYS run-wrapped now (both modes), so `effectiveSignal` is the + // RUN signal and `runId` is set in BOTH — a Stop (from /ai-chat/stop or a + // legacy disconnect's requestStop) aborts it identically. The `runId` guard + // now only defends the theoretical no-handle fallback (`begin` returned + // nothing, leaving `effectiveSignal` as the bare socket signal): there we + // keep the old warn-and-proceed rather than re-throw. if (runId && effectiveSignal.aborted) { throw err; } @@ -1040,6 +1241,9 @@ export class AiChatService implements OnModuleInit { // History-confirmed interrupt-resume flag (#198): adds the interrupt note // so the model treats the partial answer above as cut off, not finished. interrupted, + // #487: this turn superseded a still-live run — warn the model the + // previous run's last ops may still be applying (no quiescence). + superseded, // Detected between-turns human edit to the open page (#274): adds the // page_changed note + unified diff so the agent doesn't overwrite it. pageChanged, @@ -1194,29 +1398,59 @@ export class AiChatService implements OnModuleInit { // callbacks — mirroring the pre-#183 persist-at-most-once guard for the // TERMINAL status (the row may be updated many times with 'streaming' before // this fires once). + // #487: the once-gate closes ONLY AFTER a successful write, and the write is + // BOUNDED-RETRIED. Previously `finalized` was set BEFORE the write and never + // retried, so a single failed UPDATE stranded the row 'streaming' forever + // (the boot-only sweep was the only recovery). Now a transient blip is ridden + // out in place; a total give-up leaves the gate OPEN and logs, and the + // periodic reconcile (clauses b/d) later settles the row. Returns whether the + // terminal write LANDED, so the caller can error-mark the RUN on a message + // failure (the run is finalized regardless — never gated on the message). let finalized = false; + const FINALIZE_MSG_MAX_ATTEMPTS = 3; const finalizeAssistant = async ( flushed: AssistantFlush, - ): Promise => { - if (finalized) return; - finalized = true; + ): Promise => { + if (finalized) return true; const plan = planFinalizeAssistant(assistantId); - try { - // Shared dispatch (see applyFinalize): UPDATE the upfront row, or — when - // the upfront insert failed (kind 'insert') — INSERT the terminal row as - // the only safety against losing the turn entirely. - await applyFinalize( - this.aiChatMessageRepo, - plan, - { chatId, workspaceId: workspace.id, userId: user.id }, - flushed, - ); - } catch (err) { - this.logger.error( - `Failed to finalize assistant message (kind=${plan.kind})`, - err as Error, - ); + let lastError: unknown; + for (let attempt = 1; attempt <= FINALIZE_MSG_MAX_ATTEMPTS; attempt++) { + try { + // Shared dispatch (see applyFinalize): conditionally UPDATE the upfront + // row (owner-write priority), or — when the upfront insert failed (kind + // 'insert') — INSERT the terminal row as the only safety against losing + // the turn entirely. + await applyFinalize( + this.aiChatMessageRepo, + plan, + { chatId, workspaceId: workspace.id, userId: user.id }, + flushed, + ); + finalized = true; // gate closes ONLY after a successful write + return true; + } catch (err) { + lastError = err; + this.logger.warn( + `Assistant message finalize attempt ${attempt}/${FINALIZE_MSG_MAX_ATTEMPTS} ` + + `failed (kind=${plan.kind}): ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ); + if (attempt < FINALIZE_MSG_MAX_ATTEMPTS) { + await new Promise((r) => setTimeout(r, 50 * attempt)); + } + } } + // Gave up: leave the gate OPEN (no in-process second settler exists — the + // terminal callbacks are mutually exclusive) and log. The periodic reconcile + // settles the stranded row; a late owner-write is impossible for this turn, + // so the reconcile stamp (aborted+finalizeFailed) is the final state. + this.logger.error( + `Assistant message finalize GAVE UP after ${FINALIZE_MSG_MAX_ATTEMPTS} ` + + `attempts (row left 'streaming', chat ${chatId}); reconcile will settle it`, + lastError as Error, + ); + return false; }; // DIAGNOSTIC (Safari stream-drop investigation) — temporary. Measure @@ -1361,7 +1595,7 @@ export class AiChatService implements OnModuleInit { const stepExhausted = steps.length >= MAX_AGENT_STEPS; const emptyTurnMarker = !producedText && stepExhausted ? STEP_LIMIT_NO_ANSWER_MARKER : ''; - await finalizeAssistant( + const msgOk = await finalizeAssistant( flushAssistant(steps as StepLike[], emptyTurnMarker, 'completed', { finishReason: finishReason as string, usage: totalUsage as StreamUsage, @@ -1375,9 +1609,19 @@ export class AiChatService implements OnModuleInit { pageChanged, }), ); - // #184: settle the RUN as succeeded (best-effort, after the projection - // is finalized above). - if (runId) await runHooks?.onSettled?.(runId, 'completed'); + // #184/#487: the RUN is finalized ALWAYS (never gated on the message). + // If the message finalize GAVE UP, error-mark the run so the asymmetry + // "run succeeded / message streaming forever" cannot arise; the + // periodic reconcile then settles the stuck message from this run. + if (runId) { + await runHooks?.onSettled?.( + runId, + msgOk ? 'completed' : 'error', + msgOk + ? undefined + : 'Assistant message could not be persisted (finalize failed).', + ); + } // Lifecycle: release the external MCP clients leased for this turn. await closeExternalClients(); @@ -2118,7 +2362,10 @@ export function planFinalizeAssistant( * a test mock both satisfy it). */ export interface FinalizeRepo { insert(insertable: Record): Promise; - update( + // #487: the OWNER terminal write is CONDITIONAL (status='streaming' OR + // metadata.finalizeFailed) so the owner overwrites a reconcile stamp but never + // an already-proper terminal row (owner-write priority). + finalizeOwner( id: string, workspaceId: string, patch: AssistantFlush, @@ -2127,10 +2374,11 @@ export interface FinalizeRepo { /** * Apply a finalize `plan` to the repo with the terminal `flushed` payload (#183): - * UPDATE the upfront row, or INSERT a fresh terminal row as the fallback when the - * upfront insert failed. The SINGLE dispatch shared by the service's - * finalizeAssistant and its test, so the test exercises the real path instead of - * a copy (#186 review). Pure of error handling — the caller wraps it. + * conditionally UPDATE the upfront row (owner-write priority, #487), or INSERT a + * fresh terminal row as the fallback when the upfront insert failed. The SINGLE + * dispatch shared by the service's finalizeAssistant and its test, so the test + * exercises the real path instead of a copy (#186 review). Pure of error + * handling — the caller wraps it (and RETRIES it, #487). */ export async function applyFinalize( repo: FinalizeRepo, @@ -2139,7 +2387,7 @@ export async function applyFinalize( flushed: AssistantFlush, ): Promise { if (plan.kind === 'update') { - await repo.update(plan.id, base.workspaceId, flushed); + await repo.finalizeOwner(plan.id, base.workspaceId, flushed); return; } await repo.insert({ diff --git a/apps/server/src/core/ai-chat/tools/ai-chat-tools.cap.spec.ts b/apps/server/src/core/ai-chat/tools/ai-chat-tools.cap.spec.ts new file mode 100644 index 00000000..7f504f48 --- /dev/null +++ b/apps/server/src/core/ai-chat/tools/ai-chat-tools.cap.spec.ts @@ -0,0 +1,160 @@ +import { + wrapInAppToolWithCap, + inAppToolCallCapMs, + type ToolAbortSignalSink, +} from './ai-chat-tools.service'; +import type { Tool, ToolCallOptions } from 'ai'; + +/** + * #487 commit 1 — in-app tool race-on-abort + safe-points + per-call cap. + * + * Tests assert the HONEST observable property the spec names — "after Stop, NO + * new HTTP/WS call STARTS; an already-started single call may take either + * outcome" — against the REAL wrapper mechanism (the composite abort signal it + * publishes on the client + the RACE it runs), NOT a timing-dependent proxy like + * "the write didn't land". + */ + +// A minimal stand-in for the client's `toolAbortSignal` field. In production the +// wrapper publishes the composite here and the client's paginateAll / +// mutatePageContent safe-points read it; the fake "tool" below reads it the same +// way, so this exercises the real contract without a live DB / collab socket. +class FakeClient implements ToolAbortSignalSink { + private signal: AbortSignal | null = null; + setToolAbortSignal(signal: AbortSignal | null): void { + this.signal = signal; + } + getToolAbortSignal(): AbortSignal | null { + return this.signal; + } +} + +// A ToolCallOptions with just the field the wrapper reads (abortSignal). The AI +// SDK passes a fuller object; the wrapper only spreads it and reads abortSignal. +const opts = (abortSignal?: AbortSignal): ToolCallOptions => + ({ toolCallId: 't1', messages: [], abortSignal }) as unknown as ToolCallOptions; + +const tick = (ms = 5) => new Promise((r) => setTimeout(r, ms)); + +describe('#487 wrapInAppToolWithCap — race-on-abort + safe-points', () => { + it('after Stop, no NEW simulated call starts (multi-call tool)', async () => { + const client = new FakeClient(); + const started: number[] = []; + // A multi-call tool that mirrors paginateAll: it consults the client signal + // at a safe-point BEFORE starting each simulated network call. + const multiCall: Tool = { + execute: (async (_args: unknown) => { + for (let i = 0; i < 6; i++) { + // Safe-point: exactly what paginateAll / mutatePageContent do. + client.getToolAbortSignal()?.throwIfAborted(); + started.push(i); + await tick(10); + } + return 'done'; + }) as unknown as Tool['execute'], + } as Tool; + + const wrapped = wrapInAppToolWithCap(multiCall, client, 10_000); + const ac = new AbortController(); + const call = ( + wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise + )({}, opts(ac.signal)); + + // Let one or two calls start, then Stop. + await tick(12); + ac.abort(new Error('user stop')); + + await expect(call).rejects.toThrow(); // wrapper rejects promptly + const startedAtStop = started.length; + + // Give the abandoned loser ample time; its next safe-point must throw because + // the (aborted) composite is still published on the client. + await tick(60); + expect(started.length).toBe(startedAtStop); + // It must NOT have run the whole sequence (that would mean Stop did nothing). + expect(started.length).toBeLessThan(6); + }); + + it('rejects immediately on Stop even if the call never settles (discard loser)', async () => { + const client = new FakeClient(); + let settled = false; + const hang: Tool = { + execute: (async () => { + await new Promise(() => undefined); // never resolves + settled = true; + }) as unknown as Tool['execute'], + } as Tool; + const wrapped = wrapInAppToolWithCap(hang, client, 10_000); + const ac = new AbortController(); + const call = ( + wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise + )({}, opts(ac.signal)); + await tick(5); + ac.abort(); + await expect(call).rejects.toThrow(); + expect(settled).toBe(false); + }); + + it('per-call cap rejects a hung call with no Stop signal', async () => { + const client = new FakeClient(); + const hang: Tool = { + execute: (async () => { + await new Promise(() => undefined); + }) as unknown as Tool['execute'], + } as Tool; + // Tiny cap; no options.abortSignal at all (composite = cap only). + const wrapped = wrapInAppToolWithCap(hang, client, 20); + const start = Date.now(); + await expect( + (wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise)( + {}, + opts(undefined), + ), + ).rejects.toThrow(/per-call cap/); + expect(Date.now() - start).toBeLessThan(2000); + }); + + it('publishes a composite signal on the client for the duration of the call', async () => { + const client = new FakeClient(); + let seenDuringCall: AbortSignal | null = null; + const probe: Tool = { + execute: (async () => { + seenDuringCall = client.getToolAbortSignal(); + return 'ok'; + }) as unknown as Tool['execute'], + } as Tool; + const wrapped = wrapInAppToolWithCap(probe, client, 10_000); + const ac = new AbortController(); + await ( + wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise + )({}, opts(ac.signal)); + expect(seenDuringCall).not.toBeNull(); + // The published composite must reflect the turn's Stop signal. + ac.abort(); + expect((seenDuringCall as unknown as AbortSignal).aborted).toBe(true); + }); + + it('a completed call returns its raw result unchanged', async () => { + const client = new FakeClient(); + const ok: Tool = { + execute: (async () => ({ items: [1, 2, 3] })) as unknown as Tool['execute'], + } as Tool; + const wrapped = wrapInAppToolWithCap(ok, client, 10_000); + const res = await ( + wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise + )({}, opts(new AbortController().signal)); + expect(res).toEqual({ items: [1, 2, 3] }); + }); + + it('cap is env-tunable with a 2-minute default', () => { + const prev = process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS; + delete process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS; + expect(inAppToolCallCapMs()).toBe(120_000); + process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = '5000'; + expect(inAppToolCallCapMs()).toBe(5000); + process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = 'not-a-number'; + expect(inAppToolCallCapMs()).toBe(120_000); + if (prev === undefined) delete process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS; + else process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = prev; + }); +}); diff --git a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts index dc2270fd..58561ea2 100644 --- a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts +++ b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts @@ -1,5 +1,5 @@ import { Injectable, Logger } from '@nestjs/common'; -import { tool, type Tool } from 'ai'; +import { tool, type Tool, type ToolCallOptions } from 'ai'; import { z } from 'zod'; import { User } from '@docmost/db/types/entity.types'; import { TokenService } from '../../auth/services/token.service'; @@ -159,6 +159,129 @@ function __assertClientCallContract(client: DocmostClientLike): void { * existing service-account `/mcp` path already calls loopback successfully, so * this works for single-workspace self-host. */ +/** + * #487: wall-clock cap for a SINGLE in-app tool call, env-tunable via + * `AI_CHAT_INAPP_TOOL_CALL_CAP_MS`. Bounds a read tool that would otherwise + * paginate for minutes and a content write whose collab commit hangs, and is the + * per-call CAP half of the composite abort signal every in-app tool is wrapped + * with (the other half is the turn's Stop signal). Default 2 minutes: generous + * for a legitimate long read/write, tight enough that a stuck call cannot pin the + * turn. The reconcile staleness floor (#487 commit 4) is derived as + * max(2 x this cap, 15 min), so keep this well under that. + */ +export function inAppToolCallCapMs(): number { + const raw = Number(process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS); + return Number.isFinite(raw) && raw > 0 ? raw : 120_000; +} + +/** #487: the composite signal's reason as an Error (informative thrown value). */ +function inAppAbortReason(signal: AbortSignal): Error { + const r = signal.reason; + return r instanceof Error + ? r + : new Error(typeof r === 'string' ? r : 'In-app tool call aborted'); +} + +/** + * The client surface {@link wrapInAppToolWithCap} drives (#487). Both methods are + * OPTIONAL: the real loopback DocmostClient implements them (so a Stop/cap reaches + * its pagination / pre-commit safe-points), but a client that omits them still + * gets the OUTER guarantee — the race rejects on abort regardless. This keeps the + * wrapper decoupled from the exact client shape (unit-test doubles need not stub + * the plumbing). + */ +export interface ToolAbortSignalSink { + setToolAbortSignal?(signal: AbortSignal | null): void; + getToolAbortSignal?(): AbortSignal | null; +} + +/** + * #487: wrap an in-app tool so a Stop (the turn's `options.abortSignal`) OR the + * per-call wall-clock cap REJECTS the call immediately, and so that SAME + * composite signal reaches the client's pagination / pre-commit safe-points (via + * `client.setToolAbortSignal`) — making a Stop stop the NEXT HTTP/WS call from + * starting. + * + * Reuses the RACE pattern of `wrapToolWithCallTimeout` (mcp-clients.service.ts): + * the call is raced against the composite signal, so on abort we reject in the + * SAME tick and DISCARD the loser promise. Its network / collab teardown latency + * therefore never blocks the turn — the supersede timeout W=10s (#487 commit 3) + * relies on this abort->settle latency being milliseconds, not a socket teardown. + * Awaiting the client's own signal-into-write path alone would NOT satisfy this + * (the loser could still be tearing down a collab socket). + * + * The composite is SET on the client at entry and deliberately NOT restored on + * unwind: after this wrapper rejects on abort, the ABANDONED loser promise keeps + * running, and its safe-points read the client field — leaving the (aborted) + * composite there is exactly what makes the loser's NEXT call throw and stop. The + * next in-app tool call overwrites the field with its own fresh composite before + * any of its safe-points run, so a stale settled signal never leaks forward. + * SINGLE-WRITER by phase-1 assumption — see DocmostClientContext.toolAbortSignal + * for the parallel-call caveat (#487). + * + * KNOWN LIMITATION (#487): a write tool that issues SEVERAL sequential collab + * commits can be aborted BETWEEN commits, leaving a partially-applied operation. + * Cancel guarantees "no NEW call starts", NOT "the write didn't land". + */ +export function wrapInAppToolWithCap( + toolDef: Tool, + client: ToolAbortSignalSink, + capMs: number, +): Tool { + const original = toolDef.execute; + if (typeof original !== 'function') return toolDef; + const execute = async (args: unknown, options: ToolCallOptions) => { + const capController = new AbortController(); + const timer = setTimeout(() => { + capController.abort( + new Error(`In-app tool call exceeded the ${capMs}ms per-call cap`), + ); + }, capMs); + timer.unref?.(); + const composite = options?.abortSignal + ? AbortSignal.any([options.abortSignal, capController.signal]) + : capController.signal; + // Reject the MOMENT the composite fires, independent of whether `original` + // ever settles (a hung collab write / read would otherwise pin the turn). The + // losing `original` is left pending; Promise.race attaches a rejection + // handler to both inputs so a late rejection is never unhandled. + const aborted = new Promise((_, reject) => { + const fail = () => reject(inAppAbortReason(composite)); + if (composite.aborted) fail(); + else composite.addEventListener('abort', fail, { once: true }); + }); + // Publish the composite so the client's pagination / pre-commit safe-points + // observe it (see the "not restored on unwind" rationale above). Guarded: a + // client without the plumbing still gets the OUTER race guarantee below. + client.setToolAbortSignal?.(composite); + try { + return await Promise.race([ + (original as (a: unknown, o: ToolCallOptions) => Promise)( + args, + { ...options, abortSignal: composite }, + ), + aborted, + ]); + } finally { + clearTimeout(timer); + } + }; + return { ...toolDef, execute } as unknown as Tool; +} + +/** #487: apply {@link wrapInAppToolWithCap} to every tool in a set. */ +export function wrapInAppToolsWithCap( + tools: Record, + client: ToolAbortSignalSink, + capMs: number, +): Record { + const out: Record = {}; + for (const [name, t] of Object.entries(tools)) { + out[name] = wrapInAppToolWithCap(t, client, capMs); + } + return out; +} + @Injectable() export class AiChatToolsService { private readonly logger = new Logger(AiChatToolsService.name); @@ -186,7 +309,12 @@ export class AiChatToolsService { sessionId: string, workspaceId: string, aiChatId: string, - ): Promise { + // #487: the returned client also carries the tool-cancellation plumbing + // (setToolAbortSignal/getToolAbortSignal). These are host plumbing, NOT part + // of the tool-execute surface (DocmostClientMethod), so they are surfaced here + // as an intersection rather than by widening that Pick — keeping the + // positional-call drift-guard (#446) scoped to the actual tool methods. + ): Promise { const apiUrl = process.env.MCP_DOCMOST_API_URL || `http://127.0.0.1:${process.env.PORT || 3000}/api`; @@ -630,7 +758,15 @@ export class AiChatToolsService { // dependency and reuses the CASL enforcement already on `client`. When the // loaded package predates #417 (factory undefined) or the loader is mocked in // a unit test, signalling is a pure no-op and results are byte-identical. - if (!createCommentSignalTracker) return tools; + // #487: wrap every in-app tool with the race-on-abort + per-call cap guard so + // a Stop / cap rejects immediately AND reaches the client's write/pagination + // safe-points. Applied as the OUTERMOST wrapper (over the comment-signal + // wrapper below) so the race governs the whole call. The client carries the + // per-call composite signal via setToolAbortSignal. + const capMs = inAppToolCallCapMs(); + if (!createCommentSignalTracker) { + return wrapInAppToolsWithCap(tools, client, capMs); + } const tracker = createCommentSignalTracker({ probe: async (pageId: string, sinceMs: number) => { @@ -659,7 +795,11 @@ export class AiChatToolsService { }, }); - return wrapToolsWithCommentSignal(tools, tracker); + return wrapInAppToolsWithCap( + wrapToolsWithCommentSignal(tools, tracker), + client, + capMs, + ); } } diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.ts b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.ts index 78cc7064..96070e77 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat-message.repo.ts @@ -1,5 +1,6 @@ import { Injectable, Logger } from '@nestjs/common'; import { InjectKysely } from 'nestjs-kysely'; +import { sql } from 'kysely'; import { KyselyDB, KyselyTransaction } from '../../types/kysely.types'; import { dbOrTx } from '../../utils'; import { @@ -188,6 +189,144 @@ export class AiChatMessageRepo { return query.returning(this.baseFields).executeTakeFirst(); } + /** + * #487 OWNER terminal write — the streamText terminal callback's finalize. Like + * `update` but CONDITIONAL on `status='streaming' OR metadata.finalizeFailed`: + * the owner writes its real content EITHER when the row is still streaming (the + * normal case) OR when a reconcile stamp already flipped it to a terminal status + * but marked `finalizeFailed:true` — the owner's real content OVERWRITES that + * placeholder stamp (owner-write priority, #487). A row that is properly terminal + * (no finalizeFailed) is left untouched (undefined) — idempotent. The `patch` + * carries the real metadata WITHOUT finalizeFailed, so a successful write CLEARS + * the flag. Returns the updated row, or undefined when nothing matched. + */ + async finalizeOwner( + id: string, + workspaceId: string, + patch: Partial<{ + content: string | null; + toolCalls: unknown; + metadata: unknown; + status: string | null; + }>, + trx?: KyselyTransaction, + ): Promise { + const db = dbOrTx(this.db, trx); + return db + .updateTable('aiChatMessages') + .set({ ...(patch as Record), updatedAt: new Date() }) + .where('id', '=', id) + .where('workspaceId', '=', workspaceId) + .where((eb) => + eb.or([ + eb('status', '=', 'streaming'), + eb(sql`(metadata->>'finalizeFailed')`, '=', 'true'), + ]), + ) + .returning(this.baseFields) + .executeTakeFirst(); + } + + /** + * #487 RECONCILE status-only stamp — settle a stuck 'streaming' row to a + * terminal status WITHOUT the owner's real content (which lived only in the + * dead process's memory — a documented loss). CONDITIONAL on `status='streaming'` + * (never touches an already-terminal row) AND it MERGES `finalizeFailed:true` + * into metadata (preserving the partial `parts` already persisted) so a LATER + * owner-write (finalizeOwner) can still OVERWRITE this placeholder with real + * content, and so `isInterruptResume` can EXCLUDE this row (a reconcile stamp is + * not a genuine user interruption). Returns the updated row, or undefined. + */ + async stampTerminalIfStreaming( + id: string, + workspaceId: string, + status: 'aborted' | 'error' | 'completed', + trx?: KyselyTransaction, + ): Promise { + const db = dbOrTx(this.db, trx); + return db + .updateTable('aiChatMessages') + .set({ + status, + metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`, + updatedAt: new Date(), + }) + .where('id', '=', id) + .where('workspaceId', '=', workspaceId) + .where('status', '=', 'streaming') + .returning(this.baseFields) + .executeTakeFirst(); + } + + /** + * #487 reconcile clause (b): streaming assistant rows whose linked RUN has + * already reached a terminal status — an asymmetry ("run settled / message + * streaming forever") the periodic reconcile heals by stamping the message. + * Returns the message id + its run's terminal status, bounded. + */ + async findStreamingWithTerminalRun( + limit = 200, + // #487: scope to ONE chat for the opportunistic per-turn reconcile (removes + // reconcile latency from the user-visible path); omit for the periodic sweep. + chat?: { chatId: string; workspaceId: string }, + ): Promise< + Array<{ messageId: string; workspaceId: string; runStatus: string }> + > { + let query = this.db + .selectFrom('aiChatMessages as m') + .innerJoin('aiChatRuns as r', 'r.assistantMessageId', 'm.id') + .select([ + 'm.id as messageId', + 'm.workspaceId as workspaceId', + 'r.status as runStatus', + ]) + .where('m.status', '=', 'streaming') + .where('r.status', 'in', ['succeeded', 'failed', 'aborted']); + if (chat) { + query = query + .where('m.chatId', '=', chat.chatId) + .where('m.workspaceId', '=', chat.workspaceId); + } + return query.limit(limit).execute(); + } + + /** + * #487 reconcile clause (d) — historical-row safety: streaming rows older than + * `staleMs` whose chat has NO active run row (double-gated). Settle them to + * 'aborted' + finalizeFailed (so a late owner-write could still overwrite). + * Returns the count. Used ONLY by the periodic reconcile, never at boot. + */ + async sweepStreamingWithoutActiveRun( + staleMs: number, + trx?: KyselyTransaction, + ): Promise { + const db = dbOrTx(this.db, trx); + const staleBefore = new Date(Date.now() - staleMs); + const rows = await db + .updateTable('aiChatMessages as m') + .set({ + status: 'aborted', + metadata: sql`coalesce(m.metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`, + updatedAt: new Date(), + }) + .where('m.status', '=', 'streaming') + .where('m.updatedAt', '<', staleBefore) + .where((eb) => + eb.not( + eb.exists( + eb + .selectFrom('aiChatRuns as r') + .select('r.id') + .whereRef('r.chatId', '=', 'm.chatId') + .where('r.status', 'in', ['pending', 'running']), + ), + ), + ) + .returning('m.id') + .execute(); + return rows.length; + } + /** * Crash-recovery sweep (#183): flip every assistant row still left in the * 'streaming' state (a turn that died mid-write before reaching a terminal @@ -200,13 +339,20 @@ export class AiChatMessageRepo { * step, so an actively-streaming row never matches; this prevents a fresh * replica's boot-sweep from aborting a turn another replica is still streaming * in a multi-instance deploy. + * + * #487: the sweep now ALSO marks `finalizeFailed:true` so a late owner-write can + * overwrite this placeholder with real content (owner-write priority). */ async sweepStreaming(trx?: KyselyTransaction): Promise { const db = dbOrTx(this.db, trx); const staleBefore = new Date(Date.now() - SWEEP_STREAMING_STALE_MS); const rows = await db .updateTable('aiChatMessages') - .set({ status: 'aborted', updatedAt: new Date() }) + .set({ + status: 'aborted', + metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`, + updatedAt: new Date(), + }) .where('status', '=', 'streaming') .where('updatedAt', '<', staleBefore) .returning('id') diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.ts b/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.ts index 7bb6bcdb..c27480bf 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat-run.repo.ts @@ -143,6 +143,41 @@ export class AiChatRunRepo { .executeTakeFirst(); } + /** + * #487: CONDITIONAL terminal finalize — flip a run to a terminal status and + * stamp `finished_at` ONLY while it is still active (pending|running), mirroring + * the assistant message's `onlyIfStreaming` guard. A double-settle (a late or + * second writer, a supersede applying a zombie's intended, a reconcile stamp) + * matches NOTHING once the row is terminal and is a benign no-op — so a terminal + * status can never be clobbered by a later writer (last-writer-wins is gone). + * + * Returns the updated row when it WAS active (this call wrote it), else + * undefined (the row was already terminal — another writer won). The caller + * distinguishes the two to resolve the correct settle outcome. + */ + async finalizeIfActive( + id: string, + workspaceId: string, + patch: { status: string; error: string | null }, + trx?: KyselyTransaction, + ): Promise { + const db = dbOrTx(this.db, trx); + const now = new Date(); + return db + .updateTable('aiChatRuns') + .set({ + status: patch.status, + error: patch.error, + finishedAt: now, + updatedAt: now, + }) + .where('id', '=', id) + .where('workspaceId', '=', workspaceId) + .where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[]) + .returning(this.baseFields) + .executeTakeFirst(); + } + /** * Mark an EXPLICIT stop request on an active run (distinct from a browser * disconnect, which never stops a run). Stamps `stop_requested_at` ONLY while @@ -184,6 +219,31 @@ export class AiChatRunRepo { * sweeps only runs UNTOUCHED past the window. Phase 1 is single-process, so the * boot path supplies no window. */ + /** + * #487 reconcile clause (c): active (pending|running) runs UNTOUCHED past + * `staleMs` — candidates for "no live runner" abort. Staleness is measured from + * `updated_at` (the LAST-PROGRESS timestamp — recordStep bumps it), NOT + * `started_at`, so a legitimate long-running marathon (11–25 min of steady + * progress) is never a candidate. The caller filters these against its in-memory + * `active` / zombie maps ("no entry" is the PRIMARY gate — a live entry is never + * aborted) before settling any of them. Bounded. + */ + async findStaleActive( + staleMs: number, + limit = 200, + trx?: KyselyTransaction, + ): Promise> { + const db = dbOrTx(this.db, trx); + const staleBefore = new Date(Date.now() - staleMs); + return db + .selectFrom('aiChatRuns') + .select(['id', 'workspaceId', 'chatId']) + .where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[]) + .where('updatedAt', '<', staleBefore) + .limit(limit) + .execute(); + } + async sweepRunning( opts: { staleMs?: number } = {}, trx?: KyselyTransaction, diff --git a/apps/server/test/integration/ai-chat-reconcile.int-spec.ts b/apps/server/test/integration/ai-chat-reconcile.int-spec.ts new file mode 100644 index 00000000..4a38f1c2 --- /dev/null +++ b/apps/server/test/integration/ai-chat-reconcile.int-spec.ts @@ -0,0 +1,305 @@ +import { Kysely } from 'kysely'; +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 { AiChatRunService } from '../../src/core/ai-chat/ai-chat-run.service'; +import { + getTestDb, + destroyTestDb, + createWorkspace, + createUser, + createChat, + createMessage, +} from './db'; + +/** + * #487 commit 4 — bidirectional reconcile + owner-write priority, real SQL. + * + * Proves the OBSERVABLE recovery properties against docmost_test: + * - the CONDITIONAL owner-write beats a reconcile stamp, and a stamp never + * clobbers a proper terminal row; + * - a LATE owner-finalize with real content OVERWRITES a reconcile 'aborted' + * stamp (finalizeFailed); + * - each reconcile clause (b message<-run, c stale-run, d historical row) settles + * the stuck row/run, and a LIVE run entry is never touched; + * - the "kill DB on finish" recovery: after the DB comes back, neither the + * message row nor the run row stays stuck. + */ +describe('#487 reconcile + owner-write priority [integration]', () => { + let db: Kysely; + let messageRepo: AiChatMessageRepo; + let runRepo: AiChatRunRepo; + let runService: AiChatRunService; + let workspaceId: string; + let userId: string; + + beforeAll(async () => { + db = getTestDb(); + messageRepo = new AiChatMessageRepo(db as any); + runRepo = new AiChatRunRepo(db as any); + runService = new AiChatRunService(runRepo, { isCloud: () => false } as never); + workspaceId = (await createWorkspace(db)).id; + userId = (await createUser(db, workspaceId)).id; + }); + + afterAll(async () => { + await destroyTestDb(); + }); + + const newChat = async () => + (await createChat(db, { workspaceId, creatorId: userId })).id; + + const metaOf = async (id: string): Promise | null> => { + const row = await messageRepo.findById(id, workspaceId); + return (row?.metadata as Record | null) ?? null; + }; + + it('owner finalizeOwner writes a streaming row and CLEARS finalizeFailed', async () => { + const chatId = await newChat(); + const m = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + status: 'streaming', + metadata: { parts: [] }, + }); + const wrote = await messageRepo.finalizeOwner(m.id, workspaceId, { + content: 'final answer', + status: 'completed', + metadata: { parts: [{ type: 'text', text: 'final answer' }] }, + } as never); + expect(wrote!.status).toBe('completed'); + expect((await metaOf(m.id))?.finalizeFailed).toBeUndefined(); + }); + + it('a reconcile stamp NEVER clobbers a proper terminal row (finalizeOwner is a no-op there)', async () => { + const chatId = await newChat(); + const m = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + status: 'completed', + content: 'real', + metadata: { parts: [] }, + }); + // The reconcile stamp is onlyIfStreaming -> no-op on a completed row. + const stamped = await messageRepo.stampTerminalIfStreaming( + m.id, + workspaceId, + 'aborted', + ); + expect(stamped).toBeUndefined(); + expect((await messageRepo.findById(m.id, workspaceId))!.status).toBe( + 'completed', + ); + }); + + it('LATE owner-finalize with real content OVERWRITES a reconcile aborted stamp', async () => { + const chatId = await newChat(); + const m = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + status: 'streaming', + metadata: { parts: [{ type: 'text', text: 'partial' }] }, + }); + // Reconcile stamps it aborted + finalizeFailed (final text lived only in mem). + const stamped = await messageRepo.stampTerminalIfStreaming( + m.id, + workspaceId, + 'aborted', + ); + expect(stamped!.status).toBe('aborted'); + expect((await metaOf(m.id))?.finalizeFailed).toBe(true); + + // A LATE owner-write (finalizeFailed=true satisfies the OR) overwrites it with + // real content, clearing the flag — owner-write priority. + const wrote = await messageRepo.finalizeOwner(m.id, workspaceId, { + content: 'the real final answer', + status: 'completed', + metadata: { parts: [{ type: 'text', text: 'the real final answer' }] }, + } as never); + expect(wrote!.status).toBe('completed'); + expect(wrote!.content).toBe('the real final answer'); + expect((await metaOf(m.id))?.finalizeFailed).toBeUndefined(); + }); + + it('clause (c): a stale active run with NO live entry -> aborted; a LIVE entry is untouched', async () => { + // Stale run, NOT owned by this replica (no entry) -> reconcile aborts it. + const staleChat = await newChat(); + const stale = await runRepo.insert({ + chatId: staleChat, + workspaceId, + createdBy: userId, + status: 'running', + }); + await db + .updateTable('aiChatRuns') + .set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) }) + .where('id', '=', stale.id) + .execute(); + + // A live run OWNED by this replica (beginRun registers an in-memory entry), + // ALSO backdated stale — the "no entry" primary gate must protect it. + const liveChat = await newChat(); + const live = await runService.beginRun({ + chatId: liveChat, + workspaceId, + userId, + }); + await db + .updateTable('aiChatRuns') + .set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) }) + .where('id', '=', live.runId) + .execute(); + + const aborted = await runService.reconcileStaleRuns(15 * 60 * 1000); + expect(aborted).toBeGreaterThanOrEqual(1); + expect((await runRepo.findById(stale.id, workspaceId))!.status).toBe( + 'aborted', + ); + // The live entry is NEVER aborted, however stale its row looks. + expect((await runRepo.findById(live.runId, workspaceId))!.status).toBe( + 'running', + ); + expect(runService.isLocallyActive(live.runId)).toBe(true); + + // cleanup the live run + await runService.finalizeRun(live.runId, workspaceId, 'aborted'); + }); + + it('clause (b): a streaming message whose RUN is terminal is stamped by run status (succeeded -> aborted, NOT completed-empty)', async () => { + const chatId = await newChat(); + const msg = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + status: 'streaming', + metadata: { parts: [] }, + }); + // A SUCCEEDED run linked to the still-streaming message (the asymmetry). + const run = await runRepo.insert({ + chatId, + workspaceId, + createdBy: userId, + status: 'running', + assistantMessageId: msg.id, + }); + await runRepo.finalizeIfActive(run.id, workspaceId, { + status: 'succeeded', + error: null, + }); + + const stuck = await messageRepo.findStreamingWithTerminalRun(); + const mine = stuck.find((s) => s.messageId === msg.id); + expect(mine?.runStatus).toBe('succeeded'); + // Reconcile clause (b): succeeded run -> message 'aborted' (NOT 'completed'), + // the final text lived only in memory (documented loss), +finalizeFailed. + const status = mine!.runStatus === 'failed' ? 'error' : 'aborted'; + await messageRepo.stampTerminalIfStreaming(msg.id, workspaceId, status); + const row = await messageRepo.findById(msg.id, workspaceId); + expect(row!.status).toBe('aborted'); + expect((row!.metadata as Record).finalizeFailed).toBe(true); + }); + + it('clause (d): a stale streaming row with NO active run on the chat -> aborted+finalizeFailed', async () => { + const chatId = await newChat(); + const msg = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + status: 'streaming', + metadata: { parts: [] }, + }); + await db + .updateTable('aiChatMessages') + .set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) }) + .where('id', '=', msg.id) + .execute(); + + const swept = await messageRepo.sweepStreamingWithoutActiveRun( + 15 * 60 * 1000, + ); + expect(swept).toBeGreaterThanOrEqual(1); + const row = await messageRepo.findById(msg.id, workspaceId); + expect(row!.status).toBe('aborted'); + expect((row!.metadata as Record).finalizeFailed).toBe(true); + }); + + it('clause (d) is DOUBLE-GATED: a stale streaming row WITH an active run on the chat is left alone', async () => { + const chatId = await newChat(); + const msg = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + status: 'streaming', + metadata: { parts: [] }, + }); + await db + .updateTable('aiChatMessages') + .set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) }) + .where('id', '=', msg.id) + .execute(); + // An ACTIVE run on the same chat -> clause (d) must NOT touch the message. + const run = await runRepo.insert({ + chatId, + workspaceId, + createdBy: userId, + status: 'running', + }); + + await messageRepo.sweepStreamingWithoutActiveRun(15 * 60 * 1000); + expect((await messageRepo.findById(msg.id, workspaceId))!.status).toBe( + 'streaming', + ); + await runRepo.finalizeIfActive(run.id, workspaceId, { + status: 'aborted', + error: null, + }); + }); + + it('"kill DB on finish" recovery: after the DB is back, reconcile leaves NEITHER the row nor the run stuck', async () => { + // Simulate a process that seeded the assistant row + run, then died before + // finalizing EITHER (a mid-turn crash): a streaming message + a running run, + // both stale, with no in-memory entry (fresh service = fresh maps). + const chatId = await newChat(); + const msg = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + status: 'streaming', + metadata: { parts: [{ type: 'text', text: 'partial' }] }, + }); + const run = await runRepo.insert({ + chatId, + workspaceId, + createdBy: userId, + status: 'running', + assistantMessageId: msg.id, + }); + await db + .updateTable('aiChatRuns') + .set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) }) + .where('id', '=', run.id) + .execute(); + await db + .updateTable('aiChatMessages') + .set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) }) + .where('id', '=', msg.id) + .execute(); + + // Reconcile (as the periodic job would): (c) aborts the orphan run, then + // (b) settles the message from the now-terminal run. + await runService.reconcileStaleRuns(15 * 60 * 1000); + const stuck = await messageRepo.findStreamingWithTerminalRun(); + for (const s of stuck) { + const status = s.runStatus === 'failed' ? 'error' : 'aborted'; + await messageRepo.stampTerminalIfStreaming(s.messageId, s.workspaceId, status); + } + + // Neither is stuck: the run is terminal AND the message is terminal. + expect((await runRepo.findById(run.id, workspaceId))!.status).toBe('aborted'); + const row = await messageRepo.findById(msg.id, workspaceId); + expect(row!.status).toBe('aborted'); + expect((row!.metadata as Record).finalizeFailed).toBe(true); + }); +}); diff --git a/apps/server/test/integration/ai-chat-run.int-spec.ts b/apps/server/test/integration/ai-chat-run.int-spec.ts index f6a753a1..aca31c03 100644 --- a/apps/server/test/integration/ai-chat-run.int-spec.ts +++ b/apps/server/test/integration/ai-chat-run.int-spec.ts @@ -281,6 +281,52 @@ describe('AiChatRun durable lifecycle [integration]', () => { }); }); + it('#487 finalizeIfActive is CONDITIONAL: a late terminal write cannot clobber the settled status (real SQL)', async () => { + const c = (await createChat(db, { workspaceId, creatorId: userId })).id; + const run = await runRepo.insert({ + chatId: c, + workspaceId, + createdBy: userId, + status: 'running', + }); + + // First terminal write: the run IS active, so it flips + returns the row. + const first = await runRepo.finalizeIfActive(run.id, workspaceId, { + status: 'succeeded', + error: null, + }); + expect(first!.status).toBe('succeeded'); + expect(first!.finishedAt).toBeTruthy(); + + // A late/second writer tries to flip it to 'aborted' — the WHERE status IN + // ('pending','running') guard matches NOTHING now, so it is a benign no-op. + const second = await runRepo.finalizeIfActive(run.id, workspaceId, { + status: 'aborted', + error: 'late clobber attempt', + }); + expect(second).toBeUndefined(); + + // The persisted terminal status is UNCHANGED — last-writer-wins is gone. + const row = await runRepo.findById(run.id, workspaceId); + expect(row!.status).toBe('succeeded'); + expect(row!.error).toBeNull(); + }); + + it('#487 double-settle through the service collapses to one write at the SQL gate', async () => { + const c = (await createChat(db, { workspaceId, creatorId: userId })).id; + const handle = await service.beginRun({ chatId: c, workspaceId, userId }); + + // First settle writes 'aborted' via the conditional write. + await service.finalizeRun(handle.runId, workspaceId, 'aborted'); + // A late safety-net settle to 'error' is a no-op (row already terminal). + await service.finalizeRun(handle.runId, workspaceId, 'error', 'late'); + + const row = await runRepo.findById(handle.runId, workspaceId); + expect(row!.status).toBe('aborted'); + expect(service.isLocallyActive(handle.runId)).toBe(false); + expect(service.hasZombie(handle.runId)).toBe(false); + }); + it('sweepRunning() with NO args (boot sweep / variant C) aborts even a FRESH running run', async () => { // F1/DECISION C at the SQL level: the unconditional boot sweep has NO // staleness window, so a run updated just now (a fast restart) is settled too diff --git a/packages/mcp/src/client/context.ts b/packages/mcp/src/client/context.ts index 12a8b1a0..8c784f53 100644 --- a/packages/mcp/src/client/context.ts +++ b/packages/mcp/src/client/context.ts @@ -170,6 +170,43 @@ export abstract class DocmostClientContext { // cached conversion can never leak across identities. See getpage-cache.ts. protected getPageCache = new GetPageConversionCache(); + // #487: an OPTIONAL abort signal the in-app tool host sets before each tool + // call (a composite of the turn's Stop signal + a per-call wall-clock cap). It + // is checked at safe-points BETWEEN the sequential HTTP calls of a paginated + // read (paginateAll) and just before the atomic collab commit of a write (the + // mutatePage/replacePage/mutateLiveContentUnlocked seams), so a Stop / cap + // stops the NEXT network call from STARTING. An already-started single call may + // still land — a documented limitation (#487). + // + // SINGLE-WRITER by phase-1 assumption: exactly one DocmostClient is built per + // turn and shared by every tool call; the host sets this per call and does NOT + // restore the prior value on unwind (set-and-leave) — a fresh client per turn + // plus overwrite-by-the-next-call keeps it correct, and leaving a settled + // call's signal in place is what makes a discarded race-loser throw on its + // next safe-point. If the model emits PARALLEL in-app + // tool calls they share this one field, so the per-call CAP of one call is not + // guaranteed to bound another's in-flight pagination — but every composite the + // host sets carries the SAME turn Stop signal, so a Stop still aborts whichever + // signal is current. #487. + protected toolAbortSignal: AbortSignal | null = null; + + /** + * #487: set (or clear with null) the in-app tool abort signal governing the + * NEXT client call's safe-points. The host wraps each in-app tool call: it sets + * the composite (Stop + per-call cap) here before invoking the tool and leaves + * it in place afterwards (set-and-leave, NOT restored) — the next call + * overwrites it, and a fresh client is built per turn. Public so the + * server-side tool wrapper can reach it; harmless (a no-op) when never set. + */ + public setToolAbortSignal(signal: AbortSignal | null): void { + this.toolAbortSignal = signal; + } + + /** #487: the abort signal currently governing this client's safe-points. */ + public getToolAbortSignal(): AbortSignal | null { + return this.toolAbortSignal; + } + // Two construction forms: // - new DocmostClient(config) // discriminated union (current) // - new DocmostClient(baseURL, email, password) // legacy positional creds @@ -571,6 +608,10 @@ export abstract class DocmostClientContext { this.onMetricFn?.("collab_connect_timeouts_total", 1), }); try { + // #487 PRE-COMMIT safe-point (reentrant twin of mutatePageContent): a + // Stop/cap after acquiring the session but before the atomic write skips + // this commit. Same limitation applies (stops the NEXT commit only). + this.toolAbortSignal?.throwIfAborted(); return await session.mutate(transform); } catch (e) { // Drop the session on any failure so the next call reconnects fresh. @@ -602,6 +643,11 @@ export abstract class DocmostClientContext { let truncated = false; for (let page = 0; page < MAX_PAGES; page++) { + // #487 safe-point: a Stop (or the in-app tool per-call cap) that fires + // BETWEEN sequential page fetches must stop the NEXT request from starting + // — a read tool that would otherwise paginate for minutes is interrupted + // here. throwIfAborted() rejects with the signal's reason. + this.toolAbortSignal?.throwIfAborted(); const payload: Record = { ...basePayload, limit: clampedLimit, @@ -709,7 +755,15 @@ export abstract class DocmostClientContext { // #486: on a rejected collab-WS handshake, invalidate + refresh the token and // retry the write once (symmetric to the HTTP-401 reauth path). return this.writeWithCollabAuthRetry(collabToken, (token) => - mutatePageContent(pageUuid, token, apiUrl, transform), + // #487: thread the in-app tool signal to mutatePageContent's pre-commit + // safe-point so a Stop/cap during the connect/lock window skips the write. + mutatePageContent( + pageUuid, + token, + apiUrl, + transform, + this.toolAbortSignal ?? undefined, + ), ); } @@ -733,7 +787,14 @@ export abstract class DocmostClientContext { // #486: on a rejected collab-WS handshake, invalidate + refresh the token and // retry the write once (symmetric to the HTTP-401 reauth path). return this.writeWithCollabAuthRetry(collabToken, (token) => - replacePageContent(pageUuid, doc, token, apiUrl), + // #487: same pre-commit safe-point as mutatePage, for full-document writes. + replacePageContent( + pageUuid, + doc, + token, + apiUrl, + this.toolAbortSignal ?? undefined, + ), ); } diff --git a/packages/mcp/src/lib/collaboration.ts b/packages/mcp/src/lib/collaboration.ts index c7a77c5d..66d1d9ce 100644 --- a/packages/mcp/src/lib/collaboration.ts +++ b/packages/mcp/src/lib/collaboration.ts @@ -254,6 +254,12 @@ export async function mutatePageContent( collabToken: string, baseUrl: string, transform: (liveDoc: any) => any | null, + // #487: optional abort signal carrying the turn's Stop + the in-app tool + // per-call cap. Checked as the PRE-COMMIT safe-point below (after the session + // is acquired, immediately before the atomic read->write), so a Stop that + // arrives during the connect/lock window stops THIS write from landing. See the + // limitation note at the check. + signal?: AbortSignal, ): Promise { return withPageLock(pageId, async () => { if (process.env.DEBUG) { @@ -266,6 +272,13 @@ export async function mutatePageContent( const session = await acquireCollabSession(pageId, collabToken, baseUrl); try { + // #487 PRE-COMMIT safe-point: if the turn was Stopped (or the in-app tool + // per-call cap fired) after we acquired the collab session but before the + // atomic write, throw NOW so this commit never runs. KNOWN LIMITATION + // (#487): this only stops THIS commit — a write tool that already committed + // an EARLIER call this turn leaves that op applied. Cancel guarantees "no + // NEW commit starts", NOT "the write didn't land". + signal?.throwIfAborted(); return await session.mutate(transform); } catch (e) { // Drop the session on any failure so the next call reconnects fresh (this @@ -291,6 +304,8 @@ export async function replacePageContent( prosemirrorDoc: any, collabToken: string, baseUrl: string, + // #487: threaded straight to mutatePageContent's pre-commit safe-point. + signal?: AbortSignal, ): Promise { // Fail fast on a bad document instead of deferring the failure into the // collaboration write (where TiptapTransformer.toYdoc(undefined) used to @@ -307,6 +322,7 @@ export async function replacePageContent( collabToken, baseUrl, () => prosemirrorDoc, + signal, ); } diff --git a/packages/mcp/test/mock/paginate-abort-safepoint.test.mjs b/packages/mcp/test/mock/paginate-abort-safepoint.test.mjs new file mode 100644 index 00000000..af4c4e16 --- /dev/null +++ b/packages/mcp/test/mock/paginate-abort-safepoint.test.mjs @@ -0,0 +1,143 @@ +// #487 commit 1 — the in-app tool cancellation safe-point inside paginateAll. +// +// The in-app tool host sets a composite abort signal on the client +// (setToolAbortSignal) before each tool call; paginateAll checks it at a +// safe-point BEFORE every sequential page fetch, so a Stop that lands mid-read +// stops the NEXT HTTP request from STARTING (a read tool can no longer paginate +// for minutes past a Stop). This pins the HONEST observable property against the +// REAL client + a real HTTP server: "after Stop, no NEW request starts". +import { test, after } from "node:test"; +import assert from "node:assert/strict"; +import http from "node:http"; +import { DocmostClient } from "../../build/client.js"; + +function readBody(req) { + return new Promise((resolve) => { + let raw = ""; + req.on("data", (c) => (raw += c)); + req.on("end", () => resolve(raw)); + }); +} +function sendJson(res, status, obj, extra = {}) { + res.writeHead(status, { "Content-Type": "application/json", ...extra }); + res.end(JSON.stringify(obj)); +} +const openServers = []; +async function spawn(handler) { + const server = await new Promise((resolve) => { + const s = http.createServer(handler); + s.listen(0, "127.0.0.1", () => resolve(s)); + }); + openServers.push(server); + const { port } = server.address(); + return { baseURL: `http://127.0.0.1:${port}/api` }; +} +after(async () => { + await Promise.all(openServers.map((s) => new Promise((r) => s.close(r)))); +}); +function handleLogin(req, res) { + if (req.url === "/api/auth/login") { + sendJson(res, 200, { success: true }, { + "Set-Cookie": "authToken=t; Path=/; HttpOnly", + }); + return true; + } + return false; +} + +// A Stop that lands DURING pagination: the server aborts the client signal as it +// serves page 1 (more pages remain). The loop's next safe-point must throw before +// the page-2 request is sent. +test("paginateAll stops the NEXT request when the signal aborts mid-pagination", async () => { + let requests = 0; + const ac = new AbortController(); + const { baseURL } = await spawn(async (req, res) => { + await readBody(req); + if (handleLogin(req, res)) return; + if (req.url === "/api/spaces") { + requests++; + // Simulate a user Stop that lands while page 1 is in flight. + if (requests === 1) ac.abort(new Error("user stop")); + sendJson(res, 200, { + success: true, + data: { + items: [{ id: `p${requests}` }], + meta: { hasNextPage: true, nextCursor: `c${requests}` }, + }, + }); + return; + } + sendJson(res, 404, {}); + }); + + const client = new DocmostClient(baseURL, "user@example.com", "pw"); + client.setToolAbortSignal(ac.signal); + + await assert.rejects( + () => client.paginateAll("/spaces", {}), + /user stop/, + "the aborted safe-point rejects with the signal's reason", + ); + assert.equal(requests, 1, "page 2 never started after the Stop"); +}); + +// A Stop that is already in effect before the read starts: zero requests fire. +test("paginateAll starts no request when the signal is already aborted", async () => { + let requests = 0; + const { baseURL } = await spawn(async (req, res) => { + await readBody(req); + if (handleLogin(req, res)) return; + if (req.url === "/api/spaces") { + requests++; + sendJson(res, 200, { + success: true, + data: { items: [], meta: { hasNextPage: false, nextCursor: null } }, + }); + return; + } + sendJson(res, 404, {}); + }); + + const client = new DocmostClient(baseURL, "user@example.com", "pw"); + // Warm the auth so ensureAuthenticated does not itself POST after the abort. + await client.ensureAuthenticated(); + const ac = new AbortController(); + ac.abort(new Error("already stopped")); + client.setToolAbortSignal(ac.signal); + + await assert.rejects(() => client.paginateAll("/spaces", {}), /already stopped/); + assert.equal(requests, 0, "no /spaces request started once already aborted"); +}); + +// Without a tool signal (default), pagination is unaffected — the safe-point is a +// pure no-op, so pre-#487 behaviour is byte-identical. +test("paginateAll is unaffected when no tool signal is set", async () => { + let requests = 0; + const PAGES = { + "": { items: [{ id: "a" }], nextCursor: "c1" }, + c1: { items: [{ id: "b" }], nextCursor: null }, + }; + const { baseURL } = await spawn(async (req, res) => { + const raw = await readBody(req); + if (handleLogin(req, res)) return; + if (req.url === "/api/spaces") { + requests++; + const body = JSON.parse(raw || "{}"); + const page = PAGES[body.cursor ?? ""] ?? { items: [], nextCursor: null }; + sendJson(res, 200, { + success: true, + data: { + items: page.items, + meta: { hasNextPage: page.nextCursor != null, nextCursor: page.nextCursor }, + }, + }); + return; + } + sendJson(res, 404, {}); + }); + + const client = new DocmostClient(baseURL, "user@example.com", "pw"); + const all = await client.paginateAll("/spaces", {}); + assert.equal(requests, 2, "both pages fetched with no signal set"); + assert.deepEqual(all.map((p) => p.id), ["a", "b"]); +}); diff --git a/packages/mcp/test/mock/write-abort-safepoint.test.mjs b/packages/mcp/test/mock/write-abort-safepoint.test.mjs new file mode 100644 index 00000000..1002b34b --- /dev/null +++ b/packages/mcp/test/mock/write-abort-safepoint.test.mjs @@ -0,0 +1,164 @@ +// #487 F4 — the WRITE-side cancellation safe-point. +// +// Every content-mutating collab write (collaboration.mutatePageContent and the +// reentrant twin client.mutateLiveContentUnlocked used by replaceImage) checks +// the in-app tool abort signal at a PRE-COMMIT safe-point — after the collab +// session is acquired but immediately BEFORE the atomic read->write +// (session.mutate). So a Stop (or the per-call cap) that lands during the +// connect/lock window stops THIS write from landing: no new commit starts once +// aborted. paginate-abort-safepoint.test.mjs pins the READ half; this pins the +// integrity-critical WRITE half — remove the `throwIfAborted()` and the transform +// would run and the doc would be mutated past a Stop. +// +// There is no collab server in the unit env, so we swap the provider factory +// (__setCollabProviderFactory) for a fake that reports an immediate successful +// sync. That makes acquireCollabSession SUCCEED and hand back a live, ready +// session, so the ONLY thing standing between the call and session.mutate is the +// safe-point under test. The transform is instrumented to prove it never runs. +import { test, afterEach } from "node:test"; +import assert from "node:assert/strict"; +import { mutatePageContent } from "../../build/lib/collaboration.js"; +import { + __setCollabProviderFactory, + destroyAllSessions, +} from "../../build/lib/collab-session.js"; +import { DocmostClient } from "../../build/client.js"; + +const BASE_URL = "http://127.0.0.1:1/api"; +// mutatePageContent locks via withPageLock, which demands a canonical page UUID +// (resolve-then-lock invariant, #260/#449); the unlocked twin does not. +const PAGE_UUID = "11111111-1111-4111-8111-111111111111"; + +// A fake HocuspocusProvider that immediately reports a successful initial sync so +// CollabSession.open() resolves to a ready session, and stays "synced" with zero +// unsynced changes so a reached session.mutate() would resolve at once. It speaks +// only the tiny CollabProviderLike surface the session depends on. +function syncedProviderFactory(config) { + // Fire the initial-sync callback so open() settles as ready. + config.onSynced(); + return { + synced: true, + unsyncedChanges: 0, + destroy() {}, + on() {}, + off() {}, + }; +} + +// Disable the session cache so every acquire opens (and the failure path destroys) +// its own ephemeral session — no cross-test session reuse. +process.env.MCP_COLLAB_SESSION_IDLE_MS = "0"; + +afterEach(() => { + __setCollabProviderFactory(null); // restore the real factory + destroyAllSessions(); +}); + +// --- collaboration.mutatePageContent (the page-locked write path) ------------- + +test("mutatePageContent rejects at the pre-commit safe-point BEFORE session.mutate when the signal is already aborted", async () => { + __setCollabProviderFactory(syncedProviderFactory); + let transformCalls = 0; + const ac = new AbortController(); + ac.abort(new Error("user stop")); + + await assert.rejects( + () => + mutatePageContent( + PAGE_UUID, + "collab-jwt", + BASE_URL, + (liveDoc) => { + transformCalls++; + return liveDoc; + }, + ac.signal, + ), + /user stop/, + "the aborted safe-point rejects with the signal's reason before committing", + ); + assert.equal( + transformCalls, + 0, + "the transform (and therefore session.mutate) must NEVER run once aborted", + ); +}); + +test("mutatePageContent (control) DOES reach session.mutate and invoke the transform when the signal is live", async () => { + __setCollabProviderFactory(syncedProviderFactory); + let transformCalls = 0; + const ac = new AbortController(); // never aborted + + const result = await mutatePageContent( + PAGE_UUID, + "collab-jwt", + BASE_URL, + (liveDoc) => { + transformCalls++; + return null; // null -> no-op write; still proves the transform was invoked + }, + ac.signal, + ); + assert.equal( + transformCalls, + 1, + "with a live signal the safe-point is a no-op and session.mutate runs the transform", + ); + assert.ok(result && result.verify, "a MutationResult is returned"); +}); + +// --- client.mutateLiveContentUnlocked (the reentrant twin, replaceImage) ------ + +test("mutateLiveContentUnlocked rejects at the pre-commit safe-point BEFORE session.mutate when the tool signal is already aborted", async () => { + __setCollabProviderFactory(syncedProviderFactory); + const client = new DocmostClient({ + apiUrl: BASE_URL, + getToken: async () => "access", + getCollabToken: async () => "collab-jwt", + }); + let transformCalls = 0; + const ac = new AbortController(); + ac.abort(new Error("cap fired")); + client.setToolAbortSignal(ac.signal); + + await assert.rejects( + () => + client.mutateLiveContentUnlocked("page-1", "collab-jwt", (liveDoc) => { + transformCalls++; + return liveDoc; + }), + /cap fired/, + "the aborted safe-point rejects with the signal's reason before committing", + ); + assert.equal( + transformCalls, + 0, + "the transform (and therefore session.mutate) must NEVER run once aborted", + ); +}); + +test("mutateLiveContentUnlocked (control) DOES reach session.mutate and invoke the transform when the tool signal is live", async () => { + __setCollabProviderFactory(syncedProviderFactory); + const client = new DocmostClient({ + apiUrl: BASE_URL, + getToken: async () => "access", + getCollabToken: async () => "collab-jwt", + }); + let transformCalls = 0; + client.setToolAbortSignal(new AbortController().signal); // live + + const result = await client.mutateLiveContentUnlocked( + "page-1", + "collab-jwt", + (liveDoc) => { + transformCalls++; + return null; + }, + ); + assert.equal( + transformCalls, + 1, + "with a live signal the safe-point is a no-op and session.mutate runs the transform", + ); + assert.ok(result && result.verify, "a MutationResult is returned"); +});