diff --git a/packages/opencode/src/session/mcp-tools.ts b/packages/opencode/src/session/mcp-tools.ts index 64a64dcd4b..b31121ed98 100644 --- a/packages/opencode/src/session/mcp-tools.ts +++ b/packages/opencode/src/session/mcp-tools.ts @@ -2,6 +2,7 @@ import { Agent } from "@/agent/agent" import { EffectBridge } from "@/effect/bridge" import { RuntimeFlags } from "@/effect/runtime-flags" import { MCP } from "@/mcp" +import { McpCatalog } from "@/mcp/catalog" import { Permission } from "@/permission" import { Provider } from "@/provider/provider" import { ProviderTransform } from "@/provider/transform" @@ -251,6 +252,44 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In return tools }) +export const systemPrompt = Effect.fn("SessionMcpTools.systemPrompt")(function* (input: { + agent: Agent.Info + session: Session.Info +}) { + const mcp = yield* MCP.Service + const flags = yield* RuntimeFlags.Service + if (!flags.experimentalToolSearch) return undefined + + const mcpTools = yield* mcp.tools() + const mcpDisabled = Permission.disabled( + Object.keys(mcpTools), + Permission.merge(input.agent.permission, input.session.permission ?? []), + ) + const allowedTools = Object.fromEntries(Object.entries(mcpTools).filter(([key]) => !mcpDisabled.has(key))) + if (Object.keys(allowedTools).length === 0) return undefined + + const deferredDescriptors = yield* deferredToolDescriptors(allowedTools) + if ( + Token.estimate(JSON.stringify(deferredDescriptors.map(deferredToolEstimatePayload))) < + MIN_DEFERRED_MCP_SCHEMA_TOKENS + ) { + return undefined + } + + const servers = deferredServerSummaries(Object.keys(allowedTools), Object.keys(yield* mcp.clients())) + return [ + DEFERRED_TOOL_SYSTEM_PROMPT, + servers.length === 0 + ? undefined + : [ + "Deferred MCP servers available through `search_deferred_tools`:", + ...servers.map((server) => `- ${server.name}: ${server.count} tool${server.count === 1 ? "" : "s"}`), + ].join("\n"), + ] + .filter((part): part is string => part !== undefined) + .join("\n\n") +}) + function addResourceTools( tools: Record, deps: { @@ -710,6 +749,21 @@ function deferredToolEstimatePayload(descriptor: DeferredToolDescriptor) { return { id: descriptor.id, description: descriptor.description, input_schema: descriptor.inputSchema } } +function deferredServerSummaries(toolIDs: string[], serverNames: string[]) { + const prefixes = serverNames + .map((name) => ({ name, prefix: McpCatalog.sanitize(name) + "_" })) + .toSorted((a, b) => b.prefix.length - a.prefix.length || a.name.localeCompare(b.name)) + const counts = new Map() + for (const toolID of toolIDs) { + const server = prefixes.find((candidate) => toolID.startsWith(candidate.prefix)) + if (!server) continue + counts.set(server.name, (counts.get(server.name) ?? 0) + 1) + } + return Array.from(counts.entries()) + .map(([name, count]) => ({ name, count })) + .toSorted((a, b) => a.name.localeCompare(b.name)) +} + function searchDeferredTools(descriptors: DeferredToolDescriptor[], query: string, limit: number) { const terms = searchTerms(query) return descriptors diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 574ab6f937..86f0495dba 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -1254,18 +1254,24 @@ export const layer = Layer.effect( yield* plugin.trigger("experimental.chat.messages.transform", {}, { messages: msgs }) - const [skills, env, instructions, mcpInstructions, modelMsgs] = yield* Effect.all([ + const [skills, env, instructions, mcpInstructions, deferredToolInstructions, modelMsgs] = yield* Effect.all([ sys.skills(agent), sys.environment(model), instruction.system().pipe(Effect.orDie), sys.mcp(agent, session.permission), + tools.search_deferred_tools && tools.call_deferred_tool + ? SessionTools.deferredSystemPrompt({ agent, session }).pipe( + Effect.provideService(MCP.Service, mcp), + Effect.provideService(RuntimeFlags.Service, flags), + ) + : Effect.succeed(undefined), MessageV2.toModelMessagesEffect(msgs, model), ]) const system = [ ...env, ...instructions, ...(mcpInstructions ? [mcpInstructions] : []), - ...(tools.search_deferred_tools && tools.call_deferred_tool ? [SessionTools.DEFERRED_TOOL_SYSTEM_PROMPT] : []), + ...(deferredToolInstructions ? [deferredToolInstructions] : []), ...(skills ? [skills] : []), ] const format = lastUser.format ?? { type: "text" as const } diff --git a/packages/opencode/src/session/tools.ts b/packages/opencode/src/session/tools.ts index 70cd14d738..4c4c660e01 100644 --- a/packages/opencode/src/session/tools.ts +++ b/packages/opencode/src/session/tools.ts @@ -18,6 +18,7 @@ import { Session } from "./session" import { PartID } from "./schema" export const DEFERRED_TOOL_SYSTEM_PROMPT = SessionMcpTools.DEFERRED_TOOL_SYSTEM_PROMPT +export const deferredSystemPrompt = SessionMcpTools.systemPrompt export const resolve = Effect.fn("SessionTools.resolve")(function* (input: { agent: Agent.Info diff --git a/packages/opencode/test/session/tools.test.ts b/packages/opencode/test/session/tools.test.ts index e2f40e1826..6b1032f9ed 100644 --- a/packages/opencode/test/session/tools.test.ts +++ b/packages/opencode/test/session/tools.test.ts @@ -5,6 +5,7 @@ import { ProjectV2 } from "@opencode-ai/core/project" import { SessionV1 } from "@opencode-ai/core/v1/session" import { ModelV2 } from "@opencode-ai/core/model" import { ProviderV2 } from "@opencode-ai/core/provider" +import { Client } from "@modelcontextprotocol/sdk/client/index.js" import { Agent } from "@/agent/agent" import { MCP } from "@/mcp" import { Permission } from "@/permission" @@ -21,6 +22,7 @@ import { testEffect } from "../lib/effect" const model = ProviderTest.model() const sessionID = SessionID.make("ses_deferred-tools") const largeSchemaDescription = "analytics trends schema ".repeat(10_000) +const mcpClient = new Client({ name: "test", version: "0.0.0" }) const agent = { name: "build", mode: "primary", @@ -70,7 +72,7 @@ const makeIt = (input: { flags: Parameters[0]; queryD testEffect( Layer.mergeAll( Layer.mock(MCP.Service, { - clients: () => Effect.succeed({}), + clients: () => Effect.succeed({ posthog: mcpClient }), tools: () => Effect.succeed({ posthog_query_trends: tool({ @@ -158,6 +160,15 @@ describe("session.tools", () => { }), ) + deferredIt.instance("lists deferred MCP servers in the system prompt", () => + Effect.gen(function* () { + const prompt = yield* SessionTools.deferredSystemPrompt({ agent, session }) + + expect(prompt).toContain("Deferred MCP servers available through `search_deferred_tools`:") + expect(prompt).toContain("- posthog: 2 tools") + }), + ) + directIt.instance("keeps MCP tools direct when tool search is disabled", () => Effect.gen(function* () { const tools = yield* resolveTools()