mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-09-21 00:02:30 +00:00
fix: scope registered MCP tools to their chat, several are live at once
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JgzuxyafKNF2uaEeL35XQw
This commit is contained in:
co-authored by
Claude Opus 5
parent
5952a8e407
commit
5be39c27c6
@@ -134,6 +134,7 @@ import {
|
||||
} from './global/core'
|
||||
import { formatChatJobCompletion } from './datatableTools'
|
||||
import { isGlobalAiEnabled } from './global/gate'
|
||||
import { randomUUID } from '$lib/utils/uuid'
|
||||
import {
|
||||
createMcpTools,
|
||||
forgetLoadedMcpTools,
|
||||
@@ -1203,6 +1204,9 @@ export class AIChatManager {
|
||||
// this is non-empty, so a workspace with no connection pays no schema cost
|
||||
// for them on every chat-loop iteration.
|
||||
mcpServers = $state<McpServer[]>([])
|
||||
/** Scopes this manager's registered MCP tools. Several managers are live at once —
|
||||
* the docked chat plus a warm runtime per session — and each is its own conversation. */
|
||||
readonly mcpOwnerId = randomUUID()
|
||||
private mcpServersRefreshId = 0
|
||||
|
||||
// The GLOBAL prompt's path conventions and folder ACLs, for this chat's operating
|
||||
@@ -2173,7 +2177,7 @@ export class AIChatManager {
|
||||
}
|
||||
}
|
||||
const pipeline = this.pipelineAiChatHelpers
|
||||
const mcpTools = createMcpTools(this.mcpServers)
|
||||
const mcpTools = createMcpTools(this.mcpOwnerId, this.mcpServers)
|
||||
if (pipeline) {
|
||||
systemMessage.content += getPipelinePromptSection(pipeline.getPipelineContext())
|
||||
this.tools = [
|
||||
@@ -2242,8 +2246,8 @@ export class AIChatManager {
|
||||
// left behind by a workspace switch would still be callable, and would run
|
||||
// against whichever workspace the chat is on now.
|
||||
const live = new Set(this.mcpServers.map((s) => s.path))
|
||||
for (const path of loadedMcpServerPaths()) {
|
||||
if (!live.has(path)) forgetLoadedMcpTools(path)
|
||||
for (const path of loadedMcpServerPaths(this.mcpOwnerId)) {
|
||||
if (!live.has(path)) forgetLoadedMcpTools(this.mcpOwnerId, path)
|
||||
}
|
||||
if (this.mode === AIMode.GLOBAL) {
|
||||
this.configureGlobalMode()
|
||||
@@ -2702,7 +2706,7 @@ export class AIChatManager {
|
||||
// only mode that installs the MCP tools or reconciles the loaded set, so
|
||||
// anywhere else these would be schemas a mode never opted into and
|
||||
// nothing would ever drop them.
|
||||
const mcpTools = self.mode === AIMode.GLOBAL ? loadedMcpTools() : []
|
||||
const mcpTools = self.mode === AIMode.GLOBAL ? loadedMcpTools(self.mcpOwnerId) : []
|
||||
return [...self.tools, ...self.planMode.tools, ...mcpTools]
|
||||
},
|
||||
get helpers() {
|
||||
@@ -4174,7 +4178,7 @@ export class AIChatManager {
|
||||
// Carrying them forward would put their schemas in the next conversation's tool
|
||||
// list — the cost the search indirection exists to avoid — and offer the model a
|
||||
// remote tool nobody there asked for.
|
||||
forgetLoadedMcpTools()
|
||||
forgetLoadedMcpTools(this.mcpOwnerId)
|
||||
this.onChatRotated?.(this.historyManager.getCurrentChatId())
|
||||
}
|
||||
|
||||
@@ -4226,7 +4230,7 @@ export class AIChatManager {
|
||||
this.#automaticScroll = true
|
||||
this.syncArtifactsSession()
|
||||
this.planMode.resetRound()
|
||||
forgetLoadedMcpTools()
|
||||
forgetLoadedMcpTools(this.mcpOwnerId)
|
||||
this.onChatRotated?.(id)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -54,14 +54,20 @@
|
||||
// arguments, and a registered tool is resolved through the in-memory registry.
|
||||
const MCP_CALL_TOOLS = ['call_mcp_read_tool', 'call_mcp_write_tool']
|
||||
const mcpServerPath = $derived.by(() => {
|
||||
if (message.mcpServer) return message.mcpServer
|
||||
const name = message.toolName
|
||||
if (!name) return undefined
|
||||
if (MCP_CALL_TOOLS.includes(name)) {
|
||||
const server = message.parameters?.server
|
||||
return typeof server === 'string' ? server : undefined
|
||||
}
|
||||
return mcpServerForToolName(name)
|
||||
const claimed = (() => {
|
||||
if (message.mcpServer) return message.mcpServer
|
||||
const name = message.toolName
|
||||
if (!name) return undefined
|
||||
if (MCP_CALL_TOOLS.includes(name)) {
|
||||
const server = message.parameters?.server
|
||||
return typeof server === 'string' ? server : undefined
|
||||
}
|
||||
return mcpServerForToolName(aiChatManager.mcpOwnerId, name)
|
||||
})()
|
||||
// Both sources are ultimately a path the model wrote, and resolving one reads
|
||||
// that resource and asks a third party for its host's favicon. Only a server
|
||||
// the user actually connected gets marked.
|
||||
return claimed && aiChatManager.mcpServers.some((s) => s.path === claimed) ? claimed : undefined
|
||||
})
|
||||
|
||||
const isPlanReview = $derived(message.toolName === EXIT_PLAN_MODE_TOOL)
|
||||
|
||||
@@ -44,6 +44,8 @@ import {
|
||||
import { setMcpEnabled } from '$lib/components/mcp/enabledServers'
|
||||
|
||||
const SERVERS: McpServer[] = [{ path: 'u/hugo/github_mcp' }]
|
||||
// Each chat manager scopes its own registry; these tests act as one of them.
|
||||
const OWNER = 'owner-a'
|
||||
|
||||
const TOOLS = [
|
||||
{
|
||||
@@ -78,7 +80,7 @@ function createToolCallbacks() {
|
||||
}
|
||||
|
||||
function getTool(name: string) {
|
||||
const tool = createMcpTools(SERVERS).find((entry) => entry.def.function.name === name)
|
||||
const tool = createMcpTools(OWNER, SERVERS).find((entry) => entry.def.function.name === name)
|
||||
if (!tool) throw new Error(`${name} tool not found`)
|
||||
return tool
|
||||
}
|
||||
@@ -103,7 +105,7 @@ beforeEach(() => {
|
||||
|
||||
describe('tool registration', () => {
|
||||
it('registers nothing when no MCP server is connected', () => {
|
||||
expect(createMcpTools([])).toEqual([])
|
||||
expect(createMcpTools(OWNER, [])).toEqual([])
|
||||
})
|
||||
})
|
||||
|
||||
@@ -111,18 +113,18 @@ describe('loaded remote tools', () => {
|
||||
const server = SERVERS[0]
|
||||
|
||||
it('registers a remote tool under its own schema and server-scoped name', () => {
|
||||
const [name] = registerMcpTools(server, [TOOLS[0]])
|
||||
const [name] = registerMcpTools(OWNER, server, [TOOLS[0]])
|
||||
|
||||
expect(name).toBe('mcp_u_hugo_github_mcp__get_issue')
|
||||
const tool = loadedMcpTools().find((t) => t.def.function.name === name)
|
||||
const tool = loadedMcpTools(OWNER).find((t) => t.def.function.name === name)
|
||||
expect(tool?.def.function.parameters).toEqual(TOOLS[0].inputSchema)
|
||||
})
|
||||
|
||||
// The wrappers ask for confirmation by which one the model picked; a registered
|
||||
// tool carries its own hint, so the gate is per tool.
|
||||
it('gates a mutating tool and lets a read-only one through', () => {
|
||||
registerMcpTools(server, [TOOLS[0], TOOLS[1], TOOLS[2]])
|
||||
const byName = Object.fromEntries(loadedMcpTools().map((t) => [t.def.function.name, t]))
|
||||
registerMcpTools(OWNER, server, [TOOLS[0], TOOLS[1], TOOLS[2]])
|
||||
const byName = Object.fromEntries(loadedMcpTools(OWNER).map((t) => [t.def.function.name, t]))
|
||||
|
||||
expect(byName['mcp_u_hugo_github_mcp__get_issue'].requiresConfirmation).toBeUndefined()
|
||||
expect(byName['mcp_u_hugo_github_mcp__merge_pull_request'].requiresConfirmation).toBe(true)
|
||||
@@ -133,22 +135,22 @@ describe('loaded remote tools', () => {
|
||||
// A registered tool freezes a copy of `readOnlyHint`, which is what the listing
|
||||
// TTL exists to bound — so it must not outlive the listing it came from.
|
||||
it('drops loaded tools when the listing cache is cleared', () => {
|
||||
registerMcpTools(server, [TOOLS[0]])
|
||||
expect(loadedMcpTools()).toHaveLength(1)
|
||||
registerMcpTools(OWNER, server, [TOOLS[0]])
|
||||
expect(loadedMcpTools(OWNER)).toHaveLength(1)
|
||||
|
||||
clearMcpToolsCache()
|
||||
|
||||
expect(loadedMcpTools()).toEqual([])
|
||||
expect(loadedMcpTools(OWNER)).toEqual([])
|
||||
})
|
||||
|
||||
it('drops only the named server when reconciling against the live list', () => {
|
||||
registerMcpTools(server, [TOOLS[0]])
|
||||
registerMcpTools({ path: 'f/team/linear_mcp' }, [TOOLS[0]])
|
||||
expect(loadedMcpServerPaths().sort()).toEqual(['f/team/linear_mcp', 'u/hugo/github_mcp'])
|
||||
registerMcpTools(OWNER, server, [TOOLS[0]])
|
||||
registerMcpTools(OWNER, { path: 'f/team/linear_mcp' }, [TOOLS[0]])
|
||||
expect(loadedMcpServerPaths(OWNER).sort()).toEqual(['f/team/linear_mcp', 'u/hugo/github_mcp'])
|
||||
|
||||
forgetLoadedMcpTools('f/team/linear_mcp')
|
||||
forgetLoadedMcpTools(OWNER, 'f/team/linear_mcp')
|
||||
|
||||
expect(loadedMcpServerPaths()).toEqual(['u/hugo/github_mcp'])
|
||||
expect(loadedMcpServerPaths(OWNER)).toEqual(['u/hugo/github_mcp'])
|
||||
})
|
||||
|
||||
// The reported bug: the model saw only parameter names, so it invented values for
|
||||
@@ -162,7 +164,7 @@ describe('loaded remote tools', () => {
|
||||
expect(match.call).toBe('mcp_u_hugo_github_mcp__get_issue')
|
||||
expect(match.params).toBeUndefined()
|
||||
|
||||
const registered = loadedMcpTools().find((t) => t.def.function.name === match.call)
|
||||
const registered = loadedMcpTools(OWNER).find((t) => t.def.function.name === match.call)
|
||||
expect(registered?.def.function.parameters).toEqual(TOOLS[0].inputSchema)
|
||||
})
|
||||
|
||||
@@ -187,11 +189,32 @@ describe('loaded remote tools', () => {
|
||||
})
|
||||
})
|
||||
|
||||
// Several chats are live at once — the docked one plus a warm runtime per session,
|
||||
// all in GLOBAL mode. A shared registry would put one chat's remote tools into
|
||||
// another's request, and let either one's "New chat" drop the other's.
|
||||
it('scopes registered tools to their owner', () => {
|
||||
const other = 'owner-b'
|
||||
registerMcpTools(OWNER, server, [TOOLS[0]])
|
||||
registerMcpTools(other, server, [TOOLS[1]])
|
||||
|
||||
expect(loadedMcpTools(OWNER).map((t) => t.def.function.name)).toEqual([
|
||||
'mcp_u_hugo_github_mcp__get_issue'
|
||||
])
|
||||
expect(loadedMcpTools(other).map((t) => t.def.function.name)).toEqual([
|
||||
'mcp_u_hugo_github_mcp__merge_pull_request'
|
||||
])
|
||||
|
||||
forgetLoadedMcpTools(OWNER)
|
||||
|
||||
expect(loadedMcpTools(OWNER)).toEqual([])
|
||||
expect(loadedMcpTools(other)).toHaveLength(1)
|
||||
})
|
||||
|
||||
// A registered tool's schema goes into the request, where a provider rejects the
|
||||
// whole completion over one bad tool — and the tool stays registered, so every
|
||||
// later send fails too. These are shapes Windmill's own MCP server has emitted.
|
||||
it('makes a request-breaking remote schema safe before registering it', () => {
|
||||
registerMcpTools(server, [
|
||||
registerMcpTools(OWNER, server, [
|
||||
{
|
||||
name: 'broken',
|
||||
description: 'x',
|
||||
@@ -206,7 +229,7 @@ describe('loaded remote tools', () => {
|
||||
} as any
|
||||
])
|
||||
|
||||
const params = loadedMcpTools()[0].def.function.parameters as any
|
||||
const params = loadedMcpTools(OWNER)[0].def.function.parameters as any
|
||||
expect(params.type).toBe('object')
|
||||
expect(params.required).toEqual(['a'])
|
||||
expect(params.$schema).toBeUndefined()
|
||||
@@ -226,41 +249,44 @@ describe('loaded remote tools', () => {
|
||||
}
|
||||
} as any
|
||||
|
||||
expect(registerMcpTools(server, [huge])).toEqual([undefined])
|
||||
expect(loadedMcpTools()).toEqual([])
|
||||
expect(registerMcpTools(OWNER, server, [huge])).toEqual([undefined])
|
||||
expect(loadedMcpTools(OWNER)).toEqual([])
|
||||
})
|
||||
|
||||
it('falls back to an empty object schema when the remote sends no usable one', () => {
|
||||
registerMcpTools(server, [{ name: 'nada', description: 'x', inputSchema: null } as any])
|
||||
registerMcpTools(OWNER, server, [{ name: 'nada', description: 'x', inputSchema: null } as any])
|
||||
|
||||
expect(loadedMcpTools()[0].def.function.parameters).toEqual({ type: 'object', properties: {} })
|
||||
expect(loadedMcpTools(OWNER)[0].def.function.parameters).toEqual({
|
||||
type: 'object',
|
||||
properties: {}
|
||||
})
|
||||
})
|
||||
|
||||
// The emitted list carries Anthropic's cache_control breakpoint on its last entry,
|
||||
// so calling or re-registering a tool must not move it: a reorder re-processes the
|
||||
// whole cached prefix on the next iteration.
|
||||
it('keeps the emitted tool order stable across calls and re-registration', async () => {
|
||||
registerMcpTools(server, [TOOLS[0], TOOLS[1]])
|
||||
const before = loadedMcpTools().map((t) => t.def.function.name)
|
||||
registerMcpTools(OWNER, server, [TOOLS[0], TOOLS[1]])
|
||||
const before = loadedMcpTools(OWNER).map((t) => t.def.function.name)
|
||||
|
||||
// Call the first-registered one, then re-register the pair.
|
||||
await loadedMcpTools()[0].fn({
|
||||
await loadedMcpTools(OWNER)[0].fn({
|
||||
args: {},
|
||||
workspace: 'test-ws',
|
||||
helpers: {},
|
||||
toolCallbacks: createToolCallbacks(),
|
||||
toolId: 'tool-1'
|
||||
})
|
||||
registerMcpTools(server, [TOOLS[0], TOOLS[1]])
|
||||
registerMcpTools(OWNER, server, [TOOLS[0], TOOLS[1]])
|
||||
|
||||
expect(loadedMcpTools().map((t) => t.def.function.name)).toEqual(before)
|
||||
expect(loadedMcpTools(OWNER).map((t) => t.def.function.name)).toEqual(before)
|
||||
})
|
||||
|
||||
it('bounds the loaded set, evicting the least recently registered', () => {
|
||||
for (let i = 0; i < 30; i++) {
|
||||
registerMcpTools(server, [{ ...TOOLS[0], name: `tool_${i}` }])
|
||||
registerMcpTools(OWNER, server, [{ ...TOOLS[0], name: `tool_${i}` }])
|
||||
}
|
||||
const names = loadedMcpTools().map((t) => t.def.function.name)
|
||||
const names = loadedMcpTools(OWNER).map((t) => t.def.function.name)
|
||||
|
||||
expect(names).toHaveLength(25)
|
||||
expect(names).not.toContain('mcp_u_hugo_github_mcp__tool_0')
|
||||
@@ -425,7 +451,7 @@ describe('search_mcp_tools', () => {
|
||||
? Promise.reject(new Error('connection refused'))
|
||||
: Promise.resolve(TOOLS)
|
||||
)
|
||||
const tool = createMcpTools(servers).find(
|
||||
const tool = createMcpTools(OWNER, servers).find(
|
||||
(entry) => entry.def.function.name === 'search_mcp_tools'
|
||||
)!
|
||||
const result = JSON.parse(
|
||||
@@ -469,7 +495,7 @@ describe('result size cap', () => {
|
||||
it('does not reuse a tool list across a resource revision', async () => {
|
||||
callMcpToolMock.mockResolvedValue({ content: [] })
|
||||
const call = (editedAt: string) =>
|
||||
createMcpTools([{ path: 'u/hugo/github_mcp', editedAt }])
|
||||
createMcpTools(OWNER, [{ path: 'u/hugo/github_mcp', editedAt }])
|
||||
.find((t) => t.def.function.name === 'call_mcp_read_tool')!
|
||||
.fn({
|
||||
args: { server: 'u/hugo/github_mcp', tool: 'get_issue', arguments: {} },
|
||||
|
||||
@@ -34,6 +34,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
|
||||
// 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
|
||||
const MAX_RESULT_CHARS = 20_000
|
||||
// A server writes its own error text, and every enabled server can contribute
|
||||
// one, so search results are capped the same way call results are.
|
||||
@@ -80,7 +83,7 @@ async function loadServerTools(
|
||||
export function clearMcpToolsCache() {
|
||||
cacheGeneration++
|
||||
toolsCache = {}
|
||||
forgetLoadedMcpTools()
|
||||
forgetAllLoadedMcpTools()
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -97,28 +100,49 @@ export function clearMcpToolsCache() {
|
||||
* asserts `read_only` against the live server, so a stale hint cannot turn into
|
||||
* an unconfirmed write.
|
||||
*/
|
||||
const loadedTools = new Map<string, Tool<{}>>()
|
||||
|
||||
/**
|
||||
* Recency for eviction, kept beside the map rather than as its order. The emitted
|
||||
* tool list carries Anthropic's `cache_control` breakpoint on its last entry, so
|
||||
* reordering it on every call would invalidate the cached prefix — system prompt,
|
||||
* skills and transcript — for the rest of the turn.
|
||||
*/
|
||||
const lastUsed = new Map<string, number>()
|
||||
let useCounter = 0
|
||||
|
||||
function touchLoadedTool(key: string) {
|
||||
if (loadedTools.has(key)) lastUsed.set(key, ++useCounter)
|
||||
type OwnerRegistry = {
|
||||
tools: Map<string, Tool<{}>>
|
||||
/**
|
||||
* Recency for eviction, kept beside the map rather than as its order. The emitted
|
||||
* tool list carries Anthropic's `cache_control` breakpoint on its last entry, so
|
||||
* reordering it on every call would invalidate the cached prefix — system prompt,
|
||||
* skills and transcript — for the rest of the turn.
|
||||
*/
|
||||
lastUsed: Map<string, number>
|
||||
counter: number
|
||||
}
|
||||
|
||||
export function loadedMcpTools(): Tool<{}>[] {
|
||||
return [...loadedTools.values()]
|
||||
/**
|
||||
* 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.
|
||||
*/
|
||||
const registries = new Map<string, OwnerRegistry>()
|
||||
|
||||
function registryFor(owner: string): OwnerRegistry {
|
||||
let registry = registries.get(owner)
|
||||
if (!registry) {
|
||||
registry = { tools: new Map(), lastUsed: new Map(), counter: 0 }
|
||||
registries.set(owner, registry)
|
||||
}
|
||||
return registry
|
||||
}
|
||||
|
||||
function touchLoadedTool(owner: string, key: string) {
|
||||
const registry = registryFor(owner)
|
||||
if (registry.tools.has(key)) registry.lastUsed.set(key, ++registry.counter)
|
||||
}
|
||||
|
||||
export function loadedMcpTools(owner: string): Tool<{}>[] {
|
||||
return [...(registries.get(owner)?.tools.values() ?? [])]
|
||||
}
|
||||
|
||||
/** The servers that currently have a tool registered, for reconciling against the live list. */
|
||||
export function loadedMcpServerPaths(): string[] {
|
||||
return [...new Set([...loadedTools.keys()].map((key) => key.slice(0, key.lastIndexOf('::'))))]
|
||||
export function loadedMcpServerPaths(owner: string): string[] {
|
||||
const keys = [...(registries.get(owner)?.tools.keys() ?? [])]
|
||||
return [...new Set(keys.map((key) => key.slice(0, key.lastIndexOf('::'))))]
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -126,27 +150,37 @@ export function loadedMcpServerPaths(): string[] {
|
||||
* the provider. Resolved through the registry rather than by parsing the name, which
|
||||
* a long path truncates.
|
||||
*/
|
||||
export function mcpServerForToolName(toolName: string): string | undefined {
|
||||
for (const [key, tool] of loadedTools) {
|
||||
export function mcpServerForToolName(owner: string, toolName: string): string | undefined {
|
||||
for (const [key, tool] of registries.get(owner)?.tools ?? []) {
|
||||
if (tool.def.function.name === toolName) return key.slice(0, key.lastIndexOf('::'))
|
||||
}
|
||||
return undefined
|
||||
}
|
||||
|
||||
/** Drop every loaded tool, or only those belonging to one server. */
|
||||
export function forgetLoadedMcpTools(serverPath?: string) {
|
||||
/** Drop one owner's loaded tools, or only those belonging to one server. */
|
||||
export function forgetLoadedMcpTools(owner: string, serverPath?: string) {
|
||||
const registry = registries.get(owner)
|
||||
if (!registry) return
|
||||
if (serverPath === undefined) {
|
||||
loadedTools.clear()
|
||||
lastUsed.clear()
|
||||
registries.delete(owner)
|
||||
return
|
||||
}
|
||||
const prefix = `${serverPath}::`
|
||||
for (const key of [...loadedTools.keys()]) {
|
||||
for (const key of [...registry.tools.keys()]) {
|
||||
if (key.startsWith(prefix)) {
|
||||
loadedTools.delete(key)
|
||||
lastUsed.delete(key)
|
||||
registry.tools.delete(key)
|
||||
registry.lastUsed.delete(key)
|
||||
}
|
||||
}
|
||||
if (registry.tools.size === 0) registries.delete(owner)
|
||||
}
|
||||
|
||||
/**
|
||||
* 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() {
|
||||
registries.clear()
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -155,13 +189,13 @@ export function forgetLoadedMcpTools(serverPath?: string) {
|
||||
* rejected by the same stale `readOnlyHint`, leaving it with nowhere to go until
|
||||
* the entry expires.
|
||||
*/
|
||||
function forgetServerTools(workspace: string, path: string) {
|
||||
function forgetServerTools(owner: string, workspace: string, path: string) {
|
||||
cacheGeneration++
|
||||
const prefix = `${workspace}:${path}:`
|
||||
for (const key of Object.keys(toolsCache)) {
|
||||
if (key.startsWith(prefix)) delete toolsCache[key]
|
||||
}
|
||||
forgetLoadedMcpTools(path)
|
||||
forgetLoadedMcpTools(owner, path)
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -286,6 +320,7 @@ function extractResultData(result: unknown): unknown {
|
||||
}
|
||||
|
||||
async function executeTool(
|
||||
owner: string,
|
||||
workspace: string,
|
||||
server: McpServer,
|
||||
tool: McpToolDef,
|
||||
@@ -313,7 +348,7 @@ async function executeTool(
|
||||
// write tool the model is sent to must not be handed the same answer. Matching
|
||||
// loosely is safe: the worst a false positive costs is one extra listing.
|
||||
if (skippedConfirmation && status === 400 && /read-only/i.test(error)) {
|
||||
forgetServerTools(workspace, server.path)
|
||||
forgetServerTools(owner, workspace, server.path)
|
||||
}
|
||||
return bounded({
|
||||
success: false,
|
||||
@@ -420,7 +455,7 @@ const callMcpToolSchema = z.object({
|
||||
* split they accept, and that check is what keeps a mutating call behind the
|
||||
* user's confirmation — building both from one body keeps them from drifting.
|
||||
*/
|
||||
function createCallTool(servers: McpServer[], mode: 'read' | 'write'): Tool<{}> {
|
||||
function createCallTool(owner: string, servers: McpServer[], mode: 'read' | 'write'): Tool<{}> {
|
||||
const isRead = mode === 'read'
|
||||
return {
|
||||
def: createToolDef(
|
||||
@@ -464,6 +499,7 @@ function createCallTool(servers: McpServer[], mode: 'read' | 'write'): Tool<{}>
|
||||
}
|
||||
toolCallbacks.setToolStatus(toolId, { content: `Calling ${parsed.tool}...` })
|
||||
const result = await executeTool(
|
||||
owner,
|
||||
workspace,
|
||||
resolved.server,
|
||||
resolved.tool,
|
||||
@@ -546,7 +582,8 @@ function safeInputSchema(schema: unknown): Record<string, unknown> {
|
||||
return safe
|
||||
}
|
||||
|
||||
function evictLoadedTools() {
|
||||
function evictLoadedTools(registry: OwnerRegistry) {
|
||||
const { tools: loadedTools, lastUsed } = registry
|
||||
while (loadedTools.size > MAX_LOADED_TOOLS) {
|
||||
let victim: string | undefined
|
||||
let victimRank = Infinity
|
||||
@@ -573,7 +610,12 @@ function evictLoadedTools() {
|
||||
* Each tool carries its own `readOnlyHint`, so the confirmation gate is per tool
|
||||
* and the read/write wrappers' "you used the wrong one" failure cannot arise.
|
||||
*/
|
||||
export function registerMcpTools(server: McpServer, tools: McpToolDef[]): (string | undefined)[] {
|
||||
export function registerMcpTools(
|
||||
owner: string,
|
||||
server: McpServer,
|
||||
tools: McpToolDef[]
|
||||
): (string | undefined)[] {
|
||||
const registry = registryFor(owner)
|
||||
const names = tools.map((tool) => {
|
||||
const key = `${server.path}::${tool.name}`
|
||||
const name = registeredToolName(server.path, tool.name)
|
||||
@@ -587,12 +629,15 @@ export function registerMcpTools(server: McpServer, tools: McpToolDef[]): (strin
|
||||
// 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.
|
||||
loadedTools.set(key, {
|
||||
registry.tools.set(key, {
|
||||
def: {
|
||||
type: 'function',
|
||||
function: {
|
||||
name,
|
||||
description: `${tool.description ?? tool.name}\n(MCP server ${server.path})`,
|
||||
// Bounded like the schema above and like every other payload the server
|
||||
// controls: this rides in the request on every iteration, for up to
|
||||
// MAX_LOADED_TOOLS tools, and servers do ship multi-KB descriptions.
|
||||
description: `${truncate(tool.description ?? tool.name, MAX_TOOL_DESCRIPTION_CHARS)}\n(MCP server ${server.path})`,
|
||||
parameters
|
||||
}
|
||||
},
|
||||
@@ -604,12 +649,12 @@ export function registerMcpTools(server: McpServer, tools: McpToolDef[]): (strin
|
||||
confirmationMessage: `Call ${tool.name} on ${server.path}`
|
||||
}),
|
||||
fn: async ({ args, workspace, toolId, toolCallbacks }) => {
|
||||
touchLoadedTool(key)
|
||||
touchLoadedTool(owner, key)
|
||||
toolCallbacks.setToolStatus(toolId, {
|
||||
content: `Calling ${tool.name}...`,
|
||||
mcpServer: server.path
|
||||
})
|
||||
const result = await executeTool(workspace, server, tool, args ?? {}, readOnly)
|
||||
const result = await executeTool(owner, workspace, server, tool, args ?? {}, readOnly)
|
||||
const ok = JSON.parse(result).success === true
|
||||
toolCallbacks.setToolStatus(toolId, {
|
||||
content: ok ? `Called ${tool.name}` : `Call to ${tool.name} failed`,
|
||||
@@ -621,7 +666,7 @@ export function registerMcpTools(server: McpServer, tools: McpToolDef[]): (strin
|
||||
})
|
||||
return name
|
||||
})
|
||||
evictLoadedTools()
|
||||
evictLoadedTools(registry)
|
||||
return names
|
||||
}
|
||||
|
||||
@@ -630,7 +675,7 @@ export function registerMcpTools(server: McpServer, tools: McpToolDef[]): (strin
|
||||
* are not registered at all, so a workspace without an MCP connection pays no
|
||||
* per-iteration schema cost for them.
|
||||
*/
|
||||
export function createMcpTools(servers: McpServer[]): Tool<{}>[] {
|
||||
export function createMcpTools(owner: string, servers: McpServer[]): Tool<{}>[] {
|
||||
if (servers.length === 0) return []
|
||||
const serverList = servers.map((s) => s.path).join(', ')
|
||||
|
||||
@@ -696,7 +741,7 @@ export function createMcpTools(servers: McpServer[]): Tool<{}>[] {
|
||||
perServerMatches.set(match.server.path, entry)
|
||||
}
|
||||
for (const { server, tools } of perServerMatches.values()) {
|
||||
const registered = registerMcpTools(server, tools)
|
||||
const registered = registerMcpTools(owner, server, tools)
|
||||
// A tool whose schema was too large to register has no `call` name, and
|
||||
// the summary then advertises the wrapper for it instead.
|
||||
tools.forEach((tool, i) => {
|
||||
@@ -721,7 +766,7 @@ export function createMcpTools(servers: McpServer[]): Tool<{}>[] {
|
||||
return result
|
||||
}
|
||||
},
|
||||
createCallTool(servers, 'read'),
|
||||
createCallTool(servers, 'write')
|
||||
createCallTool(owner, servers, 'read'),
|
||||
createCallTool(owner, servers, 'write')
|
||||
]
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user