diff --git a/.env.example b/.env.example index f3acb7ba..6d7dec43 100644 --- a/.env.example +++ b/.env.example @@ -225,11 +225,26 @@ MCP_DOCMOST_PASSWORD= # Silence timeout (ms) for EXTERNAL-MCP transport ONLY (not the chat provider). # Tighter than AI_STREAM_TIMEOUT_MS so a byte-silent/hung MCP server is broken in -# ~1 min instead of 15. Note it also cuts a legitimately long but byte-silent -# single tool call (a slow crawl that emits nothing until done) and an SSE -# transport idling >1 min BETWEEN tool calls. Default 60000 (1 min). +# ~1 min instead of 15. It cuts a legitimately long but byte-silent single tool +# call (a slow crawl that emits nothing until done) on the HTTP (streamable) +# transport, which opens a fresh request per call. The SSE transport — one +# long-lived body across many calls — is NO LONGER governed by this timeout +# (as of #489): its idle-BETWEEN-calls window has its own, raised bodyTimeout, +# AI_MCP_SSE_BODY_TIMEOUT_MS below. Default 60000 (1 min). # AI_MCP_STREAM_TIMEOUT_MS=60000 +# bodyTimeout (ms) for the EXTERNAL-MCP SSE transport ONLY (#489). The SSE +# transport holds ONE response body open across many tool calls, so undici's +# bodyTimeout (time between body bytes) counts the LEGITIMATE silence BETWEEN the +# model's tool calls, not just a hung single call. At the tight 1-min silence +# timeout above, a normal >1-min gap between calls would break the SSE socket and +# the cache would serve a dead client until TTL — so the SSE transport gets its +# OWN, RAISED bodyTimeout. A single stuck call is still bounded by the per-call +# cap (AI_MCP_CALL_TIMEOUT_MS), and a socket that does break is healed by the +# in-run transport-error retry. The HTTP (streamable) transport keeps the tight +# timeout. Default 600000 (10 min). +# AI_MCP_SSE_BODY_TIMEOUT_MS=600000 + # Total wall-clock cap (ms) for ONE external MCP tool call (app-level, not # transport). Aborts a tool that keeps the socket warm (SSE heartbeats / trickle) # but never returns a result — which the silence timeout above never breaks. diff --git a/CHANGELOG.md b/CHANGELOG.md index 30d27413..adecf01f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -336,6 +336,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **A chat with one malformed message part no longer 500s on every turn, and a + failed send no longer duplicates the user's message.** Incoming client parts + are now whitelisted to `text` (a forged tool-result part can no longer reach + the persisted history or the model context), and the turn is converted BEFORE + the user row is inserted, so a mid-flight failure cannot leave a duplicate + user row that a retry then compounds. A single part that still fails to convert + degrades to a `[tool context omitted]` marker on that one row instead of + bricking the whole chat. (#489) +- **A transport drop to an external MCP server now heals within the same turn.** + On an undici transport error, a read-only MCP tool reconnects its server and + retries once within the run; a write is never auto-retried (it may already have + applied). One flapping server no longer nulls the shared client cache, so other + servers' cached clients are untouched. The SSE transport also gets a raised + body-timeout so a legitimate >1-min idle between the model's tool calls no + longer breaks a long-lived SSE socket (new `AI_MCP_SSE_BODY_TIMEOUT_MS`, default + 10 min; see `.env.example`). (#489) + - **The server no longer runs out of heap during long autonomous agent runs.** A new pnpm patch on `ai@6.0.134` stops the SDK from building a cumulative snapshot of the ENTIRE turn text on every streamed text-delta when no output diff --git a/apps/server/src/core/ai-chat/ai-chat.parts-resilience.spec.ts b/apps/server/src/core/ai-chat/ai-chat.parts-resilience.spec.ts new file mode 100644 index 00000000..d0b66432 --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat.parts-resilience.spec.ts @@ -0,0 +1,142 @@ +// #489 — client-parts validation + resilient history conversion. +// +// These unit tests exercise the two exported helpers against the REAL +// `convertToModelMessages` from `ai` (NOT a mock): a genuinely malformed part +// (a `null` element inside a parts array) makes the real converter throw +// ("Cannot read properties of null"), which is the actual production +// "bricked chat" mechanism this fix defends against. Asserting against the real +// converter (rather than a mock-shaped error) is the whole point — a mock would +// hide a version change in the converter's throw behaviour. +import { convertToModelMessages, type UIMessage } from 'ai'; +import { + sanitizeUserParts, + convertHistoryResilient, + TOOL_CONTEXT_OMITTED_MARKER, +} from './ai-chat.service'; + +type Row = Omit & { id: string }; + +describe('sanitizeUserParts (#489, branch: validation on receipt)', () => { + it('keeps whitelisted text parts unchanged', () => { + const drops: string[] = []; + const out = sanitizeUserParts( + [ + { type: 'text', text: 'a' }, + { type: 'text', text: 'b' }, + ] as UIMessage['parts'], + (t) => drops.push(t), + ); + expect(out).toEqual([ + { type: 'text', text: 'a' }, + { type: 'text', text: 'b' }, + ]); + expect(drops).toEqual([]); + }); + + it('drops a non-text part (a tool-part in input-available) and reports its type', () => { + const drops: string[] = []; + const out = sanitizeUserParts( + [ + { type: 'text', text: 'hi' }, + { + type: 'tool-getPage', + toolCallId: 't1', + state: 'input-available', + input: { pageId: 'p' }, + }, + ] as unknown as UIMessage['parts'], + (t) => drops.push(t), + ); + expect(out).toEqual([{ type: 'text', text: 'hi' }]); + expect(drops).toEqual(['tool-getPage']); + }); + + it('drops a null part (the shape that would poison convertToModelMessages)', () => { + const drops: string[] = []; + const out = sanitizeUserParts( + [{ type: 'text', text: 'hi' }, null] as unknown as UIMessage['parts'], + (t) => drops.push(t), + ); + expect(out).toEqual([{ type: 'text', text: 'hi' }]); + expect(drops).toEqual(['(unknown)']); + }); + + it('returns undefined when nothing survives (so a null metadata is persisted)', () => { + const out = sanitizeUserParts( + [ + { type: 'tool-x', toolCallId: 't', state: 'input-available' }, + ] as unknown as UIMessage['parts'], + () => undefined, + ); + expect(out).toBeUndefined(); + }); + + it('returns undefined for a non-array input', () => { + expect( + sanitizeUserParts(undefined as unknown as UIMessage['parts'], () => undefined), + ).toBeUndefined(); + }); +}); + +describe('convertHistoryResilient (#489, branches: happy + per-row degradation)', () => { + it('happy path: healthy history converts identically to convertToModelMessages, no degrade', async () => { + const history: Row[] = [ + { id: 'u1', role: 'user', parts: [{ type: 'text', text: 'hi' }] }, + { id: 'a1', role: 'assistant', parts: [{ type: 'text', text: 'hello' }] }, + ]; + const degrades: number[] = []; + const out = await convertHistoryResilient(history, (i) => degrades.push(i)); + const expected = await convertToModelMessages(history as UIMessage[]); + expect(out).toEqual(expected); + expect(degrades).toEqual([]); + }); + + it('REAL poison: a null part throws in the batch converter but is isolated and degraded to a marker', async () => { + // Sanity: the real converter genuinely throws on this shape. + const poisoned: Row = { + id: 'a1', + role: 'assistant', + parts: [ + { type: 'text', text: 'earlier answer' }, + null, + ] as unknown as UIMessage['parts'], + }; + await expect( + convertToModelMessages([poisoned as UIMessage]), + ).rejects.toThrow(); + + const history: Row[] = [ + { id: 'u1', role: 'user', parts: [{ type: 'text', text: 'first' }] }, + poisoned, + { id: 'u2', role: 'user', parts: [{ type: 'text', text: 'second' }] }, + ]; + const degrades: number[] = []; + const out = await convertHistoryResilient(history, (i) => degrades.push(i)); + + // Only the poisoned row (index 1) is degraded. + expect(degrades).toEqual([1]); + // Healthy rows survive verbatim. + const flat = JSON.stringify(out); + expect(flat).toContain('first'); + expect(flat).toContain('second'); + // The degraded row carries its readable text AND the truncation marker so the + // model sees that tool context was omitted (never a silent loss). + expect(flat).toContain('earlier answer'); + expect(flat).toContain(TOOL_CONTEXT_OMITTED_MARKER); + // The whole batch converted (3 model messages, none dropped). + expect(out).toHaveLength(3); + }); + + it('a fully-poisoned row (no readable text) still degrades to just the marker', async () => { + const history: Row[] = [ + { + id: 'a1', + role: 'assistant', + parts: [null] as unknown as UIMessage['parts'], + }, + ]; + const out = await convertHistoryResilient(history, () => undefined); + expect(out).toHaveLength(1); + expect(JSON.stringify(out)).toContain(TOOL_CONTEXT_OMITTED_MARKER); + }); +}); 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 a6f3e75c..22a721e8 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 @@ -97,8 +97,14 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => { }; const runService = new AiChatRunService(runRepo as never, { isCloud: () => false } as never); - // The user-message insert (the first bare await after beginRun) throws. + // The user-message insert throws. #489 runs the history load + convert BEFORE + // the insert (convert-before-insert, so a retry cannot duplicate the user row), + // so `findAllByChat` (a real repo method) is now called first — stub it to an + // empty history so the flow reaches the insert. Both awaits are AFTER beginRun, + // so the "exception after beginRun -> settled to error" invariant is unchanged; + // the throw point simply moved from insert to a later insert after a no-op load. const aiChatMessageRepo = { + findAllByChat: jest.fn().mockResolvedValue([]), insert: jest.fn().mockRejectedValue(new Error('insert boom')), }; const aiChatRepo = { 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 ca9511ef..8dd32123 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 @@ -1440,7 +1440,7 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', () } // Wire only the deps reached on the way to the pipe call, plus a spy registry. - function makeService(opts: { resumable: boolean }) { + function makeService(opts: { resumable: boolean; history?: unknown[] }) { const aiChatRepo = { findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })), insert: jest.fn(), @@ -1448,7 +1448,7 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', () const aiChatMessageRepo = { // Both the user insert and the assistant seed return the same row id. insert: jest.fn(async () => ({ id: 'msg-1' })), - findAllByChat: jest.fn(async () => []), + findAllByChat: jest.fn(async () => opts.history ?? []), update: jest.fn(async () => ({ id: 'msg-1' })), // #487: the terminal owner-write + the opportunistic reconcile query. finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })), @@ -1487,7 +1487,7 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', () } as never, streamRegistry as never, ); - return { svc, streamRegistry }; + return { svc, streamRegistry, aiChatMessageRepo }; } const body = { @@ -1570,6 +1570,86 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', () await expect(drive(svc, makeRunHooks())).rejects.toThrow('boom'); expect(streamRegistry.abortEntry).toHaveBeenCalledWith('chat-1', 'run-1'); }); + + // #489 REGRESSION (against the REAL convertToModelMessages — not mocked here): + // a persisted history row whose parts contain a `null` element makes the real + // convertToModelMessages THROW ("Cannot read properties of null"). Pre-fix that + // 500-ed every turn forever and each retry appended a duplicate user row. The + // fix converts BEFORE the insert and isolates the poisoned row per-row, degrading + // it to text with a "[tool context omitted]" marker. Assert the turn still runs, + // the marker reaches the model, and exactly ONE user row is inserted. + it('#489: a poisoned OLD-history row keeps the chat working; the marker reaches the model; one user insert', async () => { + const { svc, aiChatMessageRepo } = makeService({ + resumable: false, + history: [ + { + id: 'old-1', + role: 'assistant', + content: 'earlier answer', + // A null part is the poison: rowToUiMessage keeps it (the array is + // non-empty) and the real convertToModelMessages throws on it. + metadata: { parts: [{ type: 'text', text: 'earlier answer' }, null] }, + status: 'completed', + }, + ], + }); + // Must NOT throw — the poisoned row is degraded, not fatal. + await drive(svc, makeRunHooks()); + expect(streamTextMock).toHaveBeenCalledTimes(1); + const passedMessages = streamTextMock.mock.calls[0][0].messages; + const serialized = JSON.stringify(passedMessages); + // The model sees the truncation marker (silent tool-context loss is not ok) + // AND the row's readable text is preserved alongside it. + expect(serialized).toContain('[tool context omitted]'); + expect(serialized).toContain('earlier answer'); + // Exactly ONE user row inserted (no duplicate), inserted AFTER conversion. + const userInserts = aiChatMessageRepo.insert.mock.calls + .map((c: unknown[]) => c[0] as { role?: string }) + .filter((r) => r.role === 'user'); + expect(userInserts).toHaveLength(1); + }); + + // #489: client-supplied non-text parts (a tool-part in `input-available`, the + // exact "bricking" payload) are dropped ON RECEIPT — never persisted — so they + // can never poison future turns. Only the text survives into metadata.parts. + it('#489: a non-text client part is stripped before persist (only text survives)', async () => { + const { svc, aiChatMessageRepo } = makeService({ resumable: false }); + await svc.stream({ + user: { id: 'u1' } as never, + workspace: { id: 'ws-1' } as never, + sessionId: 's1', + body: { + chatId: 'chat-1', + messages: [ + { + id: 'm1', + role: 'user', + parts: [ + { type: 'text', text: 'hello' }, + // untrusted tool-part — must be dropped, never persisted + { + type: 'tool-getPage', + toolCallId: 't1', + state: 'input-available', + input: { pageId: 'p' }, + }, + ], + }, + ], + } as never, + res: makeRes() as never, + signal: new AbortController().signal, + model: {} as never, + role: null, + runHooks: makeRunHooks() as never, + }); + const userInsert = aiChatMessageRepo.insert.mock.calls + .map((c: unknown[]) => c[0] as { role?: string; metadata?: unknown }) + .find((r) => r.role === 'user'); + const parts = (userInsert?.metadata as { parts?: Array<{ type: string }> }) + ?.parts; + expect(parts).toEqual([{ type: 'text', text: 'hello' }]); + }); }); /** 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 02ecc9d1..66d21473 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -14,6 +14,7 @@ import { convertToModelMessages, stepCountIs, type UIMessage, + type ModelMessage, type LanguageModel, } from 'ai'; import { AiService } from '../../integrations/ai/ai.service'; @@ -1042,7 +1043,58 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { const incoming = lastUserMessage(body.messages); const incomingText = uiMessageText(incoming); - // Persist the user message before contacting the model. + // #489: sanitize client-supplied parts ON RECEIPT. The client only ever + // sends `sendMessage({ text })` (a single text part); there is no + // file/attachment path. Any other part — most dangerously a tool-part in + // `input-available` state — is untrusted data that, once persisted to + // `metadata.parts` verbatim, is REPLAYED through convertToModelMessages on + // every later turn. A malformed tool-part makes that conversion throw, + // 500-ing every future turn of the chat forever ("bricked"). Drop any + // non-whitelisted part with a warn. + const sanitizedParts = sanitizeUserParts(incoming?.parts, (type) => + this.logger.warn( + `Dropping unsupported user message part '${type}' on chat ${chatId}`, + ), + ); + + // #489: rebuild the conversation from persisted history (not the client + // payload) and CONVERT it to model messages BEFORE persisting the user row. + // Load the OLD history (WITHOUT the new row) and append the incoming turn in + // memory for the conversion. This makes the insert happen only after a + // successful conversion, so a conversion failure cannot leave a DUPLICATE + // user row behind on the client's retry (the "bricked chat" that accreted a + // dup on every 500). `findAllByChat` returns chronological order (oldest -> + // newest) and keeps a 5000-row memory-safety backstop (on overflow it keeps + // the NEWEST rows and logs a warning); that is a safety net far above any + // realistic chat, not a conversational limit. + const oldHistory = await this.aiChatMessageRepo.findAllByChat( + chatId, + workspace.id, + ); + const uiMessages: Array & { id: string }> = [ + ...oldHistory.map(rowToUiMessage), + { + id: 'pending-user', + role: 'user', + parts: (sanitizedParts && sanitizedParts.length > 0 + ? sanitizedParts + : textPart(incomingText)) as UIMessage['parts'], + }, + ]; + // convertToModelMessages is async in ai@6.0.134 (returns Promise). + // Resilient (#489): a single poisoned row in the OLD history is isolated via + // per-row conversion and degraded to plain text with a "[tool context + // omitted]" marker rather than 500-ing the whole turn (silent loss of tool + // context is not acceptable — the model must see the truncation). + const messages = await convertHistoryResilient(uiMessages, (index, err) => + this.logger.warn( + `Degraded unconvertible history row ${index} on chat ${chatId} to text: ${ + err instanceof Error ? err.message : 'unknown error' + }`, + ), + ); + + // Persist the user message only AFTER a successful conversion (#489). await this.aiChatMessageRepo.insert({ chatId, workspaceId: workspace.id, @@ -1050,31 +1102,21 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { role: 'user', content: incomingText, // jsonb column: UIMessage parts are JSON-serializable at runtime but not - // structurally `JsonValue`, so cast through unknown. - metadata: (incoming?.parts ? { parts: incoming.parts } : null) as never, + // structurally `JsonValue`, so cast through unknown. Persist the SANITIZED + // parts (never the raw client parts) so the row is always convertible. + metadata: (sanitizedParts ? { parts: sanitizedParts } : null) as never, }); - // Rebuild the conversation from persisted history (not the client payload), - // so the model always sees the authoritative server-side transcript. Load - // the FULL history in chronological order (oldest -> newest, incl. the user - // message just inserted above) so NO turns are dropped — there is no - // recent-tail window anymore. `findAllByChat` keeps a 5000-row memory-safety - // backstop (on overflow it keeps the NEWEST rows and logs a warning); that - // is a safety net far above any realistic chat, not a conversational limit. - const history = await this.aiChatMessageRepo.findAllByChat( - chatId, - workspace.id, - ); - const uiMessages = history.map(rowToUiMessage); - // convertToModelMessages is async in ai@6.0.134 (returns Promise). - const messages = await convertToModelMessages(uiMessages); - // Interrupt-resume detection (#198): the client "send now" flag is only a // hint — confirm it against the persisted history (the preceding assistant // turn must really be aborted/streaming) so a spoofed flag cannot inject the // interrupt note onto an ordinary turn. The partial output the model needs is // already in `messages` (the aborted assistant row replays via findRecent). - const interrupted = isInterruptResume(history, body.interrupted); + // Append the new user turn (shape-only) so index -2 is the prior assistant. + const interrupted = isInterruptResume( + [...oldHistory, { role: 'user', status: null, metadata: null }], + body.interrupted, + ); // Per-turn page-change detection (#274): if the open page was hand-edited by // the user since the agent's last turn ended, compute the unified diff so the @@ -2074,6 +2116,82 @@ function textPart(text: string): Array<{ type: 'text'; text: string }> { return text ? [{ type: 'text', text }] : []; } +/** + * Part types accepted on an INCOMING user turn (#489). The client only ever + * sends `sendMessage({ text })` (a single text part); there is no file/attachment + * path. Everything else on a client-supplied user message — most dangerously a + * tool-part in `input-available` state — is untrusted data that would be + * persisted to `metadata.parts` verbatim and replayed through + * `convertToModelMessages` on every later turn, potentially bricking the chat. + */ +const ALLOWED_USER_PART_TYPES: ReadonlySet = new Set(['text']); + +/** + * Keep only whitelisted parts on a client-supplied user message; report each + * dropped part's type via `onDrop` (the caller warns). Returns `undefined` when + * nothing survives (no parts / none whitelisted), so the caller persists a null + * metadata rather than an empty-parts object. Never throws. + */ +export function sanitizeUserParts( + parts: UIMessage['parts'] | undefined, + onDrop: (type: string) => void, +): UIMessage['parts'] | undefined { + if (!Array.isArray(parts)) return undefined; + const kept = parts.filter((p) => { + const type = + typeof (p as { type?: unknown })?.type === 'string' + ? (p as { type: string }).type + : ''; + if (ALLOWED_USER_PART_TYPES.has(type)) return true; + onDrop(type || '(unknown)'); + return false; + }); + return kept.length > 0 ? (kept as UIMessage['parts']) : undefined; +} + +/** Marker for a history row whose tool parts could not be replayed (#489). */ +export const TOOL_CONTEXT_OMITTED_MARKER = '[tool context omitted]'; + +/** + * Convert persisted UI history to model messages, tolerating a single poisoned + * row (#489). `convertToModelMessages` over the WHOLE array throws if ANY row is + * malformed (e.g. a tool-part left unbalanced / in `input-available` state), + * which would otherwise 500 every turn of the chat forever. On a batch failure we + * fall back to per-row conversion so the bad row is isolated: it is degraded to + * plain text carrying its readable text plus a `[tool context omitted]` marker + * (the model MUST see that its tool context was truncated — silent loss is not + * acceptable), while every healthy row converts normally. Because AI SDK v6 + * carries a tool call and its result inside the SAME assistant UIMessage's parts, + * per-row conversion preserves call/result pairing. + */ +export async function convertHistoryResilient( + uiMessages: Array & { id: string }>, + onDegrade: (index: number, err: unknown) => void, +): Promise { + try { + return await convertToModelMessages(uiMessages as UIMessage[]); + } catch { + const out: ModelMessage[] = []; + for (let i = 0; i < uiMessages.length; i++) { + const m = uiMessages[i]; + try { + out.push(...(await convertToModelMessages([m as UIMessage]))); + } catch (err) { + onDegrade(i, err); + const text = uiMessageText(m as UIMessage); + const degraded = text + ? `${text}\n\n${TOOL_CONTEXT_OMITTED_MARKER}` + : TOOL_CONTEXT_OMITTED_MARKER; + out.push({ + role: m.role === 'assistant' ? 'assistant' : 'user', + content: degraded, + } as ModelMessage); + } + } + return out; + } +} + /** * Minimal shapes of the AI SDK v6 step objects we read to rebuild UIMessage * parts (see ai@6.0.134 `StepResult`: `text`, `toolCalls` -> TypedToolCall, diff --git a/apps/server/src/core/ai-chat/external-mcp/mcp-clients.recovery.spec.ts b/apps/server/src/core/ai-chat/external-mcp/mcp-clients.recovery.spec.ts new file mode 100644 index 00000000..5cf83721 --- /dev/null +++ b/apps/server/src/core/ai-chat/external-mcp/mcp-clients.recovery.spec.ts @@ -0,0 +1,261 @@ +import { errors } from 'undici'; +import { + McpClientsService, + isRetryableConnectError, +} from './mcp-clients.service'; + +/** + * #489 — external-MCP in-run transport recovery. + * + * The transport-error classification + retry gate are exercised against the REAL + * undici error CLASSES prod throws (`errors.SocketError` / `errors.BodyTimeoutError`, + * carrying the true `UND_ERR_*` codes and class names), wrapped EXACTLY as undici's + * `fetch` wraps them — a `TypeError('fetch failed'|'terminated')` whose `.cause` is + * the undici error. These are the real classes, not hand-rolled `{code:'...'}` + * mocks: constructing the genuine class is what makes this a faithful test of the + * prod predicate (epic root-cause #4 — a mock-shaped predicate would leave the + * evict/retry path silently dead in production while CI stays green). We construct + * rather than drive a live fetch because Jest's environment degrades the live-fetch + * error to a generic `Error` cause (no undici code), which would NOT be the prod + * shape. + */ + +/** A REAL undici socket reset, wrapped as fetch wraps it. */ +function realSocketResetError(): unknown { + const err = new TypeError('fetch failed'); + (err as { cause?: unknown }).cause = new errors.SocketError('other side closed'); + return err; +} + +/** A REAL undici body timeout, wrapped as fetch wraps it. */ +function realBodyTimeoutError(): unknown { + const err = new TypeError('terminated'); + (err as { cause?: unknown }).cause = new errors.BodyTimeoutError(); + return err; +} + +type FakeServer = { + id: string; + name: string; + transport: 'http' | 'sse'; + url: string; + headersEnc: string | null; + toolAllowlist: string[] | null; + instructions: string | null; +}; + +const server = (over: Partial = {}): FakeServer => ({ + id: 's1', + name: 'srv', + transport: 'http', + url: 'http://example.test/mcp', + headersEnc: null, + toolAllowlist: null, + instructions: null, + ...over, +}); + +function buildService(servers: FakeServer[], trusted = false) { + const repo = { listEnabled: jest.fn().mockResolvedValue(servers) }; + const service = new McpClientsService(repo as never, {} as never); + // Seed a DETERMINISTIC write-class map so the retry gate is controlled here + // (the production map loads from @docmost/mcp via a dynamic ESM import). getPage + // is a read, patchNode is a write — the real classifications. + ( + service as unknown as { writeClassMapPromise: Promise } + ).writeClassMapPromise = Promise.resolve({ + getPage: 'readOnly', + patchNode: 'write', + }); + // The service only APPLIES that map to a TRUSTED internal Docmost server + // (isInternalDocmostServer, really false for every third-party row). A retry + // test needs a trusted server to exercise the readOnly-retry path at all, so it + // passes trusted=true to model a Docmost-origin server; the third-party + // double-apply test leaves it at the real value (false). + if (trusted) { + jest + .spyOn( + service as unknown as { + isInternalDocmostServer: (s: FakeServer) => boolean; + }, + 'isInternalDocmostServer', + ) + .mockReturnValue(true); + } + return { service, repo }; +} + +/** Spy the private `connect` so each call yields a controlled fake client whose + * single tool's execute is the supplied function. Returns the connect spy. */ +function stubConnect( + service: McpClientsService, + toolName: string, + execs: Array<(...a: unknown[]) => Promise>, +) { + let n = 0; + return jest + .spyOn( + service as unknown as { connect: (s: FakeServer) => Promise }, + 'connect', + ) + .mockImplementation(async () => { + const exec = execs[Math.min(n, execs.length - 1)]; + n += 1; + return { + tools: async () => ({ [toolName]: { description: 'x', execute: exec } }), + close: jest.fn().mockResolvedValue(undefined), + }; + }); +} + +const opts = (abortSignal?: AbortSignal) => + ({ toolCallId: 't', messages: [], abortSignal }) as never; + +describe('isRetryableConnectError (#489, REAL error shapes)', () => { + it('classifies a real undici socket reset and body timeout as retryable', async () => { + const socketErr = await realSocketResetError(); + const bodyErr = await realBodyTimeoutError(); + expect(isRetryableConnectError(socketErr)).toBe(true); + expect(isRetryableConnectError(bodyErr)).toBe(true); + // Unwraps a wrapped cause chain (e.g. an MCPClientError around the socket err). + const wrapped = new Error('mcp call failed'); + (wrapped as { cause?: unknown }).cause = socketErr; + expect(isRetryableConnectError(wrapped)).toBe(true); + }); + + it('does NOT classify an application-level error as a transport break', () => { + expect(isRetryableConnectError(new Error('validation failed'))).toBe(false); + expect(isRetryableConnectError({ name: 'HttpError', status: 400 })).toBe(false); + expect(isRetryableConnectError(undefined)).toBe(false); + expect(isRetryableConnectError('boom')).toBe(false); + }); +}); + +describe('McpClientsService in-run transport recovery (#489)', () => { + afterEach(() => jest.restoreAllMocks()); + + it('a readOnly tool whose transport breaks reconnects and retries WITHIN the same run', async () => { + const realErr = await realSocketResetError(); + const { service } = buildService([server()], true); + const first = jest.fn().mockRejectedValue(realErr); + const second = jest.fn().mockResolvedValue({ ok: true }); + const connectSpy = stubConnect(service, 'getPage', [first, second]); + + const toolset = await service.toolsFor('ws-1'); + const tool = toolset.tools['srv_getPage']; + const result = await (tool.execute as (a: unknown, o: unknown) => Promise)( + { pageId: 'p' }, + opts(), + ); + + // The repeat call within the run got a LIVE client and succeeded. + expect(result).toEqual({ ok: true }); + expect(first).toHaveBeenCalledTimes(1); + expect(second).toHaveBeenCalledTimes(1); + // Exactly one reconnect was minted (initial build connect + one recovery). + expect(connectSpy).toHaveBeenCalledTimes(2); + // The run accumulated BOTH leases (old + reconnected) — released together at end. + expect(toolset.clients).toHaveLength(2); + await Promise.all(toolset.clients.map((c) => c.close())); + }); + + it('a WRITE tool does NOT auto-retry on a transport error (indeterminate)', async () => { + const realErr = await realSocketResetError(); + const { service } = buildService([server()], true); + const exec = jest.fn().mockRejectedValue(realErr); + const connectSpy = stubConnect(service, 'patchNode', [exec]); + + const toolset = await service.toolsFor('ws-2'); + const tool = toolset.tools['srv_patchNode']; + await expect( + (tool.execute as (a: unknown, o: unknown) => Promise)( + { pageId: 'p' }, + opts(), + ), + ).rejects.toThrow(/MAY have already applied/); + + // Called exactly once — NO blind retry (avoids double-apply, the #435 class). + expect(exec).toHaveBeenCalledTimes(1); + // No fresh connection was minted for a write. + expect(connectSpy).toHaveBeenCalledTimes(1); + await Promise.all(toolset.clients.map((c) => c.close())); + }); + + it('does NOT retry (or reconnect) after the run is aborted (Stop)', async () => { + const realErr = await realSocketResetError(); + const { service } = buildService([server()], true); + const controller = new AbortController(); + // The transport error arrives, but the run was Stopped in the same tick. + const first = jest.fn().mockImplementation(async () => { + controller.abort(); + throw realErr; + }); + const second = jest.fn().mockResolvedValue({ ok: true }); + const connectSpy = stubConnect(service, 'getPage', [first, second]); + + const toolset = await service.toolsFor('ws-3'); + const tool = toolset.tools['srv_getPage']; + await expect( + (tool.execute as (a: unknown, o: unknown) => Promise)( + { pageId: 'p' }, + opts(controller.signal), + ), + ).rejects.toBeDefined(); + + // getPage IS readOnly, but the Stop blocks the retry — no second call, no mint. + expect(second).not.toHaveBeenCalled(); + expect(connectSpy).toHaveBeenCalledTimes(1); + await Promise.all(toolset.clients.map((c) => c.close())); + }); + + it('an app-level (non-transport) tool error is surfaced verbatim, never retried', async () => { + const { service } = buildService([server()], true); + const appErr = new Error('tool says: bad input'); + const exec = jest.fn().mockRejectedValue(appErr); + const connectSpy = stubConnect(service, 'getPage', [exec]); + + const toolset = await service.toolsFor('ws-4'); + const tool = toolset.tools['srv_getPage']; + await expect( + (tool.execute as (a: unknown, o: unknown) => Promise)( + { pageId: 'p' }, + opts(), + ), + ).rejects.toThrow('tool says: bad input'); + expect(exec).toHaveBeenCalledTimes(1); + expect(connectSpy).toHaveBeenCalledTimes(1); // no reconnect for an app error + await Promise.all(toolset.clients.map((c) => c.close())); + }); + + // #489 (review, MEDIUM) — the Docmost write-class map keys by DOCMOST tool + // names; a THIRD-PARTY server may name a WRITE tool `getPage` (a Docmost read + // name). It must NOT inherit readOnly and must NOT auto-retry on a transport + // error — a blind retry of that write is a double-apply (the #435 class). Here + // the server is UNTRUSTED (buildService default, isInternalDocmostServer=false), + // so the map is not applied and `getPage` classifies as a write. + // + // MUTATION-VERIFY: forcing the server "trusted" (buildService(..., true)) makes + // `getPage` inherit readOnly -> it WOULD reconnect+retry (connect twice) and the + // assertions below fail — i.e. removing the trust scope re-opens the bug. + it('a THIRD-PARTY WRITE tool named like a Docmost read does NOT auto-retry (no double-apply)', async () => { + const realErr = await realSocketResetError(); + // Untrusted: default trusted=false — a real third-party server. + const { service } = buildService([server()]); + const exec = jest.fn().mockRejectedValue(realErr); + const connectSpy = stubConnect(service, 'getPage', [exec, exec]); + + const toolset = await service.toolsFor('ws-5'); + const tool = toolset.tools['srv_getPage']; + await expect( + (tool.execute as (a: unknown, o: unknown) => Promise)( + { pageId: 'p' }, + opts(), + ), + ).rejects.toThrow(/MAY have already applied/); + + // Exactly one call, NO reconnect — the name collision granted no readOnly-retry. + expect(exec).toHaveBeenCalledTimes(1); + expect(connectSpy).toHaveBeenCalledTimes(1); + await Promise.all(toolset.clients.map((c) => c.close())); + }); +}); diff --git a/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.spec.ts b/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.spec.ts index 53ad6191..270b2df1 100644 --- a/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.spec.ts +++ b/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.spec.ts @@ -106,8 +106,11 @@ describe('McpClientsService.decryptHeaders', () => { describe('McpClientsService.guardedFetch (SSRF per-request guard)', () => { // The bound guardedFetch closure lives on the instance as a private field. + // #489 split it into per-transport HTTP/SSE bindings (they differ only in the + // dispatcher's bodyTimeout); the SSRF guard is identical, so testing the HTTP + // one is sufficient. const guardedFetchOf = (service: McpClientsService) => - (service as unknown as { guardedFetch: typeof fetch }).guardedFetch; + (service as unknown as { guardedFetchHttp: typeof fetch }).guardedFetchHttp; let fetchSpy: jest.SpiedFunction; diff --git a/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts b/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts index 1e688974..20397acc 100644 --- a/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts +++ b/apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts @@ -1,5 +1,6 @@ import { isIP } from 'node:net'; import { lookup as dnsLookup, type LookupAddress } from 'node:dns'; +import { pathToFileURL } from 'node:url'; import { Injectable, Logger } from '@nestjs/common'; import { type Tool, type ToolCallOptions } from 'ai'; import { createMCPClient } from '@ai-sdk/mcp'; @@ -10,9 +11,29 @@ import { streamingDispatcherOptions, mcpStreamTimeoutMs, mcpCallTimeoutMs, + mcpSseBodyTimeoutMs, } from '../../../integrations/ai/ai-streaming-fetch'; import { SecretBoxService } from '../../../integrations/crypto/secret-box'; import { isUrlAllowed, isIpAllowed } from './ssrf-guard'; +// TYPE-ONLY (erased at compile): @docmost/mcp is ESM-only and cannot be a runtime +// `require()` from this commonjs module (same constraint as docmost-client.loader). +// The write-class MAP is loaded lazily via the dynamic-import trick below. +import type { ToolWriteClass } from '@docmost/mcp'; + +// TS(commonjs) downlevels a literal `import()` to `require()`, which cannot load +// the ESM-only @docmost/mcp. Indirect through Function so the real dynamic +// `import()` survives compilation (same trick as docmost-client.loader.ts). +const esmImport = new Function( + 'specifier', + 'return import(specifier)', +) as (specifier: string) => Promise; + +/** Local read-only predicate — avoids a value import of the ESM-only package. + * Only a pure read is retry-safe after a transport break (a write is + * indeterminate). Kept in lockstep with @docmost/mcp's isRetryableWriteClass. */ +function isReadOnlyWriteClass(writeClass: ToolWriteClass | undefined): boolean { + return writeClass === 'readOnly'; +} /** A closable external MCP client handle. */ export interface Closable { @@ -81,12 +102,52 @@ const MAX_TOOL_NAME_LENGTH = 64; * close until the turn releases it, so a TTL expiry mid-turn never closes a * client a stream is still executing against. */ +/** + * Where a merged (namespaced) tool came from, so the per-run recovery wrapper + * (#489) can, on a transport error, reconnect THAT server and re-resolve the SAME + * underlying tool by its raw name. `writeClass` gates the single auto-retry (a + * read is retry-safe; a write is indeterminate). `serverIndex` indexes the + * entry's `servers` array (which server config to reconnect). + */ +interface ToolProvenance { + serverIndex: number; + rawName: string; + writeClass: ToolWriteClass | undefined; +} + +/** A live reconnected server (its fresh client + raw call-timeout-wrapped tools). */ +interface RecoveredServerState { + client: McpClient; + tools: Record; +} + +/** + * Per-run, per-server recovery binding (#489). `current` is the server's LIVE + * target for this run: `null` means "use the ORIGINAL cached client/template"; + * a non-null value is a reconnected throwaway client all this server's tools now + * call. `reconnecting` dedupes concurrent reconnects so only ONE fresh client is + * minted per death (a losing concurrent call awaits it and retries on the SAME + * new client — the CAS-by-identity rule). + */ +interface ServerBinding { + current: RecoveredServerState | null; + reconnecting?: Promise; +} + interface CacheEntry { tools: Record; clients: McpClient[]; outcomes: ServerOutcome[]; /** Prompt guidance for qualifying servers (see McpServerInstruction). */ instructions: McpServerInstruction[]; + /** + * The enabled server configs used to build this entry (#489), so the per-run + * recovery wrapper can reconnect a specific server by index. Parallel to the + * indices referenced by {@link toolMeta}. + */ + servers: AiMcpServer[]; + /** merged-tool-key -> provenance (#489), for the per-run recovery wrapper. */ + toolMeta: Record; expiresAt: number; /** Active leases (turns currently using these clients). */ refCount: number; @@ -120,20 +181,82 @@ export class McpClientsService { */ private readonly cache = new Map>(); /** - * A single shared SSRF-pinned dispatcher for ALL outbound external-MCP fetches. - * Its custom connect.lookup runs per connection, so one instance safely guards - * every server's connections (we never connect to an unvalidated IP). + * SSRF-pinned dispatchers for outbound external-MCP fetches. Both use the SAME + * custom connect.lookup (so every connection is IP-validated), but carry a + * DIFFERENT `bodyTimeout` (#489): the HTTP (streamable) transport opens a fresh + * request per call, so it keeps the tight silence timeout; the SSE transport + * holds ONE long-lived body open across many calls, so a >1-min idle BETWEEN + * calls is LEGITIMATE and must not break the socket — it gets a much larger + * bodyTimeout. (headersTimeout stays tight on both.) */ - private readonly dispatcher: Dispatcher = buildPinnedDispatcher(); - /** guardedFetch bound to the pinned dispatcher; reused by every transport. */ - private readonly guardedFetch: typeof fetch = (input, init) => - guardedFetch(this.dispatcher, input, init); + private readonly dispatcherHttp: Dispatcher = buildPinnedDispatcher( + mcpStreamTimeoutMs(), + ); + private readonly dispatcherSse: Dispatcher = buildPinnedDispatcher( + mcpSseBodyTimeoutMs(), + ); + /** guardedFetch bound to each dispatcher; picked by transport type in connect(). */ + private readonly guardedFetchHttp: typeof fetch = (input, init) => + guardedFetch(this.dispatcherHttp, input, init); + private readonly guardedFetchSse: typeof fetch = (input, init) => + guardedFetch(this.dispatcherSse, input, init); + + /** + * Memoized write-class map (#489), loaded lazily from @docmost/mcp via the + * dynamic-import trick. Keyed by tool name (=== mcpName). A tool NOT in the map + * (any third-party external MCP tool) classifies as `undefined` -> treated as a + * write by the retry gate (the safe default: never blind-retry an unknown tool). + * On any load failure the map is `{}` (every tool -> no auto-retry), so a + * missing/older @docmost/mcp build only DISABLES retries, never mis-retries. + */ + private writeClassMapPromise: Promise> | null = + null; constructor( private readonly repo: AiMcpServerRepo, private readonly secretBox: SecretBoxService, ) {} + /** + * Whether an external MCP server is the TRUSTED internal Docmost MCP server — + * the only server whose tools may be classified by the Docmost write-class map + * (#489 review). Today this is ALWAYS false: every `ai_mcp_servers` row is an + * admin-configured THIRD-PARTY endpoint (there is no builtin/self flag, sentinel + * URL, or synthetic server in this path — Docmost's OWN tools are exposed via the + * separate in-app tools path, never through this external-MCP client). So no + * third-party tool can inherit `readOnly` by a name collision with a Docmost read + * tool, and none is ever auto-retried on a transport error (which would risk a + * double-apply — the #435 class). Flip this (an explicit `kind`/`isBuiltin` + * column, or a configured self-MCP URL) if a trusted internal server is ever + * introduced. A method (not a free function) so it is a single, mockable seam. + */ + private isInternalDocmostServer(_server: AiMcpServer): boolean { + return false; + } + + /** Lazily load + memoize the shared write-class map (see the field doc). */ + private getWriteClassMap(): Promise> { + if (!this.writeClassMapPromise) { + this.writeClassMapPromise = (async () => { + try { + const entry = require.resolve('@docmost/mcp'); + const mod = (await esmImport(pathToFileURL(entry).href)) as { + SHARED_TOOL_WRITE_CLASS?: Record; + }; + return mod.SHARED_TOOL_WRITE_CLASS ?? {}; + } catch (err) { + this.logger.warn( + `Could not load MCP write-class map (auto-retry disabled): ${shortError( + err, + )}`, + ); + return {}; + } + })(); + } + return this.writeClassMapPromise; + } + /** * Build (or reuse a cached) external toolset for a workspace. Returns the * merged tools, the open client handles to release, and per-server outcomes. @@ -162,11 +285,37 @@ export class McpClientsService { } }, }; - // One release handle drives the whole leased entry; closing it releases all - // underlying clients together (they share the same lease lifecycle). + + // #489: the run accumulates a SET of leases — the primary cache lease PLUS any + // throwaway client minted by an in-run transport-recovery reconnect. They are + // NEVER released mid-run (releasing a swapped-out client while a concurrent + // in-flight call still holds it would INDUCE a second failure); the caller + // releases the WHOLE set together at turn-end. A recovery reconnect pushes its + // lease onto this live array, which the consumer closes over. + const leaseSet: Closable[] = [release]; + + // #489: per-RUN transport-recovery binding, one per server, SHARED by all of + // that server's tools so a swap by one call is seen by the next (CAS by + // identity). Kept per-run (here, not in the cached entry) because the binding + // + lease-set state is per-run. + const bindings = new Map(); + const capMs = mcpCallTimeoutMs(); + + // Wrap each cached tool with the recovery layer. On a transport error a + // declared readOnly tool reconnects its server and retries ONCE; a write is + // never blind-retried (indeterminate — may have applied before the reset). A + // tool without provenance (a minimal stub entry in a test) passes through raw. + const tools: Record = {}; + for (const [key, tool] of Object.entries(entry.tools)) { + const meta = entry.toolMeta?.[key]; + tools[key] = meta + ? this.wrapWithTransportRecovery(entry, meta, tool, leaseSet, bindings, capMs) + : tool; + } + return { - tools: entry.tools, - clients: [release], + tools, + clients: leaseSet, outcomes: entry.outcomes, instructions: entry.instructions, }; @@ -254,6 +403,16 @@ export class McpClientsService { // Per-call total wall-clock cap, read once for this build (env-overridable). const callTimeoutMs = mcpCallTimeoutMs(); const instructions: McpServerInstruction[] = []; + // merged-key -> provenance for the per-run recovery wrapper (#489). + const toolMeta: Record = {}; + // Shared Docmost write-class map (#489) — classifies a tool by its raw name. + // Loaded ONLY when at least one server is a TRUSTED internal Docmost server + // (see isInternalDocmostServer): for third-party servers the map is never + // applied (a name collision must not grant readOnly-retry), so we skip the + // dynamic ESM load entirely in that (currently universal) case. + const writeClassMap = servers.some((s) => this.isInternalDocmostServer(s)) + ? await this.getWriteClassMap() + : null; // Per-server connect+tools result, still tagged with its server so the merge // below can be applied in the SAME order as `servers` (see the parallel note). @@ -327,11 +486,23 @@ export class McpClientsService { // against names already merged from earlier servers, so no external // tool is silently overwritten on collision. The returned count drives // whether this server's prompt guidance is included (≥1 tool merged). + // #489 (review): the Docmost write-class map keys by DOCMOST tool names and + // may ONLY be trusted for a server KNOWN to be the internal Docmost MCP + // server. Every row here is an admin-configured THIRD-PARTY endpoint, so a + // third-party WRITE tool that happens to be named like a Docmost read + // (getPage, listPages, ...) must NOT inherit readOnly — that would auto-retry + // a mutation on a transport error (double-apply, the #435 class). Gate the + // map on the trust check; untrusted servers get writeClass=undefined -> the + // recovery wrapper treats them as writes and never auto-retries. + const trustWriteClass = this.isInternalDocmostServer(server); const merged = this.mergeNamespaced( tools, result.guarded, server.name, server.id, + toolMeta, + i, + trustWriteClass ? writeClassMap : null, ); outcomes.push({ name: server.name, ok: true }); // Include this server's guidance ONLY when it actually contributed at @@ -353,6 +524,8 @@ export class McpClientsService { clients, outcomes, instructions, + servers, + toolMeta, expiresAt: Date.now() + CACHE_TTL_MS, refCount: 0, evicted: false, @@ -379,18 +552,33 @@ export class McpClientsService { picked: Record, serverName: string, serverId: string, + toolMeta: Record, + serverIndex: number, + // The Docmost write-class map, or `null` for an UNTRUSTED (third-party) + // server whose tools must all default to write (never auto-retried). + writeClassMap: Record | null, ): { count: number; prefix: string } { let count = 0; - for (const [name, tool] of Object.entries(namespace(picked, serverName))) { - let key = name; + for (const { full, raw, tool } of namespace(picked, serverName)) { + let key = full; if (key in target) { const original = key; - key = disambiguate(name, serverId, (candidate) => candidate in target); + key = disambiguate(full, serverId, (candidate) => candidate in target); this.logger.debug( `External MCP tool name "${original}" collided; renamed to "${key}"`, ); } target[key] = tool; + // Record provenance so the per-run recovery wrapper (#489) can reconnect + // this tool's server and re-resolve it by its raw name. writeClass is set + // ONLY from a TRUSTED (internal-Docmost) map; for a third-party server the + // map is null -> writeClass stays undefined -> the wrapper treats the tool + // as a write and never auto-retries it (no double-apply on name collision). + toolMeta[key] = { + serverIndex, + rawName: raw, + writeClass: writeClassMap ? writeClassMap[raw] : undefined, + }; count += 1; } return { count, prefix: namespacePrefix(serverName) }; @@ -424,7 +612,10 @@ export class McpClientsService { // Defense in depth: re-validate the actual request host on EVERY fetch // AND pin the socket to a validated IP via the dispatcher's connect // lookup, closing the DNS-rebinding TOCTOU between check and connect. - fetch: this.guardedFetch, + // #489: the SSE transport uses the raised-bodyTimeout dispatcher (idle + // between calls is legit); HTTP uses the tight one. + fetch: + transportType === 'sse' ? this.guardedFetchSse : this.guardedFetchHttp, }, })) as unknown as McpClient; return client; @@ -505,6 +696,176 @@ export class McpClientsService { } } + /** + * Wrap one merged external tool with the per-run transport-recovery layer (#489). + * + * attempt 1 runs on the server's CURRENT binding (the cached client, or a client + * a sibling tool already reconnected this run). On a REAL transport error + * (undici/@ai-sdk socket/body-timeout shapes — {@link isRetryableConnectError}, + * NOT a mock) and ONLY for a declared readOnly tool, it reconnects the server + * and retries EXACTLY ONCE on the fresh client; a write is surfaced as an + * indeterminate error (it may have applied before the reset — never + * blind-retried). A single per-call cap bounds BOTH attempts + the reconnect, + * and the run's abort signal is checked before the retry AND before minting a + * fresh connection (no connection is opened for a stopped run). + */ + private wrapWithTransportRecovery( + entry: CacheEntry, + meta: ToolProvenance, + template: Tool, + leaseSet: Closable[], + bindings: Map, + capMs: number, + ): Tool { + const original = template.execute; + if (typeof original !== 'function') return template; + const service = this; + const { serverIndex, rawName, writeClass } = meta; + + let binding = bindings.get(serverIndex); + if (!binding) { + binding = { current: null }; + bindings.set(serverIndex, binding); + } + const boundBinding = binding; + + const execute = async (args: unknown, options: ToolCallOptions) => { + // The per-call cap governs the WHOLE sequence (attempt1 + reconnect + + // attempt2). Compose it with the run's abort signal so a Stop or the cap + // ends any awaited call — @ai-sdk/mcp does not settle on abort, so we RACE. + const capController = new AbortController(); + const capTimer = setTimeout(() => { + capController.abort(new Error(`MCP tool call timed out after ${capMs}ms`)); + }, capMs); + capTimer.unref?.(); + const runSignal = options?.abortSignal; + const composed = runSignal + ? AbortSignal.any([runSignal, capController.signal]) + : capController.signal; + const stopped = () => runSignal?.aborted === true || capController.signal.aborted; + + const callOn = async ( + exec: NonNullable, + ): Promise => { + const aborted = new Promise((_, reject) => { + const fail = () => reject(abortReason(composed)); + if (composed.aborted) fail(); + else composed.addEventListener('abort', fail, { once: true }); + }); + return Promise.race([exec(args, { ...options, abortSignal: composed }), aborted]); + }; + + const execFor = ( + state: RecoveredServerState | null, + ): NonNullable | undefined => + state ? (state.tools[rawName]?.execute as NonNullable) : original; + + try { + // Snapshot the target BEFORE the call so a swap by a concurrent call is + // detected by identity in the catch. + const attemptState = boundBinding.current; + const attemptExec = execFor(attemptState); + if (typeof attemptExec !== 'function') { + throw new Error(`external MCP tool "${rawName}" is not callable`); + } + try { + return await callOn(attemptExec); + } catch (err) { + // Never retry on a Stop or an exhausted cap. + if (stopped()) throw err; + // Only a genuine transport break is a recovery candidate. + if (!isRetryableConnectError(err)) throw err; + // A write tool is INDETERMINATE on a transport error (may have applied + // before the reset) — surface that; do NOT auto-retry (double-apply is + // the #435 incident class). + if (!isReadOnlyWriteClass(writeClass)) { + throw new Error( + `external MCP tool "${rawName}" hit a transport error and MAY have already ` + + `applied on the server — not retried automatically; verify state before ` + + `retrying. (${shortError(err)})`, + ); + } + // Abort check BEFORE minting a fresh connection (no socket for a + // stopped run). LIMITATION (#489, LOW): the reconnect's own connect is + // bounded by CONNECT_TIMEOUT_MS but does NOT itself observe `composed`, + // so a Stop that lands DURING the handshake is only honored at the next + // `stopped()` gate (before the retry) — a bounded ≤5s late-abort window; + // the throwaway client is closed at turn-end regardless. Threading + // `composed` into the SHARED (CAS-deduped) reconnect is deliberately + // avoided: it would let the first caller's abort tear down a reconnect a + // concurrent still-live caller depends on. + if (stopped()) throw err; + // CAS-swap by IDENTITY: mint+swap only if nobody swapped since this + // call's snapshot; a losing concurrent call awaits the same reconnect + // and retries on the SAME fresh client. + let target: RecoveredServerState; + if (boundBinding.current === attemptState) { + if (!boundBinding.reconnecting) { + boundBinding.reconnecting = (async () => { + const server = entry.servers[serverIndex]; + const fresh = await service.reconnectServer(server, capMs); + leaseSet.push(fresh.lease); // accumulate; released at turn-end + boundBinding.current = fresh.state; + return fresh.state; + })(); + // Clear the in-flight marker once it settles (success or failure) so + // a LATER death of the new client can reconnect again. + void boundBinding.reconnecting.then( + () => (boundBinding.reconnecting = undefined), + () => (boundBinding.reconnecting = undefined), + ); + } + target = await boundBinding.reconnecting; + } else { + target = boundBinding.current as RecoveredServerState; + } + // Abort check BEFORE the retry. + if (stopped()) throw err; + const retryExec = execFor(target); + if (typeof retryExec !== 'function') throw err; + return await callOn(retryExec); + } + } finally { + clearTimeout(capTimer); + } + }; + return { ...template, execute } as unknown as Tool; + } + + /** + * Reconnect ONE server for an in-run recovery (#489): open a fresh client and + * list+wrap its tools. The throwaway client is NOT cached — it is owned by the + * RUN via the returned lease (closed at turn-end), independent of the shared + * cache entry (whose TTL rebuild heals future turns). On a failure the fresh + * client is closed so its socket never leaks. + */ + private async reconnectServer( + server: AiMcpServer, + capMs: number, + ): Promise<{ state: RecoveredServerState; lease: Closable }> { + const client = await this.connectWithTimeout(server, CONNECT_TIMEOUT_MS); + let tools: Record; + try { + const raw = await withTimeout(client.tools(), CONNECT_TIMEOUT_MS); + const allow = server.toolAllowlist; + const picked = + Array.isArray(allow) && allow.length > 0 ? pick(raw, allow) : raw; + tools = wrapToolsWithCallTimeout(picked, capMs); + } catch (err) { + void client.close().catch(() => undefined); + throw err; + } + let released = false; + const lease: Closable = { + close: async () => { + if (released) return; + released = true; + await client.close().catch(() => undefined); + }, + }; + return { state: { client, tools }, lease }; + } + /** Mark an entry evicted; close its clients now if nothing is leasing them. */ private evict(entry: CacheEntry): void { clearTimeout(entry.timer); @@ -554,22 +915,21 @@ export function validateResolvedAddresses(addrs: readonly LookupAddress[]): { * certificate validation still uses the real hostname (we never rewrite the URL * to an IP literal). */ -function buildPinnedDispatcher(): Agent { - // External-MCP traffic uses a DEDICATED, shorter silence timeout +function buildPinnedDispatcher(bodyTimeoutMs: number): Agent { + // External-MCP traffic uses a DEDICATED, shorter HEADERS silence timeout // (`AI_MCP_STREAM_TIMEOUT_MS`, default 1 min) — deliberately tighter than the // chat provider's 15-min `streamTimeoutMs()` — so a byte-silent/hung MCP // upstream is broken in ~1 min instead of 15. We keep the keep-alive options - // from `streamingDispatcherOptions()` but OVERRIDE headers/body timeouts. - // Accepted trade-off: a legitimately long but byte-silent single tool call, - // and an SSE transport idling >1 min BETWEEN tool calls, are also cut here; the - // per-call total cap (wrapToolsWithCallTimeout, `AI_MCP_CALL_TIMEOUT_MS`) is the - // complementary guard for chatty-but-stuck calls that keep the socket warm yet - // never return. - const mcpSilenceMs = mcpStreamTimeoutMs(); + // from `streamingDispatcherOptions()` but OVERRIDE the timeouts. `bodyTimeout` + // is passed in per-transport (#489): tight for HTTP (fresh request per call), + // raised for SSE (one long-lived body across calls — idle BETWEEN calls is + // legit). The per-call total cap (`AI_MCP_CALL_TIMEOUT_MS`) is the complementary + // guard for chatty-but-stuck calls that keep the socket warm yet never return. + const headersMs = mcpStreamTimeoutMs(); return new Agent({ ...streamingDispatcherOptions(), - headersTimeout: mcpSilenceMs, - bodyTimeout: mcpSilenceMs, + headersTimeout: headersMs, + bodyTimeout: bodyTimeoutMs, connect: { lookup: (hostname, _options, callback) => { // Always resolve ALL addresses ourselves; do not trust the caller's @@ -669,18 +1029,22 @@ function pick( function namespace( tools: Record, serverName: string, -): Record { +): Array<{ full: string; raw: string; tool: Tool }> { const prefix = namespacePrefix(serverName); - const out: Record = {}; + const out: Array<{ full: string; raw: string; tool: Tool }> = []; + const taken: Record = {}; for (const [name, t] of Object.entries(tools)) { const safe = sanitizeName(name); let full = capName(`${prefix}_${safe}`); // Duplicate names within ONE server can still collide after sanitize/ // truncate — suffix-disambiguate so the second tool is not overwritten. - if (full in out) { - full = disambiguate(full, '', (candidate) => candidate in out); + if (full in taken) { + full = disambiguate(full, '', (candidate) => candidate in taken); } - out[full] = t; + taken[full] = true; + // Keep the RAW (un-namespaced) name alongside the merged key so the per-run + // recovery wrapper (#489) can re-resolve the same tool on a fresh client. + out.push({ full, raw: name, tool: t }); } return out; } @@ -804,6 +1168,69 @@ export function wrapToolWithCallTimeout(tool: Tool, ms: number): Tool { return { ...tool, execute } as unknown as Tool; } +/** + * undici / Node network error CODES that mean the connection broke (not an + * application-level error) — a transient transport failure a readOnly call may + * safely retry after reconnecting. Matched against the REAL error shapes (#489): + * a socket reset surfaces as `TypeError: fetch failed` whose `.cause` is an + * undici `SocketError { code:'UND_ERR_SOCKET' }`; a body-timeout as + * `TypeError: terminated` whose `.cause` is `BodyTimeoutError`. Classifying by + * these real codes/names (not by mock errors) is essential — a mock-shaped + * predicate would leave eviction silently dead in production while CI is green. + */ +const RETRYABLE_TRANSPORT_ERROR_CODES: ReadonlySet = new Set([ + 'ECONNRESET', + 'ECONNREFUSED', + 'ECONNABORTED', + 'EPIPE', + 'ETIMEDOUT', + 'EAI_AGAIN', + 'ENETUNREACH', + 'EHOSTUNREACH', + 'UND_ERR_SOCKET', + 'UND_ERR_CONNECT_TIMEOUT', + 'UND_ERR_HEADERS_TIMEOUT', + 'UND_ERR_BODY_TIMEOUT', + 'UND_ERR_CLOSED', + 'UND_ERR_DESTROYED', +]); + +/** undici error CLASS names for the same transport-break conditions. */ +const RETRYABLE_TRANSPORT_ERROR_NAMES: ReadonlySet = new Set([ + 'SocketError', + 'BodyTimeoutError', + 'HeadersTimeoutError', + 'ConnectTimeoutError', + 'ClientClosedError', + 'ClientDestroyedError', +]); + +/** + * Whether `err` is a retryable TRANSPORT break (a broken socket / body timeout), + * classified by the REAL undici/@ai-sdk error shapes (#489). undici surfaces a + * reset as `TypeError('fetch failed'|'terminated')` with the real error in + * `.cause`, and @ai-sdk/mcp may wrap it again in an `MCPClientError` (cause + * chain), so we walk `.cause` (bounded depth) checking `.code` and `.name`. An + * app-level tool error (a 4xx, a validation failure) is NOT retryable and returns + * false — only a connection-level failure heals with a reconnect. + */ +export function isRetryableConnectError(err: unknown, depth = 0): boolean { + if (!err || typeof err !== 'object' || depth > 6) return false; + const e = err as { + code?: unknown; + name?: unknown; + cause?: unknown; + }; + if (typeof e.code === 'string' && RETRYABLE_TRANSPORT_ERROR_CODES.has(e.code)) { + return true; + } + if (typeof e.name === 'string' && RETRYABLE_TRANSPORT_ERROR_NAMES.has(e.name)) { + return true; + } + if (e.cause != null) return isRetryableConnectError(e.cause, depth + 1); + return false; +} + /** The signal's reason as an Error (informative thrown value on abort/timeout). */ function abortReason(signal: AbortSignal): Error { const r = signal.reason; diff --git a/apps/server/src/integrations/ai/ai-streaming-fetch.ts b/apps/server/src/integrations/ai/ai-streaming-fetch.ts index b4a31723..9f76112c 100644 --- a/apps/server/src/integrations/ai/ai-streaming-fetch.ts +++ b/apps/server/src/integrations/ai/ai-streaming-fetch.ts @@ -129,6 +129,12 @@ const DEFAULT_MCP_STREAM_TIMEOUT_MS = 60_000; /** Default total wall-clock cap for ONE external MCP tool call (2 min). */ const DEFAULT_MCP_CALL_TIMEOUT_MS = 120_000; +/** + * Default `bodyTimeout` for the EXTERNAL-MCP SSE transport (10 min) — #489. + * Deliberately much LARGER than {@link DEFAULT_MCP_STREAM_TIMEOUT_MS}. + */ +const DEFAULT_MCP_SSE_BODY_TIMEOUT_MS = 600_000; + /** * SILENCE timeout (ms) for EXTERNAL-MCP transport ONLY. Override with * `AI_MCP_STREAM_TIMEOUT_MS`; a missing/invalid/non-positive value falls back to @@ -164,6 +170,26 @@ export function mcpCallTimeoutMs(): number { return positiveEnv('AI_MCP_CALL_TIMEOUT_MS', DEFAULT_MCP_CALL_TIMEOUT_MS); } +/** + * `bodyTimeout` (ms) for the EXTERNAL-MCP **SSE** transport ONLY — #489. Override + * with `AI_MCP_SSE_BODY_TIMEOUT_MS`; a missing/invalid/non-positive value falls + * back to {@link DEFAULT_MCP_SSE_BODY_TIMEOUT_MS} (10 min). + * + * The SSE transport holds ONE long-lived response body open across many tool + * calls, so undici's `bodyTimeout` (time between body bytes) counts the LEGITIMATE + * silence BETWEEN calls, not just a hung single call. At the tight HTTP silence + * timeout ({@link mcpStreamTimeoutMs}, 1 min) a normal >1-min gap between the + * model's tool calls would break the SSE socket, and the cache would then serve a + * dead client until TTL. So the SSE transport gets its OWN, RAISED bodyTimeout; + * the per-call total cap ({@link mcpCallTimeoutMs}) still bounds a single stuck + * call, and the app-level transport-error retry heals a socket that does break. + * The HTTP (streamable) transport keeps the tight timeout — it opens a fresh + * request per call, so idle-between-calls does not apply there. + */ +export function mcpSseBodyTimeoutMs(): number { + return positiveEnv('AI_MCP_SSE_BODY_TIMEOUT_MS', DEFAULT_MCP_SSE_BODY_TIMEOUT_MS); +} + /** * undici `Agent` options for streaming AI traffic — the (generous, finite) * silence timeouts plus the keep-alive recycle window. Shared by the chat diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 15fed6ac..298c88d5 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -30,6 +30,13 @@ export { destroyAllSessions } from "./lib/collab-session.js"; // internals directly; it goes through loadDocmostMcp()). export { SHARED_TOOL_SPECS } from "./tool-specs.js"; export type { SharedToolSpec } from "./tool-specs.js"; +// #489 — write-class registry consumed by the in-app external-MCP retry gate. +export { + SHARED_TOOL_WRITE_CLASS, + isRetryableWriteClass, + assertEverySpecDeclaresWriteClass, +} from "./tool-specs.js"; +export type { ToolWriteClass } from "./tool-specs.js"; // Re-export the build-time REGISTRY_STAMP (issue #447): a deterministic hash of // the tool-specs registry content, generated into src/registry-stamp.generated.ts diff --git a/packages/mcp/src/tool-specs.ts b/packages/mcp/src/tool-specs.ts index 941eead9..0bc9fd24 100644 --- a/packages/mcp/src/tool-specs.ts +++ b/packages/mcp/src/tool-specs.ts @@ -127,6 +127,18 @@ export interface SharedToolSpec { mcpName: string; /** camelCase key in the ai-SDK tools object (the in-app layer). */ inAppKey: string; + /** + * Write-class of the tool (#489), declared on EVERY spec (a registration-time + * assert enforces completeness; `satisfies Record` + * makes it a compile error to omit). 'readOnly' = a pure read that mutates + * NOTHING durable, so it is safe to auto-retry once after a transport break. + * 'write' = anything that mutates a page/comment/share/diagram/etc — a + * transport error is INDETERMINATE (the server may have applied it before the + * connection reset), so it is NEVER blind-retried (a retry would double-apply, + * the #435 incident class). Consumed by the external-MCP retry path + * (mcp-clients.service.ts) to gate its single auto-retry. + */ + writeClass: 'readOnly' | 'write'; /** Single canonical model-facing description used by both layers. */ description: string; /** @@ -240,6 +252,7 @@ export const SHARED_TOOL_SPECS = { getWorkspace: { mcpName: 'getWorkspace', inAppKey: 'getWorkspace', + writeClass: 'readOnly', description: 'Fetch metadata about the current workspace (name, settings).', tier: 'core', catalogLine: 'getWorkspace — fetch current workspace metadata (name, settings).', @@ -249,6 +262,7 @@ export const SHARED_TOOL_SPECS = { listSpaces: { mcpName: 'listSpaces', inAppKey: 'listSpaces', + writeClass: 'readOnly', description: 'List the spaces the current user can access. Returns the array of ' + 'spaces (id, name, slug, ...).', @@ -260,6 +274,7 @@ export const SHARED_TOOL_SPECS = { listShares: { mcpName: 'listShares', inAppKey: 'listShares', + writeClass: 'readOnly', description: 'List all public shares in the workspace with page titles and public URLs.', tier: 'deferred', @@ -272,6 +287,7 @@ export const SHARED_TOOL_SPECS = { getPageJson: { mcpName: 'getPageJson', inAppKey: 'getPageJson', + writeClass: 'readOnly', description: 'Get page details with the raw ProseMirror JSON content (lossless: ' + 'includes block ids, callouts, tables, link/image attributes) plus the ' + @@ -289,6 +305,7 @@ export const SHARED_TOOL_SPECS = { getOutline: { mcpName: 'getOutline', inAppKey: 'getOutline', + writeClass: 'readOnly', description: "Return a COMPACT outline of a page's top-level blocks ({index, type, " + 'id, level, firstText}; tables add rows/cols/header; lists add item ' + @@ -309,6 +326,7 @@ export const SHARED_TOOL_SPECS = { getNode: { mcpName: 'getNode', inAppKey: 'getNode', + writeClass: 'readOnly', description: "Fetch a single block for editing. `nodeId` is a block id from the page " + 'outline or page-JSON view (works for headings/paragraphs/callouts/images), OR ' + @@ -350,6 +368,7 @@ export const SHARED_TOOL_SPECS = { searchInPage: { mcpName: 'searchInPage', inAppKey: 'searchInPage', + writeClass: 'readOnly', description: 'Find every occurrence of a string (or regex) INSIDE one page and get ' + 'WHERE each is — instead of pulling blocks one-by-one with getNode. ' + @@ -413,6 +432,7 @@ export const SHARED_TOOL_SPECS = { deleteNode: { mcpName: 'deleteNode', inAppKey: 'deleteNode', + writeClass: 'write', description: 'Remove a single block by its attrs.id (from the page outline or ' + 'page-JSON view) WITHOUT resending the whole document.', @@ -438,6 +458,7 @@ export const SHARED_TOOL_SPECS = { patchNode: { mcpName: 'patchNode', inAppKey: 'patchNode', + writeClass: 'write', description: 'Replace a single content block identified by its attrs.id, WITHOUT ' + 'resending the whole document; the replacement keeps the same block id. ' + @@ -505,6 +526,7 @@ export const SHARED_TOOL_SPECS = { insertNode: { mcpName: 'insertNode', inAppKey: 'insertNode', + writeClass: 'write', description: 'Insert content before/after another block (by attrs.id or anchor text) ' + 'or append it at the end (top level). For before/after you MUST provide ' + @@ -597,6 +619,7 @@ export const SHARED_TOOL_SPECS = { sharePage: { mcpName: 'sharePage', inAppKey: 'sharePage', + writeClass: 'write', // CANONICAL: merges the MCP copy's URL-format + idempotency detail with the // in-app copy's reversibility note; keeps the security framing both had. description: @@ -626,6 +649,7 @@ export const SHARED_TOOL_SPECS = { unsharePage: { mcpName: 'unsharePage', inAppKey: 'unsharePage', + writeClass: 'write', description: 'Remove the public share of a page (revokes the public URL).', tier: 'deferred', catalogLine: "unsharePage — revoke a page's public share (removes the public URL).", @@ -640,6 +664,7 @@ export const SHARED_TOOL_SPECS = { diffPageVersions: { mcpName: 'diffPageVersions', inAppKey: 'diffPageVersions', + writeClass: 'readOnly', description: 'Diff two versions of a page and return a Docmost-equivalent change set ' + '(inserted/deleted text, integrity counts for images/links/tables/' + @@ -672,6 +697,7 @@ export const SHARED_TOOL_SPECS = { listPageHistory: { mcpName: 'listPageHistory', inAppKey: 'listPageHistory', + writeClass: 'readOnly', description: "List a page's saved versions (Docmost auto-snapshots on every save), " + 'newest first, cursor-paginated. Returns { items, nextCursor }; each ' + @@ -693,6 +719,7 @@ export const SHARED_TOOL_SPECS = { restorePageVersion: { mcpName: 'restorePageVersion', inAppKey: 'restorePageVersion', + writeClass: 'write', description: 'Restore a page to a saved version: writes that version\'s content back ' + 'as the page\'s current content (Docmost has no restore endpoint, so ' + @@ -713,6 +740,7 @@ export const SHARED_TOOL_SPECS = { importPageMarkdown: { mcpName: 'importPageMarkdown', inAppKey: 'importPageMarkdown', + writeClass: 'write', // IN-APP ONLY (issue #411): the external /mcp surface no longer exposes // importPageMarkdown — the registry loop in index.ts skips inAppOnly specs, // so this stays available to the in-app agent (round-tripping an EXPORTED @@ -742,6 +770,7 @@ export const SHARED_TOOL_SPECS = { copyPageContent: { mcpName: 'copyPageContent', inAppKey: 'copyPageContent', + writeClass: 'write', description: "Replace targetPageId's content with a copy of sourcePageId's content, " + 'entirely server-side — the document is NOT sent through the model. The ' + @@ -770,6 +799,7 @@ export const SHARED_TOOL_SPECS = { editPageText: { mcpName: 'editPageText', inAppKey: 'editPageText', + writeClass: 'write', description: "Surgical find/replace inside a page's text, preserving all block " + 'ids and marks. A find MAY cross bold/italic/link boundaries; the ' + @@ -819,6 +849,7 @@ export const SHARED_TOOL_SPECS = { stashPage: { mcpName: 'stashPage', inAppKey: 'stashPage', + writeClass: 'readOnly', description: 'Serialize a whole page (the full ProseMirror JSON, as getPageJson ' + 'returns) into an ephemeral in-memory blob and return ONLY a short ' + @@ -880,6 +911,7 @@ export const SHARED_TOOL_SPECS = { getPage: { mcpName: 'getPage', inAppKey: 'getPage', + writeClass: 'readOnly', description: 'Fetch a single page as Markdown by its id. Returns the page title and ' + 'its Markdown content. The converter is canonical (round-trips text and ' + @@ -919,6 +951,7 @@ export const SHARED_TOOL_SPECS = { listPages: { mcpName: 'listPages', inAppKey: 'listPages', + writeClass: 'readOnly', description: 'List the most recent pages (ordered by updatedAt, descending), ' + 'optionally scoped to a single space. Returns a bounded list (default ' + @@ -965,6 +998,7 @@ export const SHARED_TOOL_SPECS = { getTree: { mcpName: 'getTree', inAppKey: 'getTree', + writeClass: 'readOnly', description: "Get a space's page hierarchy (or one subtree) as a nested tree in a " + 'SINGLE request — completely and without loss. Each node is ' + @@ -1009,6 +1043,7 @@ export const SHARED_TOOL_SPECS = { getPageContext: { mcpName: 'getPageContext', inAppKey: 'getPageContext', + writeClass: 'readOnly', description: 'Given a pageId, get its LOCATION and immediate surroundings (metadata ' + 'only, no page content) in one call — answers "where am I / what is ' + @@ -1038,6 +1073,7 @@ export const SHARED_TOOL_SPECS = { createPage: { mcpName: 'createPage', inAppKey: 'createPage', + writeClass: 'write', description: 'Create a new page with a Markdown body in a space, optionally under a ' + 'parent page (omit parentPageId to create at the space root). Returns ' + @@ -1088,6 +1124,7 @@ export const SHARED_TOOL_SPECS = { movePage: { mcpName: 'movePage', inAppKey: 'movePage', + writeClass: 'write', description: 'Move a page under a new parent page, or to the space root when no ' + 'parent is given. Reversible: move it back at any time.', @@ -1181,6 +1218,7 @@ export const SHARED_TOOL_SPECS = { renamePage: { mcpName: 'renamePage', inAppKey: 'renamePage', + writeClass: 'write', description: 'Rename a page (change its title only; the body is untouched, never ' + 'resent). Reversible: rename back at any time.', @@ -1202,6 +1240,7 @@ export const SHARED_TOOL_SPECS = { deletePage: { mcpName: 'deletePage', inAppKey: 'deletePage', + writeClass: 'write', description: 'Move a page to the trash — SOFT delete only: the page can be restored ' + 'from trash and nothing is ever permanently deleted.', @@ -1235,6 +1274,7 @@ export const SHARED_TOOL_SPECS = { updatePageJson: { mcpName: 'updatePageJson', inAppKey: 'updatePageJson', + writeClass: 'write', description: "Replace a page's content with a raw ProseMirror JSON document (lossless " + 'write: preserves the block ids, callouts, tables and attributes you pass ' + @@ -1291,6 +1331,7 @@ export const SHARED_TOOL_SPECS = { updatePageMarkdown: { mcpName: 'updatePageMarkdown', inAppKey: 'updatePageMarkdown', + writeClass: 'write', description: "Replace a page's body with new Markdown content (and optionally its " + 'title). The whole body is re-imported from the markdown (block ids ' + @@ -1325,6 +1366,7 @@ export const SHARED_TOOL_SPECS = { exportPageMarkdown: { mcpName: 'exportPageMarkdown', inAppKey: 'exportPageMarkdown', + writeClass: 'readOnly', // CANONICAL: the MCP copy (a strict superset of the terse in-app wording). description: 'Export a page to a single self-contained Docmost-flavoured Markdown ' + @@ -1373,6 +1415,7 @@ export const SHARED_TOOL_SPECS = { createComment: { mcpName: 'createComment', inAppKey: 'createComment', + writeClass: 'write', // CANONICAL: the in-app copy (the more-maintained one). It keeps the same // rules as the MCP copy — inline-only, top-level requires a `selection`, no // page-level comments, replies inherit the anchor, suggestedText must be @@ -1505,6 +1548,7 @@ export const SHARED_TOOL_SPECS = { listComments: { mcpName: 'listComments', inAppKey: 'listComments', + writeClass: 'readOnly', // CANONICAL: the two copies are near-identical; the MCP copy is the // superset (it keeps the "(pagination is handled internally)" note the // in-app copy dropped), so it is used verbatim. @@ -1532,6 +1576,7 @@ export const SHARED_TOOL_SPECS = { resolveComment: { mcpName: 'resolveComment', inAppKey: 'resolveComment', + writeClass: 'write', // CANONICAL: the MCP copy's richer wording, minus its reference // to `deleteComment` (a sibling tool that does NOT exist in the in-app // layer) — rephrased transport-neutrally per the registry convention. @@ -1572,6 +1617,7 @@ export const SHARED_TOOL_SPECS = { checkNewComments: { mcpName: 'checkNewComments', inAppKey: 'checkNewComments', + writeClass: 'readOnly', // CANONICAL: the MCP copy (the more detailed of the two). The MCP layer's // execute-side guard that rejects an unparseable `since` timestamp stays in // its execute body (per-layer logic), not in the shared schema. @@ -1651,6 +1697,7 @@ export const SHARED_TOOL_SPECS = { tableInsertRow: { mcpName: 'tableInsertRow', inAppKey: 'tableInsertRow', + writeClass: 'write', description: 'Insert a row of plain-text cells into a table. `table` is `#` ' + 'from the page outline, or a block id inside it. `cells` is the text per ' + @@ -1684,6 +1731,7 @@ export const SHARED_TOOL_SPECS = { tableDeleteRow: { mcpName: 'tableDeleteRow', inAppKey: 'tableDeleteRow', + writeClass: 'write', description: 'Delete the row at 0-based `index` from a table (`table` is `#` ' + 'from the page outline, or a block id inside it). Refuses to delete the ' + @@ -1707,6 +1755,7 @@ export const SHARED_TOOL_SPECS = { tableUpdateCell: { mcpName: 'tableUpdateCell', inAppKey: 'tableUpdateCell', + writeClass: 'write', description: 'Set the plain-text content of cell [row, col] (0-based) in a table ' + '(`table` is `#` from the page outline, or a block id inside it). ' + @@ -1747,6 +1796,7 @@ export const SHARED_TOOL_SPECS = { insertFootnote: { mcpName: 'insertFootnote', inAppKey: 'insertFootnote', + writeClass: 'write', description: 'Insert an AUTHOR-INLINE footnote: you specify only WHERE (anchorText) ' + 'and WHAT (text). The footnote marker is placed right after anchorText in ' + @@ -1784,6 +1834,7 @@ export const SHARED_TOOL_SPECS = { insertImage: { mcpName: 'insertImage', inAppKey: 'insertImage', + writeClass: 'write', description: 'Download an image from a web (http/https) URL and insert it into ' + 'a page in one step. By default ' + @@ -1828,6 +1879,7 @@ export const SHARED_TOOL_SPECS = { replaceImage: { mcpName: 'replaceImage', inAppKey: 'replaceImage', + writeClass: 'write', description: 'Replace an existing image on a page with a new image fetched from a web ' + '(http/https) URL: uploads the new file as a NEW ' + @@ -1866,6 +1918,7 @@ export const SHARED_TOOL_SPECS = { drawioGet: { mcpName: 'drawioGet', inAppKey: 'drawioGet', + writeClass: 'readOnly', description: 'Read a draw.io diagram on a page as mxGraph XML (default) or as its raw ' + '`.drawio.svg`. `node` is the drawio node\'s attrs.id (from getOutline / ' + @@ -1899,6 +1952,7 @@ export const SHARED_TOOL_SPECS = { drawioCreate: { mcpName: 'drawioCreate', inAppKey: 'drawioCreate', + writeClass: 'write', description: 'Create a draw.io diagram from mxGraph XML and insert it as a diagram ' + 'block. `xml` is a bare `` OR a list of `` elements ' + @@ -1969,6 +2023,7 @@ export const SHARED_TOOL_SPECS = { drawioUpdate: { mcpName: 'drawioUpdate', inAppKey: 'drawioUpdate', + writeClass: 'write', description: 'Replace a draw.io diagram\'s content with new mxGraph XML (same lint ' + 'pipeline as drawioCreate). `baseHash` is MANDATORY: pass the hash from ' + @@ -2020,6 +2075,7 @@ export const SHARED_TOOL_SPECS = { drawioEditCells: { mcpName: 'drawioEditCells', inAppKey: 'drawioEditCells', + writeClass: 'write', description: 'Make TARGETED, id-based edits to an existing draw.io diagram instead of ' + 'resending the whole XML (a full-XML diff is fragile — draw.io reorders ' + @@ -2078,6 +2134,7 @@ export const SHARED_TOOL_SPECS = { drawioFromGraph: { mcpName: 'drawioFromGraph', inAppKey: 'drawioFromGraph', + writeClass: 'write', description: 'Build a draw.io diagram from a SEMANTIC graph — you describe nodes, groups ' + 'and edges by MEANING and the server picks every coordinate, color and icon ' + @@ -2202,6 +2259,7 @@ export const SHARED_TOOL_SPECS = { drawioFromMermaid: { mcpName: 'drawioFromMermaid', inAppKey: 'drawioFromMermaid', + writeClass: 'write', description: 'Convert Mermaid `flowchart` text into an EDITABLE draw.io diagram (LLMs ' + 'write Mermaid reliably). Best for STANDARD flowcharts/decision trees: ' + @@ -2248,6 +2306,7 @@ export const SHARED_TOOL_SPECS = { drawioShapes: { mcpName: 'drawioShapes', inAppKey: 'drawioShapes', + writeClass: 'readOnly', description: 'Look up VERIFIED draw.io stencil style-strings so you never guess a ' + '`shape=mxgraph.*` name (a wrong name renders as an EMPTY BOX). Searches a ' + @@ -2294,6 +2353,7 @@ export const SHARED_TOOL_SPECS = { drawioGuide: { mcpName: 'drawioGuide', inAppKey: 'drawioGuide', + writeClass: 'readOnly', description: 'Progressive-disclosure draw.io authoring reference. Call with a `section` ' + 'to pull one focused, <=4KB chapter instead of bloating context: ' + @@ -2323,3 +2383,53 @@ export const SHARED_TOOL_SPECS = { inlineBothHosts: true, }, } satisfies Record; + +// --- write-class registry (#489) ------------------------------------------ + +/** A tool's retry-safety class. 'readOnly' may be auto-retried once after a + * transport break; 'write' is indeterminate and must never be blind-retried. */ +export type ToolWriteClass = 'readOnly' | 'write'; + +/** + * Name → write-class map for the shared registry, keyed by mcpName (=== inAppKey). + * The external-MCP retry path (mcp-clients.service.ts) looks a tool up here by its + * RAW (un-namespaced) name to decide whether a transport failure may be retried. + * A tool NOT in this map (a third-party external MCP tool) is treated as 'write' + * by the consumer — the safe default (never blind-retry an unknown tool). + */ +export const SHARED_TOOL_WRITE_CLASS: Record = + Object.fromEntries( + Object.values(SHARED_TOOL_SPECS).map((spec) => [spec.mcpName, spec.writeClass]), + ); + +/** Whether a write-class permits a single automatic retry after a transport + * break. Only a pure read is retry-safe; everything mutating is indeterminate. */ +export function isRetryableWriteClass( + writeClass: ToolWriteClass | undefined, +): boolean { + return writeClass === 'readOnly'; +} + +/** + * Registration-time assert (#489): EVERY spec must declare a valid write-class. + * `satisfies Record` already makes an omission a compile + * error, but this guards a raw/cast construction path and documents the invariant + * at the point of use. Runs once on import — both hosts import this module, so + * both get the check. Throws (fails startup) rather than silently mis-gating a + * retry in production. + */ +export function assertEverySpecDeclaresWriteClass(): void { + for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) { + const wc = (spec as SharedToolSpec).writeClass; + if (wc !== 'readOnly' && wc !== 'write') { + throw new Error( + `tool-specs: spec "${key}" must declare writeClass ('readOnly' | 'write'), got ${JSON.stringify( + wc, + )}`, + ); + } + } +} + +// Enforce at module load (registration time) on both hosts. +assertEverySpecDeclaresWriteClass(); diff --git a/packages/mcp/test/unit/tool-specs.test.mjs b/packages/mcp/test/unit/tool-specs.test.mjs index 41f16606..aad5354a 100644 --- a/packages/mcp/test/unit/tool-specs.test.mjs +++ b/packages/mcp/test/unit/tool-specs.test.mjs @@ -2,7 +2,12 @@ import { test } from "node:test"; import assert from "node:assert/strict"; import { z } from "zod"; -import { SHARED_TOOL_SPECS } from "../../build/tool-specs.js"; +import { + SHARED_TOOL_SPECS, + SHARED_TOOL_WRITE_CLASS, + isRetryableWriteClass, + assertEverySpecDeclaresWriteClass, +} from "../../build/tool-specs.js"; // The shared registry is consumed by BOTH the zod-v3 MCP server and the zod-v4 // in-app AI-SDK service, so every spec must carry the cross-layer wiring @@ -43,6 +48,41 @@ test("mcpName and inAppKey are each unique across the registry", () => { } }); +// #489 — every spec must declare its write-class so the external-MCP retry path +// can gate a single auto-retry ONLY on a pure read (a blind retry of a write = +// double-apply). The declaration is enforced at registration time. +test("#489: every spec declares a valid writeClass ('readOnly' | 'write')", () => { + for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) { + assert.ok( + spec.writeClass === "readOnly" || spec.writeClass === "write", + `${key}: missing/invalid writeClass: ${JSON.stringify(spec.writeClass)}`, + ); + } + // The registration-time assert must not throw for the shipped registry. + assert.doesNotThrow(() => assertEverySpecDeclaresWriteClass()); +}); + +test("#489: SHARED_TOOL_WRITE_CLASS maps every mcpName to its class; helper gates on readOnly", () => { + const specs = Object.values(SHARED_TOOL_SPECS); + assert.equal(Object.keys(SHARED_TOOL_WRITE_CLASS).length, specs.length); + for (const spec of specs) { + assert.equal(SHARED_TOOL_WRITE_CLASS[spec.mcpName], spec.writeClass); + } + // Only a readOnly tool is retry-eligible; a write tool and an unknown tool are not. + assert.equal(isRetryableWriteClass("readOnly"), true); + assert.equal(isRetryableWriteClass("write"), false); + assert.equal(isRetryableWriteClass(undefined), false); +}); + +test("#489: representative reads are readOnly and representative writes are write", () => { + for (const name of ["getPage", "getTree", "searchInPage", "listComments"]) { + assert.equal(SHARED_TOOL_SPECS[name].writeClass, "readOnly", `${name} should be readOnly`); + } + for (const name of ["patchNode", "createPage", "deletePage", "createComment", "drawioCreate"]) { + assert.equal(SHARED_TOOL_SPECS[name].writeClass, "write", `${name} should be write`); + } +}); + test("buildShape (when present) returns a usable ZodRawShape with a real zod", () => { for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) { if (!spec.buildShape) continue;