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;