From 197e492ddf728657de08186c83c5197ffa6d0767 Mon Sep 17 00:00:00 2001 From: Guilhem Lemouel Date: Thu, 10 Sep 2026 17:47:02 +0200 Subject: [PATCH] refactor: bound the registered schema set and drop churn the review found Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01JgzuxyafKNF2uaEeL35XQw --- backend/windmill-mcp/src/common/transform.rs | 11 +--- .../copilot/chat/ToolExecutionDisplay.svelte | 13 ++-- .../chat/WebSearchSourcesDisplay.svelte | 15 +++-- .../copilot/chat/global/mcpTools.test.ts | 23 ++++++++ .../copilot/chat/global/mcpTools.ts | 59 ++++++++++--------- .../copilot/chat/toolSchema.test.ts | 13 ---- frontend/src/lib/components/mcp/iconCache.ts | 9 +-- .../src/lib/components/mcp/mcpIcon.test.ts | 5 +- frontend/src/lib/components/mcp/mcpIcon.ts | 5 +- 9 files changed, 75 insertions(+), 78 deletions(-) diff --git a/backend/windmill-mcp/src/common/transform.rs b/backend/windmill-mcp/src/common/transform.rs index 2402fc2a92..4514787bc0 100644 --- a/backend/windmill-mcp/src/common/transform.rs +++ b/backend/windmill-mcp/src/common/transform.rs @@ -69,11 +69,7 @@ pub fn transform_hub_path(version_id: u64, summary: &str) -> String { /// Returns `(type_str, is_hub, is_hashed)`. /// Hashed names use an uppercase first character as the signal. pub fn parse_tool_prefix(name: &str) -> Result<(&str, bool, bool), String> { - let is_hashed = name - .chars() - .next() - .map(|c| c.is_ascii_uppercase()) - .unwrap_or(false); + let is_hashed = name.chars().next().map(|c| c.is_ascii_uppercase()).unwrap_or(false); let lower = name.to_ascii_lowercase(); let (type_str, is_hub) = if lower.starts_with("hs-") { ("script", true) @@ -432,10 +428,7 @@ mod tests { #[test] fn test_extract_path_prefix_handles_hs_prefix() { // Hs- is 3 chars, not 2 — ensure the prefix is stripped correctly - let hashed = transform_hub_path( - 12345, - "a]very long hub script summary that exceeds the limit", - ); + let hashed = transform_hub_path(12345, "a]very long hub script summary that exceeds the limit"); let (_, is_hub, is_hashed) = parse_tool_prefix(&hashed).unwrap(); assert!(is_hub); assert!(is_hashed); diff --git a/frontend/src/lib/components/copilot/chat/ToolExecutionDisplay.svelte b/frontend/src/lib/components/copilot/chat/ToolExecutionDisplay.svelte index 5ac02974d5..cb12affcb3 100644 --- a/frontend/src/lib/components/copilot/chat/ToolExecutionDisplay.svelte +++ b/frontend/src/lib/components/copilot/chat/ToolExecutionDisplay.svelte @@ -250,16 +250,13 @@ {/if} {/snippet} - + {#snippet serverMark()} {#if mcpServerPath && aiChatManager.operatingWorkspace} - + {#await resolveMcpServerMark(aiChatManager.operatingWorkspace, mcpServerPath) then mark} {#if message.mcpIconSrc || mark.icon} diff --git a/frontend/src/lib/components/copilot/chat/WebSearchSourcesDisplay.svelte b/frontend/src/lib/components/copilot/chat/WebSearchSourcesDisplay.svelte index 8b954362b1..ff71c44160 100644 --- a/frontend/src/lib/components/copilot/chat/WebSearchSourcesDisplay.svelte +++ b/frontend/src/lib/components/copilot/chat/WebSearchSourcesDisplay.svelte @@ -36,19 +36,18 @@ } } + const failedFavicons = new SvelteSet() + // Favicons come from Google's public favicon service, which discloses each // consulted hostname to a third party from the user's browser — an accepted - // tradeoff for a page the user just searched for (blocked/air-gapped - // environments degrade to the Globe icon via onerror). Hit gstatic directly - // rather than www.google.com/s2/favicons: the app is served with COEP - // require-corp, and the s2 redirect hop carries no Cross-Origin-Resource-Policy - // header, so the browser blocks the image. The gstatic endpoint itself responds - // with CORP: cross-origin. + // tradeoff for now (blocked/air-gapped environments degrade to the Globe + // icon via onerror). Hit gstatic directly rather than www.google.com/s2/ + // favicons: the app is served with COEP require-corp, and the s2 redirect + // hop carries no Cross-Origin-Resource-Policy header, so the browser blocks + // the image. The gstatic endpoint itself responds with CORP: cross-origin. function faviconUrl(hostname: string): string { return `https://t3.gstatic.com/faviconV2?client=SOCIAL&type=FAVICON&fallback_opts=TYPE,SIZE,URL&url=https://${encodeURIComponent(hostname)}&size=64` } - - const failedFavicons = new SvelteSet()
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 087543624f..b76ddebfd3 100644 --- a/frontend/src/lib/components/copilot/chat/global/mcpTools.test.ts +++ b/frontend/src/lib/components/copilot/chat/global/mcpTools.test.ts @@ -328,6 +328,29 @@ describe('loaded remote tools', () => { expect(loadedMcpTools(OWNER).map((t) => t.def.function.name)).toEqual(before) }) + // The per-tool bound alone would allow 25 large schemas in every request, which is + // the cost the whole search-then-register indirection exists to avoid. + it('bounds the total size of the registered schemas', () => { + const big = (name: string) => ({ + name, + description: 'x', + inputSchema: { + type: 'object', + properties: { a: { type: 'string', description: 'x'.repeat(7_000) } } + } + }) + for (let i = 0; i < 12; i++) { + registerMcpTools(OWNER, mcpRegistryGeneration(OWNER), server, [big(`tool_${i}`)] as any) + } + + const total = loadedMcpTools(OWNER).reduce( + (n, t) => n + JSON.stringify(t.def.function.parameters).length, + 0 + ) + expect(total).toBeLessThanOrEqual(40_000) + expect(loadedMcpTools(OWNER).length).toBeLessThan(12) + }) + it('bounds the loaded set, evicting the least recently registered', () => { for (let i = 0; i < 30; i++) { registerMcpTools(OWNER, mcpRegistryGeneration(OWNER), server, [ diff --git a/frontend/src/lib/components/copilot/chat/global/mcpTools.ts b/frontend/src/lib/components/copilot/chat/global/mcpTools.ts index 029970cbe4..e3b1c68315 100644 --- a/frontend/src/lib/components/copilot/chat/global/mcpTools.ts +++ b/frontend/src/lib/components/copilot/chat/global/mcpTools.ts @@ -17,6 +17,9 @@ import { pickMcpIconSrc } from '$lib/components/mcp/mcpIcon' * call wrappers remain as a fallback for anything not registered, and take a * free-form argument object — which is what the model has to guess into, and * why the registered tool is preferred wherever there is one. + * + * Registrations do not outlive the page: a reloaded transcript shows the calls it made + * but no longer offers those tools, and the model searches again to get them back. */ type McpToolDef = GetMcpToolsResponse[number] @@ -35,6 +38,9 @@ const MAX_TOOL_NAME_CHARS = 64 // controls. Past this, the tool is left unregistered and the model reaches it through // `call_mcp_*` instead — the schema still arrives, but only when a call fails. const MAX_TOOL_SCHEMA_CHARS = 8_000 +// And a ceiling on the set, since the per-tool bound alone would allow 25 large ones — +// far more context than the search indirection exists to save. +const MAX_LOADED_SCHEMA_CHARS = 40_000 // Same reasoning for the description, which rides alongside it. Roomier than the // search summary's 200: this one is what the model chooses the tool from. const MAX_TOOL_DESCRIPTION_CHARS = 2_000 @@ -227,7 +233,7 @@ export function forgetLoadedMcpTools(owner: string, serverPath?: string) { * Drop every owner's loaded tools. Used when a server resource itself changed, which * makes the frozen copy every chat holds stale at once — not for one chat rotating. */ -export function forgetAllLoadedMcpTools() { +function forgetAllLoadedMcpTools() { allGeneration++ registries.clear() } @@ -321,8 +327,8 @@ function summarizeTool(server: McpServer, tool: McpToolDef, callName?: string) { return { server: server.path, tool: tool.name, - // Once the tool is registered its real schema is in front of the model, so the - // bare parameter names it used to guess from are dropped rather than repeated. + // A registered tool puts its real schema in front of the model, so bare parameter + // names are only worth sending for one that is not. ...(callName ? { call: callName } : params.length > 0 ? { params } : {}), description: truncate(tool.description ?? '', MAX_DESCRIPTION_CHARS), mode: isReadOnly(tool) ? 'read' : 'write' @@ -619,16 +625,11 @@ 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. * - * Until now a remote schema only ever reached the model as tool-result text; a - * registered tool puts it in the request itself, where a provider rejects the whole - * completion rather than the one bad tool — and the chat would keep failing, because - * the tool stays registered. Windmill's own MCP server has shipped both of the shapes - * guarded here, so "the server is trusted" is not an argument. - * - * Deliberately shallow: this is a guard against a request-breaking schema, 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. + * 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. */ function safeInputSchema(schema: unknown): Record { const empty = { type: 'object', properties: {} } @@ -640,33 +641,38 @@ function safeInputSchema(schema: unknown): Record { !Array.isArray(rest.properties) ? (rest.properties as Record) : {} - // `required` is `uniqueItems`, and must name properties that exist: a strict - // validator rejects a repeat, and a name with no property reads as a parameter - // the model can never satisfy. + // A name with no property reads as a parameter the model can never satisfy. + // (`normalizeToolParameterSchema` below is what de-duplicates, at every depth.) const required = Array.isArray(rest.required) ? [ ...new Set( rest.required.filter( - (name): name is string => typeof name === 'string' && name in properties + (name): name is string => typeof name === 'string' && Object.hasOwn(properties, name) ) ) ] : [] - // Cloned, not spread: the spread would leave `properties` and any `items`/`allOf` - // carried through `rest` pointing at the listing cache's own objects, and the - // normalize pass below rewrites nested nodes in place. A registered tool is - // supposed to be a frozen copy — sharing that structure makes it one in name only. + // Cloned so the normalize pass below, which rewrites nested nodes in place, cannot + // reach the listing cache these objects came from. const safe = structuredClone({ ...rest, type: 'object', properties, required }) - // The same pass `createToolDef` runs on every other tool, so a remote schema is not - // the one that reaches a provider unnormalized. It recurses, which this does not: - // Windmill's own MCP server emits `format: ""` on untyped fields. normalizeToolParameterSchema(safe) return safe } +function registrySchemaChars(registry: OwnerRegistry): number { + let total = 0 + for (const tool of registry.tools.values()) { + total += JSON.stringify(tool.def.function.parameters ?? {}).length + } + return total +} + function evictLoadedTools(registry: OwnerRegistry) { const { tools: loadedTools, lastUsed } = registry - while (loadedTools.size > MAX_LOADED_TOOLS) { + while ( + loadedTools.size > MAX_LOADED_TOOLS || + (loadedTools.size > 1 && registrySchemaChars(registry) > MAX_LOADED_SCHEMA_CHARS) + ) { let victim: string | undefined let victimRank = Infinity for (const key of loadedTools.keys()) { @@ -716,9 +722,8 @@ export function registerMcpTools( // than none, since the model cannot tell which half it is missing. return undefined } - // Set without deleting first: re-registering refreshes the schema in place, - // and Map.set keeps an existing key's position, so the emitted tool list — and - // the cache breakpoint on its last entry — does not move. + // `Map.set` keeps an existing key's position, so re-registering refreshes the + // schema without moving the emitted list (see `lastUsed`). registry.lastUsed.set(key, ++registry.counter) registry.editedAt.set(key, server.editedAt) registry.tools.set(key, { diff --git a/frontend/src/lib/components/copilot/chat/toolSchema.test.ts b/frontend/src/lib/components/copilot/chat/toolSchema.test.ts index f4b2e79574..06d455ca9a 100644 --- a/frontend/src/lib/components/copilot/chat/toolSchema.test.ts +++ b/frontend/src/lib/components/copilot/chat/toolSchema.test.ts @@ -29,17 +29,4 @@ describe('normalizeToolParameterSchema', () => { expect(schema.properties.tags.items.required).toEqual(['name']) expect(schema.anyOf[0].required).toEqual(['a']) }) - - it('strips an empty or null format at every depth', () => { - const schema: Record = { - type: 'object', - format: '', - properties: { a: { type: 'string', format: null } } - } - - normalizeToolParameterSchema(schema) - - expect(schema.format).toBeUndefined() - expect(schema.properties.a.format).toBeUndefined() - }) }) diff --git a/frontend/src/lib/components/mcp/iconCache.ts b/frontend/src/lib/components/mcp/iconCache.ts index 7aeadfbfbb..ef748e8441 100644 --- a/frontend/src/lib/components/mcp/iconCache.ts +++ b/frontend/src/lib/components/mcp/iconCache.ts @@ -11,9 +11,7 @@ import { providerKey } from './providerIcon' */ type Entry = { key: string | null; editedAt?: string } -// Versioned so a shape change starts entries over rather than reading a stale one, -// at one resource read each. -const STORE_KEY = 'mcp_provider_icons_v2' +const STORE_KEY = 'mcp_provider_icons' function read(): Record> { try { @@ -56,10 +54,7 @@ export function rememberProviderKey( ): string | null { const key = providerKey(url) ?? null const store = read() - store[workspace] = { - ...(store[workspace] ?? {}), - [path]: { key, editedAt } - } + store[workspace] = { ...(store[workspace] ?? {}), [path]: { key, editedAt } } try { localStorage.setItem(STORE_KEY, JSON.stringify(store)) } catch {} diff --git a/frontend/src/lib/components/mcp/mcpIcon.test.ts b/frontend/src/lib/components/mcp/mcpIcon.test.ts index 0f017a0bf7..e156b47363 100644 --- a/frontend/src/lib/components/mcp/mcpIcon.test.ts +++ b/frontend/src/lib/components/mcp/mcpIcon.test.ts @@ -16,9 +16,8 @@ describe('pickMcpIconSrc', () => { ) }) - // An https src would make the browser fetch from a host the server names, which is - // the disclosure the favicon lookup was removed for. COEP blocks the response, not - // the request, so it does not save us. + // An https src would make the browser fetch from a host the server names. COEP + // blocks the response, not the request, so it does not prevent the disclosure. it('refuses a fetchable src, whatever the scheme', () => { expect(pickMcpIconSrc([{ src: 'https://example.com/logo.png' }])).toBeUndefined() expect(pickMcpIconSrc([{ src: 'http://example.com/i.png' }])).toBeUndefined() diff --git a/frontend/src/lib/components/mcp/mcpIcon.ts b/frontend/src/lib/components/mcp/mcpIcon.ts index 33c70a5e79..35df32ea13 100644 --- a/frontend/src/lib/components/mcp/mcpIcon.ts +++ b/frontend/src/lib/components/mcp/mcpIcon.ts @@ -18,9 +18,8 @@ const MAX_ICON_SRC_CHARS = 8_000 * user's browser fetch from a host the server names, disclosing their IP, user agent * and the moment they looked — and a per-user URL turns it into a read receipt. COEP * `require-corp` does not prevent that: it blocks the *response*, so the request has - * already left. That is the same disclosure the favicon lookup was removed for, and - * against an arbitrary host rather than one known party. A `data:` src carries its - * bytes over the MCP connection the user already made and fetches nothing. + * already left. A `data:` src carries its bytes over the MCP connection the user + * already made and fetches nothing. * * The rest is the spec's own list, enforced rather than trusted, since `src` is chosen * by a third party: no SVG (it can carry script), and a data URI judged by the type it