diff --git a/packages/opencode/src/session/mcp-tools.ts b/packages/opencode/src/session/mcp-tools.ts index 7e89c9367f2..a79742ecdae 100644 --- a/packages/opencode/src/session/mcp-tools.ts +++ b/packages/opencode/src/session/mcp-tools.ts @@ -43,16 +43,19 @@ export const DEFERRED_TOOL_SYSTEM_PROMPT = `Deferred tools are separate from dir - Use \`call_deferred_tool\` only with a \`tool_id\` copied from \`search_deferred_tools\` results. - Never use deferred tools for direct tools already listed in the current tool list, including shell, file, search, edit, web, task, LSP, question, or apply_patch tools.` -interface Input { - agent: Agent.Info +interface Input extends DeferredInput { model: Provider.Model - session: Session.Info processor: Pick bypassAgentCheck: boolean - messages: SessionV1.WithParts[] promptOps: TaskPromptOps } +interface DeferredInput { + agent: Agent.Info + session: Session.Info + messages: SessionV1.WithParts[] +} + interface DeferredToolDescriptor { id: string description: string @@ -60,11 +63,6 @@ interface DeferredToolDescriptor { searchText: string } -interface Output { - tools: Record - deferredSystemPrompt?: string -} - const SearchDeferredToolsParameters = Schema.Struct({ query: Schema.String.annotate({ description: @@ -104,7 +102,6 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In const permission = yield* Permission.Service const mcp = yield* MCP.Service const truncate = yield* Truncate.Service - const flags = yield* RuntimeFlags.Service const context = (args: Record, options: ToolExecutionOptions): Context => ({ sessionID: input.session.id, @@ -238,36 +235,20 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In if (hasMcpResourceServer) addResourceTools(tools, { input, run, mcp, plugin, truncate, context }) const mcpTools = yield* mcp.tools() - const mcpDisabled = Permission.disabled( - Object.keys(mcpTools), - Permission.merge(input.agent.permission, input.session.permission ?? []), - ) - const userTools = currentUserToolOverrides(input.messages) - const allowedMcpTools = Object.fromEntries( - Object.entries(mcpTools).filter(([key]) => userTools?.[key] !== false && !mcpDisabled.has(key)), - ) - const deferredDescriptors = - flags.experimentalToolSearch && Object.keys(allowedMcpTools).length > 0 - ? yield* deferredToolDescriptors(allowedMcpTools) - : [] - const deferMcpTools = - deferredDescriptors.length > 0 && deferredToolSchemaTokens(deferredDescriptors) >= MIN_DEFERRED_MCP_SCHEMA_TOKENS + const deferred = yield* deferredState(input, mcpTools) - if (deferMcpTools) { + if (deferred) { addDeferredTools(tools, { input, run, plugin, truncate, context, - deferredDescriptors, - allowedMcpTools, + deferredDescriptors: deferred.descriptors, + allowedMcpTools: deferred.allowedMcpTools, executeMcpTool, }) - return { - tools, - deferredSystemPrompt: deferredSystemPrompt(allowedMcpTools, clients), - } satisfies Output + return tools } for (const [key, item] of Object.entries(mcpTools)) { @@ -281,9 +262,37 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In tools[key] = item } - return { tools } satisfies Output + return tools }) +export const systemPrompt = Effect.fn("SessionMcpTools.systemPrompt")(function* (input: DeferredInput) { + const mcp = yield* MCP.Service + const deferred = yield* deferredState(input, yield* mcp.tools()) + if (!deferred) return undefined + return deferredSystemPrompt(deferred.allowedMcpTools, yield* mcp.clients()) +}) + +function deferredState(input: DeferredInput, mcpTools: Record) { + return Effect.gen(function* () { + const flags = yield* RuntimeFlags.Service + if (!flags.experimentalToolSearch) return undefined + + const mcpDisabled = Permission.disabled( + Object.keys(mcpTools), + Permission.merge(input.agent.permission, input.session.permission ?? []), + ) + const userTools = currentUserToolOverrides(input.messages) + const allowedMcpTools = Object.fromEntries( + Object.entries(mcpTools).filter(([key]) => userTools?.[key] !== false && !mcpDisabled.has(key)), + ) + if (Object.keys(allowedMcpTools).length === 0) return undefined + + const descriptors = yield* deferredToolDescriptors(allowedMcpTools) + if (deferredToolSchemaTokens(descriptors) < MIN_DEFERRED_MCP_SCHEMA_TOKENS) return undefined + return { allowedMcpTools, descriptors } + }) +} + function deferredSystemPrompt(allowedMcpTools: Record, clients: Record) { const servers = deferredServerSummaries(Object.keys(allowedMcpTools), Object.keys(clients)) return [ diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 9f8c9efbb3f..41b57a30bf4 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -55,6 +55,7 @@ import { eq } from "drizzle-orm" import { SessionTable } from "@opencode-ai/core/session/sql" import { SessionReminders } from "./reminders" import { SessionTools } from "./tools" +import { SessionMcpTools } from "./mcp-tools" import { LLMEvent } from "@opencode-ai/llm" // @ts-ignore @@ -1223,7 +1224,7 @@ export const layer = Layer.effect( const bypassAgentCheck = lastUserMsg?.parts.some((p) => p.type === "agent") ?? false const promptOps = yield* ops() - const resolvedTools = yield* SessionTools.resolve({ + const tools = yield* SessionTools.resolve({ agent, session, model, @@ -1239,7 +1240,6 @@ export const layer = Layer.effect( Effect.provideService(Truncate.Service, truncate), Effect.provideService(RuntimeFlags.Service, flags), ) - const tools = resolvedTools.tools if (lastUser.format?.type === "json_schema") { tools["StructuredOutput"] = createStructuredOutputTool({ @@ -1260,7 +1260,10 @@ export const layer = Layer.effect( sys.environment(model), instruction.system().pipe(Effect.orDie), sys.mcp(agent, session.permission), - Effect.succeed(resolvedTools.deferredSystemPrompt), + SessionMcpTools.systemPrompt({ agent, session, messages: msgs }).pipe( + Effect.provideService(MCP.Service, mcp), + Effect.provideService(RuntimeFlags.Service, flags), + ), MessageV2.toModelMessagesEffect(msgs, model), ]) const system = [ diff --git a/packages/opencode/src/session/tools.ts b/packages/opencode/src/session/tools.ts index 4a86ddeb7ca..a1dc8db3f7d 100644 --- a/packages/opencode/src/session/tools.ts +++ b/packages/opencode/src/session/tools.ts @@ -108,12 +108,8 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: { }) } - const mcp = yield* SessionMcpTools.resolve(input) - Object.assign(tools, mcp.tools) - return { - tools, - deferredSystemPrompt: mcp.deferredSystemPrompt, - } + Object.assign(tools, yield* SessionMcpTools.resolve(input)) + return tools }) export * as SessionTools from "./tools" diff --git a/packages/opencode/test/session/tools.test.ts b/packages/opencode/test/session/tools.test.ts index 6405c6893b6..f6a220631c6 100644 --- a/packages/opencode/test/session/tools.test.ts +++ b/packages/opencode/test/session/tools.test.ts @@ -11,6 +11,7 @@ import { MCP } from "@/mcp" import { Permission } from "@/permission" import { RuntimeFlags } from "@/effect/runtime-flags" import { SessionTools } from "@/session/tools" +import { SessionMcpTools } from "@/session/mcp-tools" import { MessageID, SessionID } from "@/session/schema" import { Session } from "@/session/session" import { Plugin } from "@/plugin" @@ -132,7 +133,7 @@ const belowThresholdIt = makeIt({ queryDescription: "Natural language analytics query", }) -function resolveToolResult(input: { messages?: SessionV1.WithParts[] } = {}) { +function resolveTools(input: { messages?: SessionV1.WithParts[] } = {}) { return SessionTools.resolve({ agent, model, @@ -144,10 +145,6 @@ function resolveToolResult(input: { messages?: SessionV1.WithParts[] } = {}) { }) } -function resolveTools(input: { messages?: SessionV1.WithParts[] } = {}) { - return resolveToolResult(input).pipe(Effect.map((result) => result.tools)) -} - describe("session.tools", () => { deferredIt.instance("defers MCP tools behind fixed search and call tools", () => Effect.gen(function* () { @@ -184,7 +181,7 @@ describe("session.tools", () => { deferredIt.instance("lists deferred MCP servers in the system prompt", () => Effect.gen(function* () { - const prompt = (yield* resolveToolResult()).deferredSystemPrompt + const prompt = yield* SessionMcpTools.systemPrompt({ agent, session, messages: [] }) expect(prompt).toContain("Deferred MCP servers available through `search_deferred_tools`:") expect(prompt).toContain("- posthog: 2 tools")