From f1bf23c18727d2b694d422cc97b05f38f306d6d2 Mon Sep 17 00:00:00 2001 From: Guilhem Lemouel Date: Thu, 10 Sep 2026 20:17:18 +0200 Subject: [PATCH] fix: recover the chat when a registered MCP schema is rejected Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01JgzuxyafKNF2uaEeL35XQw --- .../copilot/chat/AIChatManager.svelte.ts | 40 +++++++++++++- .../lib/components/copilot/chat/chatLoop.ts | 2 +- .../copilot/chat/global/mcpTools.test.ts | 52 +++++++++++++++++++ .../copilot/chat/global/mcpTools.ts | 50 +++++++++++------- frontend/src/lib/components/mcp/serverMark.ts | 27 +++++----- 5 files changed, 138 insertions(+), 33 deletions(-) diff --git a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts index b1458b8602..13378da13b 100644 --- a/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts +++ b/frontend/src/lib/components/copilot/chat/AIChatManager.svelte.ts @@ -103,7 +103,12 @@ import { import type { Selection } from 'monaco-editor' import type AIChatInput from './AIChatInput.svelte' import { prepareApiSystemMessage, prepareApiUserMessage } from './api/core' -import { closeInterruptedToolBatch, runChatLoop, truncateToToolPairedPrefix } from './chatLoop' +import { + closeInterruptedToolBatch, + getErrorStatus, + runChatLoop, + truncateToToolPairedPrefix +} from './chatLoop' import { sanitizeToolCallArguments } from './toolCallArguments' import { billedTokens, normalizeContextUsage, type ChatTokenUsage } from './tokenUsage' import { logAiUsage } from '$lib/utils/aiUsageReporter' @@ -138,6 +143,7 @@ import { randomUUID } from '$lib/utils/uuid' import { createMcpTools, forgetLoadedMcpTools, + invalidateMcpRegistrations, loadedMcpServers, loadedMcpTools, loadMcpServers, @@ -1208,6 +1214,8 @@ export class AIChatManager { * the docked chat plus a warm runtime per session — and each is its own conversation. */ readonly mcpOwnerId = randomUUID() private mcpServersRefreshId = 0 + /** The connected set the last refresh settled on, to notice when it changes. */ + private mcpServersSignature = '' // The GLOBAL prompt's path conventions and folder ACLs, for this chat's operating // workspace (`GlobalPromptIdentity`). Resolved asynchronously alongside skills, never @@ -2255,6 +2263,15 @@ export class AIChatManager { // a registered call bypasses the listing cache, so a connection edited elsewhere // would otherwise keep running against the schema it was frozen with. const live = new Map(this.mcpServers.map((s) => [s.path, s.editedAt])) + // The reconcile below only reaches servers that already registered something. A + // search still awaiting its listing has registered nothing yet, so a server turned + // off during that await would register after the fact and be advertised on the next + // iteration — bumping the generation makes that search drop its results instead. + const signature = this.mcpServers.map((s) => `${s.path}@${s.editedAt ?? ''}`).join(',') + if (signature !== this.mcpServersSignature) { + this.mcpServersSignature = signature + invalidateMcpRegistrations(this.mcpOwnerId) + } for (const { path, editedAt } of loadedMcpServers(this.mcpOwnerId)) { if (!live.has(path) || live.get(path) !== editedAt) { forgetLoadedMcpTools(this.mcpOwnerId, path) @@ -2265,6 +2282,26 @@ export class AIChatManager { } } + /** + * A registered MCP tool carries a schema a third party wrote, and a provider that + * refuses it refuses every later request in the conversation the same way — with an + * error naming the request, not the tool. So a rejection drops the registered tools: + * the next send goes out with the search tool and the wrappers only, and the model + * can register again. A false positive costs one re-search. Statuses about the + * account rather than the body (auth, quota) are left alone. + */ + private dropMcpToolsOnRejectedRequest = (err: unknown) => { + const status = getErrorStatus(err) + if (status === undefined || status < 400 || status >= 500) return + if (status === 401 || status === 403 || status === 429) return + if (loadedMcpTools(this.mcpOwnerId).length === 0) return + console.warn('Dropping registered MCP tools after a rejected request', err) + forgetLoadedMcpTools(this.mcpOwnerId) + if (this.mode === AIMode.GLOBAL) { + this.configureGlobalMode() + } + } + // The workspace's AI provider resources and their models exist only at run time, so they are // appended once the catalog resolves. The chat loop re-reads this.systemMessage on every // iteration, so a send that beats the fetch still picks them up on the next one. @@ -2807,6 +2844,7 @@ export class AIChatManager { console.error('chatRequest error', err) callbacks.onMessageEnd() this.cancelLoadingTools('Error') + this.dropMcpToolsOnRejectedRequest(err) if (!abortController.signal.aborted) { throw err } diff --git a/frontend/src/lib/components/copilot/chat/chatLoop.ts b/frontend/src/lib/components/copilot/chat/chatLoop.ts index 110e4f1154..955a77e7a6 100644 --- a/frontend/src/lib/components/copilot/chat/chatLoop.ts +++ b/frontend/src/lib/components/copilot/chat/chatLoop.ts @@ -240,7 +240,7 @@ function getErrorText(err: unknown): string { } } -function getErrorStatus(err: unknown): number | undefined { +export function getErrorStatus(err: unknown): number | undefined { if (!isRecord(err)) { return undefined } diff --git a/frontend/src/lib/components/copilot/chat/global/mcpTools.test.ts b/frontend/src/lib/components/copilot/chat/global/mcpTools.test.ts index b76ddebfd3..fc9816d407 100644 --- a/frontend/src/lib/components/copilot/chat/global/mcpTools.test.ts +++ b/frontend/src/lib/components/copilot/chat/global/mcpTools.test.ts @@ -35,6 +35,7 @@ import { clearMcpToolsCache, createMcpTools, forgetLoadedMcpTools, + invalidateMcpRegistrations, loadedMcpServers, mcpRegistryGeneration, loadedMcpTools, @@ -416,6 +417,16 @@ describe('loaded remote tools', () => { expect(loadedMcpTools(OWNER)).toEqual([]) }) + // Turning a server off while a search is still awaiting its listing leaves nothing + // registered for the reconcile to find, so the registration itself has to be rejected. + it('drops results of a search that finished after the connected set changed', () => { + const generation = mcpRegistryGeneration(OWNER) + invalidateMcpRegistrations(OWNER) + + expect(registerMcpTools(OWNER, generation, server, [TOOLS[0]])).toEqual([undefined]) + expect(loadedMcpTools(OWNER)).toEqual([]) + }) + // The key is `${server}::${tool}` and a remote tool name is arbitrary — namespaced // names are ordinary. Splitting on the last `::` reported a server that does not // exist, so the reconcile dropped the tool at the start of every send. @@ -598,6 +609,47 @@ describe('search_mcp_tools', () => { ]) }) + // Registering the second server can evict tools the first one just registered, and a + // `call` name whose tool is gone answers `Unknown tool call` on the next iteration. + it('advertises no call name for a match a later server evicted', async () => { + const wide = (name: string) => ({ + name, + description: 'issue', + inputSchema: { + type: 'object', + properties: { a: { type: 'string', description: 'x'.repeat(7_000) } } + } + }) + const servers: McpServer[] = [{ path: 'u/hugo/a_mcp' }, { path: 'u/hugo/b_mcp' }] + getMcpToolsMock.mockImplementation(({ path }: { path: string }) => + Promise.resolve( + Array.from({ length: 5 }, (_, i) => wide(`${path === 'u/hugo/a_mcp' ? 'a' : 'b'}_${i}`)) + ) + ) + const tool = createMcpTools(OWNER, servers).find( + (entry) => entry.def.function.name === 'search_mcp_tools' + )! + const result = JSON.parse( + await tool.fn({ + args: { query: 'issue' }, + workspace: 'test-ws', + helpers: {}, + toolCallbacks: createToolCallbacks(), + toolId: 'tool-1' + }) + ) + + const registered = new Set(loadedMcpTools(OWNER).map((t) => t.def.function.name)) + expect(registered.size).toBeLessThan(10) + const advertised = result.matches + .map((m: { call?: string }) => m.call) + .filter((c: string | undefined) => c !== undefined) + expect(advertised.length).toBe(registered.size) + for (const call of advertised) { + expect(registered.has(call)).toBe(true) + } + }) + it('still returns matches when one server is unreachable', async () => { const servers: McpServer[] = [...SERVERS, { path: 'u/hugo/broken_mcp' }] getMcpToolsMock.mockImplementation(({ path }: { path: string }) => diff --git a/frontend/src/lib/components/copilot/chat/global/mcpTools.ts b/frontend/src/lib/components/copilot/chat/global/mcpTools.ts index ac796cd91f..2f866ef431 100644 --- a/frontend/src/lib/components/copilot/chat/global/mcpTools.ts +++ b/frontend/src/lib/components/copilot/chat/global/mcpTools.ts @@ -92,18 +92,15 @@ export function clearMcpToolsCache() { } /** - * Remote tools promoted to first-class chat tools for the current conversation, - * keyed `${server}::${tool}`. `chatLoop` re-reads its tool list on every iteration, - * so one registered while a call is running is callable on the next one, and the - * whole set is dropped when the conversation rotates — a tool the model was never - * told about in a fresh chat should not be in its list, nor its schema in the bill. + * Remote tools promoted to first-class chat tools for the current conversation, keyed + * `${server}::${tool}`. `chatLoop` re-reads its tool list every iteration, so one + * registered mid-turn is callable on the next, and the whole set is dropped when the + * conversation rotates. * - * A registered tool is a frozen copy of an input schema *and* of `readOnlyHint`, - * which is the thing `TOOLS_CACHE_TTL_MS` exists to bound — so these are dropped - * by the same triggers that drop a listing rather than left to age on a timer. - * What keeps that safe in between is the backend, not this map: `executeTool` - * asserts `read_only` against the live server, so a stale hint cannot turn into - * an unconfirmed write. + * Each entry freezes an input schema and a `readOnlyHint`, which is what + * `TOOLS_CACHE_TTL_MS` bounds — so entries are dropped by the triggers that drop a + * listing rather than aged on a timer. `executeTool` re-asserts `read_only` against + * the live server in between, so a stale hint cannot become an unconfirmed write. */ type OwnerRegistry = { tools: Map> @@ -117,11 +114,9 @@ type OwnerRegistry = { } /** - * Keyed by owner, because several chats are live at once: the docked chat is a - * singleton and every warm session runtime builds its own manager, all in GLOBAL - * mode. One shared map would put a session's registrations into the docked chat's - * request, and let either one's "New chat" wipe a tool the other advertised an - * iteration ago. + * Keyed by owner: the docked chat and every warm session runtime each build a manager + * in GLOBAL mode. One shared map would put a session's registrations into the docked + * chat's request, and let either one's "New chat" wipe what the other advertised. */ const registries = new Map() /** @@ -142,6 +137,17 @@ function invalidateOwner(owner: string) { generations.set(owner, (generations.get(owner) ?? 0) + 1) } +/** + * Reject the registrations of searches that are still awaiting their listing, without + * touching what is already registered. Called when the connected server set changes: + * a search started before a server was turned off has registered nothing yet, so the + * reconcile over registered tools cannot reach it, and it would otherwise install that + * server's tools after the fact. + */ +export function invalidateMcpRegistrations(owner: string) { + invalidateOwner(owner) +} + /** * The server half of a `${server}::${tool}` key. Split on the FIRST `::`: a resource * path cannot contain `:` (the backend's path validation forbids it), but a remote @@ -617,8 +623,8 @@ function uniqueRegisteredName( /** * A remote server's `inputSchema`, made safe to send as a provider tool definition. - * A schema a provider rejects fails the whole completion, not the one tool, and the - * tool stays registered — so the chat keeps failing until it is dropped. + * A schema a provider rejects fails the whole completion, not the one tool; the manager + * drops the registered tools on a rejected request so the next send goes out clean. * * A guard, not a validator: anything it cannot make sense of collapses to "accepts any * object", which costs the model its argument names but keeps the chat alive. @@ -840,6 +846,14 @@ export function createMcpTools(owner: string, servers: McpServer[]): Tool<{}>[] if (name !== undefined) callNames.set(tool, name) }) } + // Each registration checks its own names against eviction, but registering a + // later server can evict a name an earlier one just returned. Advertising an + // evicted name yields `Unknown tool call`; dropping it here sends that match + // to the wrapper instead. + const live = new Set(loadedMcpTools(owner).map((t) => t.def.function.name)) + for (const [tool, name] of callNames) { + if (!live.has(name)) callNames.delete(tool) + } const result = boundedSearch({ matches: top.map((s) => summarizeTool(s.server, s.tool, callNames.get(s.tool))), hint: 'Call each `call` name directly; use the wrappers only for a match without one.', diff --git a/frontend/src/lib/components/mcp/serverMark.ts b/frontend/src/lib/components/mcp/serverMark.ts index 6ea0d6a390..6542c558a9 100644 --- a/frontend/src/lib/components/mcp/serverMark.ts +++ b/frontend/src/lib/components/mcp/serverMark.ts @@ -3,8 +3,7 @@ import { ResourceService } from '$lib/gen' import { cachedProviderMark } from './iconCache' import { loadProviderIcon, providerKey } from './providerIcon' -/** Windmill's own icon for a connected server's integration, when it ships one. The - * server's published icon is preferred over this and is resolved separately. */ +/** Windmill's own icon for a connected server's integration, when it ships one. */ export type McpServerMark = { icon?: Component } // One resolution per server per session, shared by every transcript row naming it — @@ -15,7 +14,13 @@ export function resolveMcpServerMark(workspace: string, path: string): Promise { + marks.delete(key) + return {} + }) marks.set(key, pending) } return pending @@ -24,14 +29,10 @@ export function resolveMcpServerMark(workspace: string, path: string): Promise { const cached = cachedProviderMark(workspace, path) if (cached) return { icon: await loadProviderIcon(cached.key) } - try { - // Deliberately not written back to the shared cache: that entry is keyed by - // `editedAt` for the server list's sake, and storing one from here — where the - // row is a past call and `editedAt` is unknown — would make every list re-read. - const resource = await ResourceService.getResource({ workspace, path }) - const url = (resource.value as { url?: unknown } | undefined)?.url - return { icon: await loadProviderIcon(providerKey(url)) } - } catch { - return {} - } + // Deliberately not written back to the shared cache: that entry is keyed by + // `editedAt` for the server list's sake, and storing one from here — where the + // row is a past call and `editedAt` is unknown — would make every list re-read. + const resource = await ResourceService.getResource({ workspace, path }) + const url = (resource.value as { url?: unknown } | undefined)?.url + return { icon: await loadProviderIcon(providerKey(url)) } }