From a075593fc4267b078ac26b50474b000dcbe4e2c9 Mon Sep 17 00:00:00 2001 From: Ali Ibrahim Jr <48456829+IBJunior@users.noreply.github.com> Date: Mon, 9 Nov 2026 11:58:33 +0100 Subject: [PATCH] fix(mcp): keep loaded tools when their server goes away, and never stop silently Deleting an MCP server cascade-deleted its catalog row, which is where the disclosure middleware rebuilt loaded tools from each turn. The tools vanished from the model's list, the model called one from memory with no schema, and the unparseable call ended the turn with no error shown. - Disclosure saves each loaded tool's definition in the thread (`toolDefinitions`, first definition wins) and builds the loaded tools from it, so the tool list never shrinks. A call to a tool whose server was removed or turned off returns `server_unavailable` instead of falling through to the tool node. Threads from before keep the catalog fallback. - The stream also takes tool results from the run's state, so a result a middleware returned itself (server_unavailable, auth_required, unknown_tool) arrives live and is no longer reported as pending approval. - invalidToolCallMiddleware re-asks the model when a tool call's arguments don't parse, with a correction never saved to the thread (invalid_arguments, or unknown_tool for a tool not offered). After two retries the reply ends with a visible notice, and the stream shows the notice live. Invalid calls are read from both `invalid_tool_calls` and v1 `invalid_tool_call` content blocks. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 18 ++- docs/MCP_DISCLOSURE.md | 32 ++++- packages/mcp-disclosure/src/catalogTools.ts | 30 ++-- .../mcp-disclosure/src/middleware.test.ts | 99 +++++++++++-- packages/mcp-disclosure/src/middleware.ts | 134 +++++++++++++++--- packages/mcp-disclosure/src/types.ts | 10 ++ src/lib/agent/index.ts | 4 +- src/lib/agent/invalidCallNotice.ts | 28 ++++ src/lib/agent/invalidToolCalls.test.ts | 121 ++++++++++++++++ src/lib/agent/invalidToolCalls.ts | 111 +++++++++++++++ src/services/messageStream.test.ts | 113 +++++++++++++++ src/services/messageStream.ts | 72 ++++++++-- 12 files changed, 704 insertions(+), 68 deletions(-) create mode 100644 src/lib/agent/invalidCallNotice.ts create mode 100644 src/lib/agent/invalidToolCalls.test.ts create mode 100644 src/lib/agent/invalidToolCalls.ts diff --git a/AGENTS.md b/AGENTS.md index def4807..7dc2dba 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -390,7 +390,9 @@ cameron_analytics` migration, and `analytics.test.ts` fails if the two differ. M delegates to the SDK's `auth()` — never hardcode a scope in the client metadata. - **Progressive disclosure** ([docs/MCP_DISCLOSURE.md](docs/MCP_DISCLOSURE.md), read it before changing this code): the model sees a short list of capabilities plus `load_tools`, and a tool's - definition only once a search has loaded it; nothing connects until a tool is called. The logic + definition only once a search has loaded it; nothing connects until a tool is called. A loaded + tool's definition is saved in the thread (`toolDefinitions`), so it stays offered after its + server is removed or turned off, and a call to it returns `server_unavailable`. The logic is the workspace package `packages/mcp-disclosure`, which must not import app code or the MCP client library; `src/lib/mcp/connector.ts` is the only file using the adapters. @@ -466,6 +468,20 @@ cameron_analytics` migration, and `analytics.test.ts` fails if the two differ. M client-only artifact. Waiting whenever an artifact is merely absent is indistinguishable from waiting for one that will never come — it parked the pump on every ordinary tool and pushed all results to the end of the turn. `messageStream.test.ts` pins both this and the pending contract. +- **A result a middleware returned itself is streamed from the state.** The disclosure middleware + answers some calls without running the tool node (`server_unavailable`, `auth_required`, + `unknown_tool`), so they never appear on `run.toolCalls`. `streamMessages` also takes results + from `run.values` snapshots and the final state, for calls this run requested, each sent once + (`pushResult`). Without it the result showed only after a reload, and the call was reported as + `pendingToolCallIds`, the approval signal. +- **A malformed tool call is retried, never dropped silently.** Arguments that don't parse land in + `invalid_tool_calls` (or, in LangChain's v1 output format, only as `invalid_tool_call` content + blocks: read both, via `findInvalidCalls`). The agent never runs such a call, so with no valid + calls the turn just ended. `invalidToolCallMiddleware` (`src/lib/agent/invalidToolCalls.ts`, + installed just inside disclosure so it sees the full tool list) re-asks the model up to + `MAX_INVALID_CALL_RETRIES` times with a correction (`invalid_arguments`, or `unknown_tool` for a + tool not offered) that is never saved to the thread; after that the response ends with a + visible notice. The stream shows the same notice live for every malformed attempt. - **Dedup and React keys are `type:id`, not `id`.** Ids are unique only within a type (an AI message's synthetic id and a tool message's call id are different namespaces). diff --git a/docs/MCP_DISCLOSURE.md b/docs/MCP_DISCLOSURE.md index 587e444..7533897 100644 --- a/docs/MCP_DISCLOSURE.md +++ b/docs/MCP_DISCLOSURE.md @@ -151,7 +151,8 @@ load_tools({ query: string; server?: string; limit?: number }); // limit default re-sent, and it can never be removed; it also loses real tool names and `tool_choice`. **Why nothing is removed.** Changing the tool list costs a prompt-cache miss, so loaded tools stay -for the thread and `activeTools` keeps load order, leaving the existing array as a stable prefix. +for the thread — even after their server is removed or turned off, since their definitions live in +the thread, not the catalog (deleting a server cascade-deletes its catalog row) and `activeTools` keeps load order, leaving the existing array as a stable prefix. Each load still costs one miss on that prefix: measure it with the cache attributes in [OBSERVABILITY.md](OBSERVABILITY.md). A thread only grows by tools the model asked for. @@ -185,13 +186,18 @@ interrupt carrying the authorization URL, with a Connect button in the chat) is No MCP tool is registered with `createAgent`. Instead: -1. **Stand-ins from the catalog.** `catalogTools()` builds one tool per cached entry - (`server_tool`, its description, and its `inputSchema` exactly as the server listed it). `wrapModelCall` adds the loaded stand-ins to `request.tools` — - LangChain accepts tools added there as long as a `wrapToolCall` runs them. A stand-in throws if it is ever run directly. +1. **Stand-ins from the thread's saved definitions.** When `load_tools` loads a tool, its + definition (`server_tool`, its description, and its `inputSchema` exactly as the server listed + it) is saved in the thread's `toolDefinitions`. `wrapModelCall` builds a stand-in from each + (`standInTool()`) and adds them to `request.tools` — LangChain accepts tools added there as long + as a `wrapToolCall` runs them. A stand-in throws if it is ever run directly. A thread from + before definitions were saved has names only, and falls back to the catalog's stand-ins. 2. **The real tool on call.** `wrapToolCall` asks `resolveTool(serverId, toolName)` for the real tool and passes it on as `request.tool`, which the tool node runs; result formatting stays the client library's. A tool the catalog does not list is refused (`unknown_tool`) without - connecting. A connection failure comes back as an error result (`connect_failed`, + connecting; with no catalog this request, the live server decides. A loaded tool whose server + was removed or turned off is answered with `server_unavailable` (naming the server and the + Connections page), never passed on to the tool node, which knows no MCP tools. A connection failure comes back as an error result (`connect_failed`, `auth_required`, `timeout`) and does not load the tool. 3. **Auto-load.** A call to a catalog tool that was not loaded still runs it and loads it (`autoLoadOnCall`), which only happens when the model guesses a name. @@ -245,8 +251,20 @@ Code defaults in `packages/mcp-disclosure/src/types.ts`; the pool's idle timeout ## State (per thread, via the checkpointer) `activeTools: string[]` holds `server_tool` names in load order, through an appending, de-duplicating -reducer, so two `load_tools` calls in one step both apply. A tool whose server was disabled or -removed is left out when the list is built; no state change is needed. +reducer, so two `load_tools` calls in one step both apply. + +`toolDefinitions` maps each loaded name to `{ server, tool }`, the definition as it was loaded. Its +reducer keeps the **first** definition saved for a name, so what the model is shown never changes, +even if the server later lists a different schema (the real tool still validates the call against +its current one). `load_tools` also back-fills a definition for a tool loaded before this field +existed. Neither field is ever pruned. + +Why it is in the state and not the catalog: a tool removed from the model's list mid-thread is +still in the conversation, so the model calls it from memory with no schema. That produced an +unparseable call (an `invalid_tool_call`) which the agent does not run, ending the turn with no +error shown. Such calls are now retried by the app's `invalidToolCallMiddleware` (see +`AGENTS.md`, Streaming Architecture), which tells the model a tool it isn't offered is not +available. ## The catalog (`mcp_server_catalog`) diff --git a/packages/mcp-disclosure/src/catalogTools.ts b/packages/mcp-disclosure/src/catalogTools.ts index 7726a60..9279098 100644 --- a/packages/mcp-disclosure/src/catalogTools.ts +++ b/packages/mcp-disclosure/src/catalogTools.ts @@ -1,6 +1,6 @@ import { DynamicStructuredTool, type StructuredToolInterface } from "@langchain/core/tools"; import { prefixedToolName } from "./policy"; -import type { DisclosureServer } from "./types"; +import type { DisclosureServer, McpToolDefinition } from "./types"; export class CatalogToolNotRunnable extends Error {} @@ -16,20 +16,22 @@ export function catalogTools(servers: DisclosureServer[]): Map { - const name = prefixedToolName(server.name, t.name); - return new DynamicStructuredTool({ - name, - description: t.description ?? "", - schema: t.inputSchema as Record, - func: async () => { - throw new CatalogToolNotRunnable( - `${name} is a catalog stand-in; it must be run through the disclosure middleware.`, - ); - }, - }); - }), + server.catalog.tools.map((t) => standInTool(prefixedToolName(server.name, t.name), t)), ); } return byServer; } + +/** One stand-in. The same definition always yields the same tool the model sees. */ +export function standInTool(name: string, definition: McpToolDefinition): StructuredToolInterface { + return new DynamicStructuredTool({ + name, + description: definition.description ?? "", + schema: definition.inputSchema as Record, + func: async () => { + throw new CatalogToolNotRunnable( + `${name} is a catalog stand-in; it must be run through the disclosure middleware.`, + ); + }, + }); +} diff --git a/packages/mcp-disclosure/src/middleware.test.ts b/packages/mcp-disclosure/src/middleware.test.ts index aa8d06c..11e7816 100644 --- a/packages/mcp-disclosure/src/middleware.test.ts +++ b/packages/mcp-disclosure/src/middleware.test.ts @@ -13,11 +13,14 @@ import type { DisclosureOptions, DisclosureServer, ResolveMcpTool } from "./type type State = { activeTools: string[] }; type Step = { tools?: [string, Record][]; text?: string }; +let modelSeq = 0; + /** Plays a script of tool calls and records what each model call was shown. */ class ScriptedModel extends BaseChatModel { seen: { tools: string[]; system: string }[] = []; private bound: string[] = []; private n = 0; + private seq = ++modelSeq; constructor(private script: Step[]) { super({}); @@ -37,7 +40,8 @@ class ScriptedModel extends BaseChatModel { const system = messages.find((m) => m.getType() === "system"); this.seen.push({ tools: [...this.bound], system: system?.text ?? "" }); const step = this.script.shift() ?? { text: "done" }; - const id = `m${++this.n}`; + // Unique across models, so a second agent on the same thread appends rather than replaces. + const id = `s${this.seq}-m${++this.n}`; const message = new AIMessage({ id, content: step.text ?? "", @@ -312,9 +316,10 @@ describe("disclosure", () => { expect(result).toMatchObject({ ok: false, error: "tool_not_loaded" }); }); - it("drops a loaded tool whose server is no longer available", async () => { - const checkpointer = new MemorySaver(); - const make = (using: DisclosureServer[], model: ScriptedModel) => + describe("a loaded tool outlives its server", () => { + const config = { configurable: { thread_id: "t1" } }; + const input = { messages: [{ role: "user", content: "hi" }] }; + const make = (checkpointer: MemorySaver, using: DisclosureServer[], model: ScriptedModel) => createAgent({ model, tools: [fakeTool("log_expense")], @@ -327,15 +332,85 @@ describe("disclosure", () => { }), ], }); - const config = { configurable: { thread_id: "t1" } }; - const input = { messages: [{ role: "user", content: "hi" }] }; - const first = new ScriptedModel([{ tools: [["load_tools", { query: "notes" }]] }]); - await make(servers, first).invoke(input, config); - expect(first.seen.at(-1)!.tools).toContain("notes_search_notes"); + const loadNotes = async (checkpointer: MemorySaver) => { + const first = new ScriptedModel([{ tools: [["load_tools", { query: "notes" }]] }]); + await make(checkpointer, servers, first).invoke(input, config); + expect(first.seen.at(-1)!.tools).toContain("notes_search_notes"); + }; + const lastToolResult = (out: { messages: BaseMessage[] }) => + JSON.parse(String(out.messages.filter((m) => m.getType() === "tool").at(-1)!.content)); + + it("keeps showing it after the server is removed, and saves its definition", async () => { + const checkpointer = new MemorySaver(); + await loadNotes(checkpointer); + const state = (await make(checkpointer, servers, new ScriptedModel([])).getState( + config, + )) as unknown as { values: { toolDefinitions: Record } }; + expect(state.values.toolDefinitions.notes_search_notes).toEqual({ + server: "notes", + tool: expect.objectContaining({ name: "search_notes", description: "fake search_notes" }), + }); + + // Removed (or turned off): the server is no longer in the list the agent is built with. + for (const using of [[servers[0]], []]) { + const next = new ScriptedModel([]); + await make(checkpointer, using, next).invoke(input, config); + expect(next.seen[0].tools).toContain("notes_search_notes"); + } + }); - const second = new ScriptedModel([]); - await make([servers[0]], second).invoke(input, config); - expect(second.seen[0].tools).not.toContain("notes_search_notes"); + it("answers a call to it with server_unavailable, without connecting", async () => { + const checkpointer = new MemorySaver(); + await loadNotes(checkpointer); + connected.length = 0; + const next = new ScriptedModel([{ tools: [["notes_search_notes", {}]] }]); + const out = await make(checkpointer, [servers[0]], next).invoke(input, config); + expect(lastToolResult(out)).toMatchObject({ + ok: false, + error: "server_unavailable", + server: "notes", + }); + expect(connected).toEqual([]); + }); + + it("handles a thread saved before definitions were, by name only", async () => { + const checkpointer = new MemorySaver(); + const agent = make(checkpointer, servers, new ScriptedModel([])); + await agent.invoke(input, config); + // The old state shape: a loaded name, no saved definition. + await ( + agent as unknown as { updateState: (c: object, v: object) => Promise } + ).updateState(config, { activeTools: ["notes_search_notes"] }); + + // Server still there: shown from the catalog, as before. + const present = new ScriptedModel([]); + await make(checkpointer, servers, present).invoke(input, config); + expect(present.seen[0].tools).toContain("notes_search_notes"); + + // Server gone: nothing to show it from, but a call still gets a clear answer. + const gone = new ScriptedModel([{ tools: [["notes_search_notes", {}]] }]); + const out = await make(checkpointer, [servers[0]], gone).invoke(input, config); + expect(gone.seen[0].tools).not.toContain("notes_search_notes"); + expect(lastToolResult(out)).toMatchObject({ + error: "server_unavailable", + tool: "notes_search_notes", + }); + }); + + it("keeps the definition first saved, even if the catalog changes", async () => { + const checkpointer = new MemorySaver(); + await loadNotes(checkpointer); + const changed = makeServer("notes", { toolNames: ["search_notes"] }); + changed.catalog!.tools[0].description = "a newer description"; + const next = new ScriptedModel([{ tools: [["load_tools", { query: "search_notes" }]] }]); + await make(checkpointer, [servers[0], changed], next).invoke(input, config); + const state = (await make(checkpointer, servers, new ScriptedModel([])).getState( + config, + )) as unknown as { values: { toolDefinitions: Record } }; + expect(state.values.toolDefinitions.notes_search_notes.tool).toMatchObject({ + description: "fake search_notes", + }); + }); }); }); diff --git a/packages/mcp-disclosure/src/middleware.ts b/packages/mcp-disclosure/src/middleware.ts index d005e62..a19192b 100644 --- a/packages/mcp-disclosure/src/middleware.ts +++ b/packages/mcp-disclosure/src/middleware.ts @@ -4,7 +4,7 @@ import type { StructuredToolInterface } from "@langchain/core/tools"; import { Command, ReducedValue, StateSchema, isCommand } from "@langchain/langgraph"; import { z } from "zod"; import { truncate } from "./catalog"; -import { catalogTools } from "./catalogTools"; +import { catalogTools, standInTool } from "./catalogTools"; import { catalogedServers, createToolServerIndex, @@ -19,6 +19,7 @@ import { type DisclosureOptions, type DisclosureServer, type ResolveMcpTool, + type SavedToolDefinition, } from "./types"; export const LOAD_TOOLS_TOOL = "load_tools"; @@ -31,16 +32,51 @@ function appendNew(current: string[], next: string[]): string[] { return [...new Set([...current, ...next])]; } -/** Persisted per thread by the checkpointer. A reducer, so two loads in one step both apply. */ +type SavedDefinitions = Record; + +/** The first definition saved for a name wins: what the model is shown never changes. */ +function keepFirst(current: SavedDefinitions, next: SavedDefinitions): SavedDefinitions { + return { ...next, ...current }; +} + +const savedDefinitionsSchema = z + .record( + z.string(), + z.object({ + server: z.string(), + tool: z.object({ + name: z.string(), + description: z.string().optional(), + inputSchema: z.unknown().optional(), + annotations: z.unknown().optional(), + }), + }), + ) + .default(() => ({})); + +/** Persisted per thread by the checkpointer. Reducers, so two loads in one step both apply. */ export const disclosureStateSchema = new StateSchema({ - /** `server__tool` names available to the model, in the order they were loaded. */ + /** `server_tool` names available to the model, in the order they were loaded. */ activeTools: new ReducedValue( z.array(z.string()).default(() => []), { inputSchema: z.array(z.string()).default(() => []), reducer: appendNew }, ), + /** + * Each loaded tool's definition, saved when it was loaded, so the tool stays in the list after + * its server is removed or turned off. Threads from before this field hold names only. + */ + toolDefinitions: new ReducedValue(savedDefinitionsSchema, { + inputSchema: savedDefinitionsSchema, + reducer: keepFirst, + }), }); -type DisclosureState = { activeTools?: string[] }; +type DisclosureState = { activeTools?: string[]; toolDefinitions?: SavedDefinitions }; + +const saved = (hit: SearchableTool): SavedToolDefinition => ({ + server: hit.server.name, + tool: hit.definition, +}); export interface DisclosureMiddlewareConfig { /** Every enabled server, with its catalog entry (`null` = unavailable for this request). */ @@ -95,7 +131,7 @@ export function planLoadTools( servers: DisclosureServer[], index: ToolSearchIndex, state: DisclosureState, -): { payload: Record; add: string[] } { +): { payload: Record; add: string[]; definitions: SavedDefinitions } { const available = catalogedServers(servers); const query = typeof args.query === "string" ? args.query : ""; const ref = typeof args.server === "string" ? args.server.trim() : ""; @@ -110,9 +146,10 @@ export function planLoadTools( availableServers: available.map((s) => s.name).sort(), }, add: [], + definitions: {}, }; } - if (scoped?.authRequired) return { payload: authRequired(scoped), add: [] }; + if (scoped?.authRequired) return { payload: authRequired(scoped), add: [], definitions: {} }; const exact = exactMatch(query, index, scoped ? [scoped] : available); const hits = exact ? [exact] : index.search(query, scoped?.id); @@ -133,12 +170,15 @@ export function planLoadTools( ), }, add: [], + definitions: {}, }; } const usable = hits.filter((h) => !h.server.authRequired); const blocked = [...new Set(hits.filter((h) => h.server.authRequired).map((h) => h.server))]; - if (usable.length === 0) return { payload: authRequired(blocked[0]), add: [] }; + if (usable.length === 0) { + return { payload: authRequired(blocked[0]), add: [], definitions: {} }; + } const active = new Set(state.activeTools ?? []); const top = usable.slice(0, asLimit(args.limit)); @@ -157,6 +197,10 @@ export function planLoadTools( : {}), }, add: fresh.map((h) => h.name), + // Also back-fills a tool loaded before definitions were saved. + definitions: Object.fromEntries( + top.filter((h) => !state.toolDefinitions?.[h.name]).map((h) => [h.name, saved(h)]), + ), }; } @@ -189,6 +233,20 @@ function buildLoadToolsTool(servers: DisclosureServer[], index: ToolSearchIndex) type ToolCallLike = { toolCall: { id?: string; name: string } }; +/** + * What the model is told when a loaded tool's server was removed or turned off. A thread from + * before definitions were saved knows only the tool's name, not its server's. + */ +const serverUnavailable = (server: string | undefined, tool: string) => ({ + ok: false, + error: "server_unavailable", + ...(server ? { server } : { tool }), + message: + `${server ?? `The server behind ${tool}`} is no longer available: it was removed or ` + + `turned off. Ask the user to add it back or turn it on in the Connections page, then try ` + + `again.`, +}); + const errorMessage = (request: ToolCallLike, body: object) => new ToolMessage({ content: json(body), @@ -204,7 +262,8 @@ const errorMessage = (request: ToolCallLike, body: object) => * No MCP tool is registered with the agent and nothing connects up front: loaded tools are * stand-ins built from the cached catalog, and a call runs on the real tool from `resolveTool`, * which connects to that one server on demand. Loaded tools are never removed, because changing - * the tool list costs a prompt-cache miss. + * the tool list costs a prompt-cache miss: their definitions are saved in the thread, so they + * outlive their server, and a call to one whose server is gone returns `server_unavailable`. */ export function mcpServerDisclosureMiddleware(config: DisclosureMiddlewareConfig) { const options: DisclosureOptions = { ...DEFAULT_DISCLOSURE_OPTIONS, ...config.options }; @@ -218,6 +277,14 @@ export function mcpServerDisclosureMiddleware(config: DisclosureMiddlewareConfig .map((t): [string, StructuredToolInterface] => [t.name, t]), ); const totalTools = standIns.size; + // Keyed by name and definition: a thread that saved an older schema keeps that one. + const savedStandIns = new Map(); + const standInFor = (name: string, definition: SavedToolDefinition) => { + const key = `${name}\u0000${JSON.stringify(definition.tool)}`; + let tool = savedStandIns.get(key); + if (!tool) savedStandIns.set(key, (tool = standInTool(name, definition.tool))); + return tool; + }; const catalogBlock = renderCatalogBlock(servers); /** Runs an MCP tool call on the real tool, connecting to its server if needed. */ @@ -229,7 +296,8 @@ export function mcpServerDisclosureMiddleware(config: DisclosureMiddlewareConfig server: DisclosureServer, ): Promise => { const name = request.toolCall.name.slice(server.name.length + TOOL_NAME_SEPARATOR.length); - if (!server.catalog?.tools.some((t) => t.name === name)) { + // With no catalog this request, the live server decides whether the tool exists. + if (server.catalog && !server.catalog.tools.some((t) => t.name === name)) { return errorMessage(request, { ok: false, error: "unknown_tool", @@ -261,17 +329,24 @@ export function mcpServerDisclosureMiddleware(config: DisclosureMiddlewareConfig tools: catalogBlock ? [buildLoadToolsTool(servers, index)] : [], wrapModelCall: (request, handler) => { - // A tool whose server was disabled or removed since loading is simply left out. - const loaded = (request.state.activeTools ?? []).flatMap((name) => standIns.get(name) ?? []); + // From the thread's saved definition, so removing a server doesn't shrink the list. A + // thread from before definitions were saved falls back to the catalog. + const definitions = request.state.toolDefinitions ?? {}; + const loaded = (request.state.activeTools ?? []).flatMap((name) => { + const definition = definitions[name]; + return definition ? [standInFor(name, definition)] : (standIns.get(name) ?? []); + }); log( `[mcp-disclosure] loaded=${loaded.length}/${totalTools} ` + `tools=${request.tools.length + loaded.length}`, ); - if (!catalogBlock) return handler(request); + if (!catalogBlock && loaded.length === 0) return handler(request); return handler({ ...request, tools: [...request.tools, ...loaded], - systemMessage: request.systemMessage.concat(`\n\n${catalogBlock}`), + ...(catalogBlock + ? { systemMessage: request.systemMessage.concat(`\n\n${catalogBlock}`) } + : {}), }); }, @@ -279,13 +354,29 @@ export function mcpServerDisclosureMiddleware(config: DisclosureMiddlewareConfig const name = request.toolCall.name; if (name === LOAD_TOOLS_TOOL) { const result = await handler(request); - const { add } = planLoadTools(request.toolCall.args, servers, index, request.state); - if (add.length === 0 || isCommand(result)) return result; - return new Command({ update: { activeTools: add, messages: [result] } }); + const { add, definitions } = planLoadTools( + request.toolCall.args, + servers, + index, + request.state, + ); + if ((add.length === 0 && Object.keys(definitions).length === 0) || isCommand(result)) { + return result; + } + return new Command({ + update: { activeTools: add, toolDefinitions: definitions, messages: [result] }, + }); } const server = serverOf(name); - if (!server) return handler(request); + if (!server) { + // A loaded tool whose server is gone: say so, rather than fall through to "unknown tool". + const definition = request.state.toolDefinitions?.[name]; + if (definition || (request.state.activeTools ?? []).includes(name)) { + return errorMessage(request, serverUnavailable(definition?.server, name)); + } + return handler(request); + } const loaded = (request.state.activeTools ?? []).includes(name); if (!loaded && !options.autoLoadOnCall) { @@ -299,7 +390,14 @@ export function mcpServerDisclosureMiddleware(config: DisclosureMiddlewareConfig const result = await runMcpTool(request, handler, server); if (loaded || isCommand(result) || result.status === "error") return result; log(`[mcp-disclosure] auto-loaded ${name} on call`); - return new Command({ update: { activeTools: [name], messages: [result] } }); + const hit = index.tool(name); + return new Command({ + update: { + activeTools: [name], + toolDefinitions: hit ? { [name]: saved(hit) } : {}, + messages: [result], + }, + }); }, }); } diff --git a/packages/mcp-disclosure/src/types.ts b/packages/mcp-disclosure/src/types.ts index e694e71..48eaaef 100644 --- a/packages/mcp-disclosure/src/types.ts +++ b/packages/mcp-disclosure/src/types.ts @@ -8,6 +8,16 @@ export interface McpToolDefinition { annotations?: unknown; } +/** + * A loaded tool as the thread saved it. The model keeps being shown this definition even after + * its server is removed or turned off, so the tool list never shrinks. + */ +export interface SavedToolDefinition { + /** The server's name when the tool was loaded; a call is routed by it. */ + server: string; + tool: McpToolDefinition; +} + /** What a catalog refresh reads from a live server. */ export interface ServerManifest { tools: McpToolDefinition[]; diff --git a/src/lib/agent/index.ts b/src/lib/agent/index.ts index abb6da2..6347112 100644 --- a/src/lib/agent/index.ts +++ b/src/lib/agent/index.ts @@ -16,6 +16,7 @@ import { skillTools } from "./tools/skills"; import { listSkills } from "@/lib/skills/registry"; import { createAgent, humanInTheLoopMiddleware } from "langchain"; import { mcpServerDisclosureMiddleware } from "@agentailor/mcp-disclosure"; +import { invalidToolCallMiddleware } from "./invalidToolCalls"; import { loadDisclosureServers } from "@/lib/mcp/catalog"; import { createToolResolver } from "@/lib/mcp/pool"; // Lives in ./capabilities so the page can read it without importing this file's deps. @@ -96,7 +97,8 @@ async function buildAgent(cfg?: AgentConfigOptions) { // Shows the model a list of capabilities and a tool once a search has loaded it, and runs a call // on a pooled connection. load_tools is not gated: it only changes visibility. const disclosure = mcpServerDisclosureMiddleware(mcp); - const middleware = [disclosure, ...approval]; + // Inside disclosure, so the retry sees the full tool list, loaded MCP tools included. + const middleware = [disclosure, invalidToolCallMiddleware(), ...approval]; const agent = createAgent({ model: llm, diff --git a/src/lib/agent/invalidCallNotice.ts b/src/lib/agent/invalidCallNotice.ts new file mode 100644 index 0000000..b9b5e1a --- /dev/null +++ b/src/lib/agent/invalidCallNotice.ts @@ -0,0 +1,28 @@ +export type InvalidCall = { id?: string; name?: string; error?: string }; + +type MessageLike = { invalid_tool_calls?: unknown[]; content?: unknown }; + +/** + * A message's tool calls whose arguments didn't parse. Read from both places LangChain puts them: + * the `invalid_tool_calls` field, and — in its v1 output format, which leaves that field empty — + * `invalid_tool_call` content blocks. + */ +export function findInvalidCalls(message: MessageLike | undefined): InvalidCall[] { + const fromBlocks = Array.isArray(message?.content) + ? (message.content as { type?: string }[]).filter((b) => b?.type === "invalid_tool_call") + : []; + const all = [...((message?.invalid_tool_calls ?? []) as InvalidCall[]), ...fromBlocks]; + const seen = new Set(); + return (all as InvalidCall[]).filter((c) => { + const key = c.id ?? `${c.name}:${c.error}`; + if (seen.has(key)) return false; + seen.add(key); + return true; + }); +} + +/** One line naming the tool calls that could not run because their arguments didn't parse. */ +export function invalidCallNotice(calls: { name?: string }[]): string { + const names = [...new Set(calls.map((c) => c.name || "a tool"))].join(", "); + return `(The call to ${names} had arguments that were not valid JSON, so it was not run.)`; +} diff --git a/src/lib/agent/invalidToolCalls.test.ts b/src/lib/agent/invalidToolCalls.test.ts new file mode 100644 index 0000000..c0aa9b2 --- /dev/null +++ b/src/lib/agent/invalidToolCalls.test.ts @@ -0,0 +1,121 @@ +import { describe, expect, it, vi } from "vitest"; +import { AIMessage, HumanMessage, ToolMessage, type BaseMessage } from "@langchain/core/messages"; +import { invalidToolCallMiddleware, MAX_INVALID_CALL_RETRIES } from "./invalidToolCalls"; + +const valid = (text = "done") => new AIMessage({ content: text }); + +const malformed = (extra: Partial[0] & object> = {}) => + new AIMessage({ + content: "I'll pull all your notes now.", + invalid_tool_calls: [ + { + id: "bad-1", + name: "notes_list_notes", + args: "", + error: "Failed to parse tool call arguments as JSON", + type: "invalid_tool_call", + }, + ], + ...extra, + }); + +type Request = { messages: BaseMessage[]; tools: { name: string }[] }; +type Wrap = (request: Request, handler: (r: Request) => Promise) => Promise; + +/** Runs the middleware's wrapModelCall against a scripted sequence of model responses. */ +async function run(responses: AIMessage[], tools = ["notes_list_notes", "run_sql"]) { + const wrap = invalidToolCallMiddleware().wrapModelCall as unknown as Wrap; + const handler = vi.fn(async (_request: Request) => responses.shift()!); + const request = { + messages: [new HumanMessage("can you list all my notes")], + tools: tools.map((name) => ({ name })), + }; + const result = await wrap(request, handler); + return { result, handler, request }; +} + +const resultsSentOn = (handler: ReturnType, call: number) => + (handler.mock.calls[call][0] as Request).messages + .filter((m): m is ToolMessage => ToolMessage.isInstance(m)) + .map((m) => ({ id: m.tool_call_id, ...JSON.parse(String(m.content)) })); + +describe("invalidToolCallMiddleware", () => { + it("passes a well-formed response through with one model call", async () => { + const { result, handler } = await run([valid()]); + expect(handler).toHaveBeenCalledTimes(1); + expect(result.content).toBe("done"); + }); + + it("re-asks with the parse error, and keeps only the repaired response", async () => { + const repaired = new AIMessage({ + content: "", + tool_calls: [{ id: "good-1", name: "notes_list_notes", args: {}, type: "tool_call" }], + }); + const { result, handler, request } = await run([malformed(), repaired]); + + expect(handler).toHaveBeenCalledTimes(2); + const retry = (handler.mock.calls[1][0] as Request).messages; + // The original messages, then the malformed step answered — never saved to the thread. + expect(retry.slice(0, request.messages.length)).toEqual(request.messages); + expect(retry.at(-2)).toMatchObject({ + tool_calls: [{ id: "bad-1", name: "notes_list_notes", args: {} }], + }); + expect(resultsSentOn(handler, 1)).toEqual([ + expect.objectContaining({ id: "bad-1", ok: false, error: "invalid_arguments" }), + ]); + expect(result).toBe(repaired); + expect(request.messages).toHaveLength(1); + }); + + it("finds a call that only appears as a v1 content block, and echoes back text only", async () => { + // LangChain's v1 output format: the malformed call is a content block, the field is empty. + const v1 = new AIMessage({ + content: [ + { type: "text", text: "I'll pull all your notes now." }, + { + type: "invalid_tool_call", + id: "bad-1", + name: "notes_list_notes", + args: "", + error: "Failed to parse tool call arguments as JSON", + }, + ] as unknown as string, + }); + expect(v1.invalid_tool_calls ?? []).toEqual([]); + const { handler } = await run([v1, valid()]); + + expect(handler).toHaveBeenCalledTimes(2); + const echoed = (handler.mock.calls[1][0] as Request).messages.at(-2) as AIMessage; + expect(echoed.content).toBe("I'll pull all your notes now."); + expect(resultsSentOn(handler, 1)).toEqual([ + expect.objectContaining({ id: "bad-1", error: "invalid_arguments" }), + ]); + }); + + it("says a tool isn't available when the model calls one it isn't offered", async () => { + // A tool remembered from earlier in the thread whose server is gone: no schema, so no args. + const { handler } = await run([malformed(), valid()], ["run_sql"]); + expect(resultsSentOn(handler, 1)).toEqual([ + expect.objectContaining({ id: "bad-1", error: "unknown_tool" }), + ]); + }); + + it("answers a valid sibling as not run, since the step is retried whole", async () => { + const mixed = malformed({ + tool_calls: [{ id: "ok-1", name: "run_sql", args: { query: "select 1" }, type: "tool_call" }], + }); + const { handler } = await run([mixed, valid()]); + expect(resultsSentOn(handler, 1).map((r) => [r.id, r.error])).toEqual([ + ["ok-1", "not_run"], + ["bad-1", "invalid_arguments"], + ]); + }); + + it("gives up after the retry limit with a visible notice instead of a silent stop", async () => { + const attempts = Array.from({ length: MAX_INVALID_CALL_RETRIES + 1 }, () => malformed()); + const { result, handler } = await run(attempts); + expect(handler).toHaveBeenCalledTimes(MAX_INVALID_CALL_RETRIES + 1); + expect(String(result.content)).toMatch(/^I'll pull all your notes now\.\n\n\(The call to/); + expect(String(result.content)).toContain("notes_list_notes"); + }); +}); diff --git a/src/lib/agent/invalidToolCalls.ts b/src/lib/agent/invalidToolCalls.ts new file mode 100644 index 0000000..5dc38c4 --- /dev/null +++ b/src/lib/agent/invalidToolCalls.ts @@ -0,0 +1,111 @@ +import { createMiddleware } from "langchain"; +import { AIMessage, ToolMessage, type BaseMessage } from "@langchain/core/messages"; +import { findInvalidCalls, invalidCallNotice, type InvalidCall } from "./invalidCallNotice"; + +/** Extra model calls allowed per step to repair malformed tool-call arguments. */ +export const MAX_INVALID_CALL_RETRIES = 2; + +/** + * The model's malformed step and the answer to each of its calls, sent back on the retry only. + * Every call is answered — the invalid ones with the parse error, the valid ones as not run — + * because the step is retried whole, and each provider requires an answer for every call. + */ +function correction( + response: AIMessage, + invalid: InvalidCall[], + offered: Set, +): BaseMessage[] { + const calls = [ + ...(response.tool_calls ?? []).map((c) => ({ ...c, valid: true })), + ...invalid.map((c, i) => ({ + id: c.id || `invalid-${i}`, + name: c.name || "unknown_tool", + args: {}, + valid: false, + error: c.error, + })), + ]; + // Text only: the response's own content can carry the malformed call as a block, and the + // provider would then see it twice. + const text = + typeof response.content === "string" + ? response.content + : response.content + .filter((b) => b.type === "text") + .map((b) => (b as { text?: string }).text ?? "") + .join(""); + return [ + new AIMessage({ + content: text, + tool_calls: calls.map(({ id, name, args }) => ({ id, name, args, type: "tool_call" })), + }), + ...calls.map( + (c) => + new ToolMessage({ + tool_call_id: c.id ?? "", + name: c.name, + status: "error", + content: JSON.stringify( + c.valid + ? { + ok: false, + error: "not_run", + message: "Not run: another call in this step was malformed. Request it again.", + } + : offered.has(c.name) + ? { + ok: false, + error: "invalid_arguments", + message: + `Your arguments for ${c.name} were not valid JSON` + + `${"error" in c && c.error ? ` (${c.error})` : ""}. ` + + `Call it again with a JSON object matching its schema.`, + } + : // Usually a tool remembered from earlier in the thread, no longer offered. + { + ok: false, + error: "unknown_tool", + message: + `${c.name} is not one of your tools right now, so it can't be called. ` + + `Don't call it again; if the user needs it, tell them it isn't available.`, + }, + ), + }), + ), + ]; +} + +function withNotice(response: AIMessage, notice: string): AIMessage { + const content = + typeof response.content === "string" + ? `${response.content}${response.content ? "\n\n" : ""}${notice}` + : [...response.content, { type: "text" as const, text: notice }]; + return new AIMessage({ ...response, content }); +} + +/** + * A tool call whose arguments don't parse lands in `invalid_tool_calls`, which the agent never + * runs; with no valid calls the run just ends, and the user sees no error. This re-asks the model + * with the parse error instead. The correction is never saved: only the final response joins the + * thread, so the approval gate and the history see one clean step. If it is still malformed after + * {@link MAX_INVALID_CALL_RETRIES}, the step ends with a visible notice. + */ +export function invalidToolCallMiddleware() { + return createMiddleware({ + name: "InvalidToolCallRetry", + wrapModelCall: async (request, handler) => { + const offered = new Set(request.tools.map((t) => (t as { name?: string }).name ?? "")); + let response = await handler(request); + for (let attempt = 0; attempt < MAX_INVALID_CALL_RETRIES; attempt++) { + const invalid = findInvalidCalls(response); + if (invalid.length === 0) return response; + response = await handler({ + ...request, + messages: [...request.messages, ...correction(response, invalid, offered)], + }); + } + const invalid = findInvalidCalls(response); + return invalid.length === 0 ? response : withNotice(response, invalidCallNotice(invalid)); + }, + }); +} diff --git a/src/services/messageStream.test.ts b/src/services/messageStream.test.ts index 4a4c586..cbf5e35 100644 --- a/src/services/messageStream.test.ts +++ b/src/services/messageStream.test.ts @@ -235,3 +235,116 @@ describe("streamMessages artifacts", () => { expect(tool?.data).not.toHaveProperty("artifact"); }); }); + +describe("streamMessages results a middleware returned itself", () => { + // Such a call never runs the tool node, so it is not on `run.toolCalls`: only in the state. + const answered = (callId: string) => + new ToolMessage({ + content: '{"ok":false,"error":"server_unavailable"}', + tool_call_id: callId, + name: "notes_list_notes", + status: "error", + }); + const call: ToolCall = { id: "call-1", name: "notes_list_notes", args: {}, type: "tool_call" }; + + const pendingIds = (chunks: MessageResponse[]) => + chunks.flatMap((c) => (c.data as { pendingToolCallIds?: string[] }).pendingToolCallIds ?? []); + const results = (chunks: MessageResponse[]) => chunks.filter((c) => c.type === "tool"); + + it("streams it from a state snapshot and never reports it pending", async () => { + const chunks = await collect({ + messages: fromArray([aiMessage([""], [call])]), + toolCalls: fromArray([]), + // Released after the AI message is in, as in a real run. + values: (async function* () { + await new Promise((r) => setTimeout(r, 0)); + yield { messages: [answered("call-1")] }; + })(), + output: Promise.resolve({ messages: [answered("call-1")] }), + }); + expect(results(chunks)).toHaveLength(1); + expect(results(chunks)[0].data).toMatchObject({ + tool_call_id: "call-1", + name: "notes_list_notes", + status: "error", + }); + expect(pendingIds(chunks)).toEqual([]); + }); + + it("finds it in the final state when no snapshot carried it", async () => { + const chunks = await collect({ + messages: fromArray([aiMessage([""], [call])]), + toolCalls: fromArray([]), + values: fromArray([]), + output: Promise.resolve({ messages: [answered("call-1")] }), + }); + expect(results(chunks)).toHaveLength(1); + expect(pendingIds(chunks)).toEqual([]); + }); + + it("sends a result once when both the execution stream and the state carry it", async () => { + const chunks = await collect({ + messages: fromArray([aiMessage([""], [call])]), + toolCalls: fromArray([toolCall("notes_list_notes", "call-1")]), + values: fromArray([{ messages: [answered("call-1")] }]), + output: Promise.resolve({ messages: [answered("call-1")] }), + }); + expect(results(chunks)).toHaveLength(1); + }); + + it("ignores earlier results in the thread, and still reports a paused call", async () => { + const chunks = await collect({ + messages: fromArray([aiMessage([""], [call])]), + toolCalls: fromArray([]), + values: fromArray([{ messages: [answered("call-from-an-earlier-turn")] }]), + output: Promise.resolve({ messages: [answered("call-from-an-earlier-turn")] }), + }); + expect(results(chunks)).toEqual([]); + expect(pendingIds(chunks)).toEqual(["call-1"]); + }); +}); + +describe("streamMessages malformed tool calls", () => { + it("also finds one carried only as a v1 content block", async () => { + const chunks = await collect({ + messages: fromArray([ + { + text: fromArray(["I'll pull all your notes now."]), + output: Promise.resolve({ + invalid_tool_calls: [], + content: [ + { type: "text", text: "I'll pull all your notes now." }, + { type: "invalid_tool_call", id: "bad-1", name: "notes_list_notes", args: "" }, + ], + }), + }, + ]), + toolCalls: fromArray([]), + values: fromArray([]), + output: Promise.resolve({ messages: [] }), + }); + const text = chunks + .map((c) => String((c.data as { content?: unknown }).content ?? "")) + .join(""); + expect(text).toContain("The call to notes_list_notes had arguments that were not valid JSON"); + }); + + it("shows a notice for a call whose arguments didn't parse, which never runs", async () => { + const chunks = await collect({ + messages: fromArray([ + { + text: fromArray(["I'll pull all your notes now."]), + output: Promise.resolve({ invalid_tool_calls: [{ name: "notes_list_notes" }] }), + }, + ]), + toolCalls: fromArray([]), + values: fromArray([]), + output: Promise.resolve({ messages: [] }), + }); + const text = chunks + .map((c) => String((c.data as { content?: unknown }).content ?? "")) + .join(""); + expect(text).toContain("I'll pull all your notes now."); + expect(text).toContain("The call to notes_list_notes had arguments that were not valid JSON"); + }); +}); diff --git a/src/services/messageStream.ts b/src/services/messageStream.ts index 01c971e..42e42f6 100644 --- a/src/services/messageStream.ts +++ b/src/services/messageStream.ts @@ -1,6 +1,7 @@ import { ToolMessage } from "@langchain/core/messages"; -import type { MessageResponse, ToolCall } from "@/types/message"; +import type { MessageResponse, ToolCall, ToolMessageData } from "@/types/message"; import { producesArtifact } from "@/lib/agent/artifactTools"; +import { findInvalidCalls, invalidCallNotice } from "@/lib/agent/invalidCallNotice"; /** * Turns a LangGraph run's independent projections into the ordered MessageResponse chunks the SSE @@ -16,7 +17,12 @@ import { producesArtifact } from "@/lib/agent/artifactTools"; export interface AgentRun { messages: AsyncIterable<{ text: AsyncIterable; - output: PromiseLike<{ id?: string; tool_calls?: ToolCall[] }>; + output: PromiseLike<{ + id?: string; + tool_calls?: ToolCall[]; + invalid_tool_calls?: unknown[]; + content?: unknown; + }>; }>; toolCalls: AsyncIterable<{ name: string; @@ -49,6 +55,34 @@ export function streamMessages(run: AgentRun): AsyncGenerator(); const executedCallIds = new Set(); + // A result can reach us twice (execution stream and state), so each is sent once. Sending one + // also marks its call executed: it has an answer, so it is not awaiting approval. + const pushResult = (callId: string, data: ToolMessageData) => { + if (executedCallIds.has(callId)) return; + executedCallIds.add(callId); + push({ type: "tool", data }); + }; + + /** + * Results a middleware returned itself (an MCP call to a removed server, say) never run the + * tool node, so they are not on `run.toolCalls`; they are only in the state. Only calls this + * run requested count — the snapshot also holds the thread's earlier results. + */ + const pushResultsFrom = (snapshot: unknown) => { + for (const m of (snapshot as { messages?: unknown[] })?.messages ?? []) { + if (!ToolMessage.isInstance(m) || !m.tool_call_id) continue; + if (!requestedCalls.has(m.tool_call_id)) continue; + pushResult(m.tool_call_id, { + id: m.tool_call_id, + content: typeof m.content === "string" ? m.content : JSON.stringify(m.content), + status: m.status ?? "success", + tool_call_id: m.tool_call_id, + name: m.name ?? requestedCalls.get(m.tool_call_id)!.name, + ...(m.artifact !== undefined ? { artifact: m.artifact } : {}), + }); + } + }; + const pumpMessages = async () => { let counter = 0; for await (const msg of run.messages) { @@ -60,6 +94,11 @@ export function streamMessages(run: AgentRun): AsyncGenerator 0) { + push({ type: "ai", data: { id, content: `\n\n${invalidCallNotice(invalid)}` } }); + } if (assembled?.tool_calls && assembled.tool_calls.length > 0) { for (const call of assembled.tool_calls) { if (call.id) requestedCalls.set(call.id, { id, name: call.name }); @@ -80,7 +119,10 @@ export function streamMessages(run: AgentRun): AsyncGenerator { - for await (const snapshot of run.values) collectArtifacts(snapshot); + for await (const snapshot of run.values) { + collectArtifacts(snapshot); + pushResultsFrom(snapshot); + } }; const pumpToolCalls = async () => { @@ -96,18 +138,16 @@ export function streamMessages(run: AgentRun): AsyncGenerator pushResultsFrom(output)) .then(pushPendingApprovals) .catch((e) => { failure = e;