From 02696a7a5ee8b6eaf893214d9e78b5fa60fb28ee Mon Sep 17 00:00:00 2001 From: Aiden Cline Date: Fri, 3 Jul 2026 03:21:42 -0500 Subject: [PATCH] refactor(opencode): consume native MCP tools in code mode --- packages/opencode/src/tool/code-mode.ts | 133 ++++++++---------- packages/opencode/test/tool/code-mode.test.ts | 99 ++++++------- 2 files changed, 100 insertions(+), 132 deletions(-) diff --git a/packages/opencode/src/tool/code-mode.ts b/packages/opencode/src/tool/code-mode.ts index 2de2681cdb..534c33a750 100644 --- a/packages/opencode/src/tool/code-mode.ts +++ b/packages/opencode/src/tool/code-mode.ts @@ -1,6 +1,5 @@ import * as Tool from "./tool" -import type { Tool as AITool, ToolExecutionOptions } from "ai" -import type { Tool as MCPToolDef } from "@modelcontextprotocol/sdk/types.js" +import { CallToolResultSchema, type CallToolResult } from "@modelcontextprotocol/sdk/types.js" import { Cause, Effect, Schema } from "effect" import { CodeMode, @@ -48,57 +47,34 @@ type CatalogEntry = { key: string server: string local: string - description: string - tool: AITool - inputSchema: JsonSchema - outputSchema?: JsonSchema + tool: MCP.McpTool } const toJsonSchema = (schema: unknown): JsonSchema => schema as JsonSchema -function groupByServer( - mcpTools: Record, - servers: readonly string[], - mcpDefs: Record = {}, -): Map { +function groupByServer(mcpTools: Record, servers: readonly string[]): Map { const byLongest = [...servers].sort((a, b) => b.length - a.length) const groups = new Map() for (const key of Object.keys(mcpTools).sort((a, b) => a.localeCompare(b))) { const server = byLongest.find((name) => key.startsWith(name + "_")) ?? (key.includes("_") ? key.slice(0, key.indexOf("_")) : key) const local = server && key.startsWith(server + "_") ? key.slice(server.length + 1) : key - const tool = mcpTools[key]! - const def = mcpDefs[key] - const schema = (tool.inputSchema as { jsonSchema?: unknown } | undefined)?.jsonSchema const entry: CatalogEntry = { path: `${server}.${local}`, key, server, local, - description: tool.description ?? def?.description ?? "", - tool, - inputSchema: def?.inputSchema - ? toJsonSchema(def.inputSchema) - : schema && typeof schema === "object" - ? toJsonSchema(schema) - : { type: "object", properties: {} }, - ...(def?.outputSchema ? { outputSchema: toJsonSchema(def.outputSchema) } : {}), + tool: mcpTools[key]!, } groups.set(server, [...(groups.get(server) ?? []), entry]) } return groups } -export function describeCatalog( - mcpTools: Record, - mcpDefs: Record, - servers: readonly string[], -): string { +export function describeCatalog(mcpTools: Record, servers: readonly string[]): string { return CodeMode.make({ tools: toolTree( - [...groupByServer(mcpTools, servers, mcpDefs).values()] - .flat() - .filter((entry) => entry.tool.execute !== undefined), + [...groupByServer(mcpTools, servers).values()].flat(), () => () => Effect.fail(toolError("Tool preview is not executable.")), ), }).instructions() @@ -112,10 +88,7 @@ const lastSegment = (uri: string) => { const dataUrl = (mime: string, base64: string) => `data:${mime};base64,${base64}` -function projectMcpResult(raw: unknown, collect: (attachment: Attachment) => void): unknown { - if (raw === null || typeof raw !== "object") return raw - const record = raw as { structuredContent?: unknown; content?: unknown } - const content = Array.isArray(record.content) ? record.content : [] +function projectMcpResult(result: CallToolResult, collect: (attachment: Attachment) => void): unknown { const text: string[] = [] let files = 0 let images = 0 @@ -124,53 +97,42 @@ function projectMcpResult(raw: unknown, collect: (attachment: Attachment) => voi if (attachment.mime.startsWith("image/")) images += 1 collect(attachment) } - for (const item of content) { - if (!item || typeof item !== "object") continue - const block = item as Record + for (const block of result.content) { switch (block.type) { case "text": - if (typeof block.text === "string") text.push(block.text) + text.push(block.text) break case "image": case "audio": - if (typeof block.data === "string" && typeof block.mimeType === "string") { - push({ type: "file", mime: block.mimeType, url: dataUrl(block.mimeType, block.data) }) - } + push({ type: "file", mime: block.mimeType, url: dataUrl(block.mimeType, block.data) }) break case "resource": { - const res = block.resource as Record | undefined - if (res && typeof res === "object") { - const mime = typeof res.mimeType === "string" ? res.mimeType : "application/octet-stream" - const uri = typeof res.uri === "string" ? res.uri : undefined - if (typeof res.blob === "string") { - push({ type: "file", mime, url: dataUrl(mime, res.blob), filename: uri ? lastSegment(uri) : undefined }) - } else if (typeof res.text === "string") { - text.push(res.text) - } + if ("text" in block.resource) { + text.push(block.resource.text) + break } + const mime = block.resource.mimeType ?? "application/octet-stream" + push({ type: "file", mime, url: dataUrl(mime, block.resource.blob), filename: lastSegment(block.resource.uri) }) break } case "resource_link": - if (typeof block.uri === "string") { - push({ - type: "file", - mime: typeof block.mimeType === "string" ? block.mimeType : "application/octet-stream", - url: block.uri, - filename: typeof block.name === "string" ? block.name : lastSegment(block.uri), - }) - } + push({ + type: "file", + mime: block.mimeType ?? "application/octet-stream", + url: block.uri, + filename: block.name, + }) break } } - if (record.structuredContent !== undefined && record.structuredContent !== null) return record.structuredContent + if (result.structuredContent !== undefined && result.structuredContent !== null) return result.structuredContent if (text.length > 0) return text.join("\n") if (files > 0) { const noun = files === images ? "image" : "file" return `[${files} ${noun}${files === 1 ? "" : "s"} attached to the result]` } - if (Array.isArray(record.content)) return null // MCP-shaped result with nothing extractable - return raw + return null } type Run = (input: unknown) => Effect.Effect @@ -180,32 +142,51 @@ function toolTree(catalog: readonly CatalogEntry[], run: (entry: CatalogEntry) = for (const entry of catalog) { const namespace = (tree[entry.server] ??= {}) namespace[entry.local] = SandboxTool.make({ - description: entry.description, - input: entry.inputSchema, - output: entry.outputSchema, + description: entry.tool.def.description ?? "", + input: toJsonSchema(entry.tool.def.inputSchema), + output: entry.tool.def.outputSchema ? toJsonSchema(entry.tool.def.outputSchema) : undefined, run: run(entry), }) } return tree } -const invokeChildTool = Effect.fn("CodeMode.invokeChildTool")(function* (input: { +const invokeChildTool = Effect.fn("CodeMode.invokeChildTool")(function* (input: { plugin: Plugin.Interface entry: CatalogEntry - args: any + args: Record callID: string - options: ToolExecutionOptions ctx: Tool.Context - execute: (args: any, options: ToolExecutionOptions) => R | PromiseLike }) { yield* input.plugin.trigger( "tool.execute.before", { tool: input.entry.key, sessionID: input.ctx.sessionID, callID: input.callID }, { args: input.args }, ) - const result: R = yield* Effect.gen(function* () { + const result: CallToolResult = yield* Effect.gen(function* () { yield* input.ctx.ask({ permission: input.entry.key, metadata: {}, patterns: ["*"], always: ["*"] }) - return yield* Effect.promise(() => Promise.resolve(input.execute(input.args, input.options))) + // Deliberately mirrors McpCatalog.convertTool's transport call so the MCP service stays free of tool-loop concerns. + return yield* Effect.promise(async () => { + const raw = await input.entry.tool.client.callTool( + { name: input.entry.tool.def.name, arguments: input.args }, + CallToolResultSchema, + { + resetTimeoutOnProgress: true, + signal: input.ctx.abort, + timeout: input.entry.tool.timeout, + // The MCP SDK only sends a progress token when this hook is present, enabling timeout resets. + onprogress: () => {}, + }, + ) + if (raw.isError) + throw new Error( + raw.content + .flatMap((item) => (item.type === "text" ? [item.text] : [])) + .filter((text) => text.trim()) + .join("\n\n") || "MCP tool returned an error", + ) + return raw + }) }).pipe( Effect.withSpan("Tool.execute", { attributes: { @@ -248,9 +229,7 @@ export const CodeModeTool = Tool.define( const ruleset = Permission.merge(agent.permission, session.permission ?? []) const mcpTools = Permission.visibleTools(yield* mcp.tools(), ruleset) const servers = Object.keys(yield* mcp.clients()).map(McpCatalog.sanitize) - const catalog = [...groupByServer(mcpTools, servers, yield* mcp.defs()).values()] - .flat() - .filter((entry) => entry.tool.execute !== undefined) + const catalog = [...groupByServer(mcpTools, servers).values()].flat() const calls: CallEntry[] = [] const attachments: Attachment[] = [] @@ -262,16 +241,14 @@ export const CodeModeTool = Tool.define( const callTool = (entry: CatalogEntry) => (input: unknown) => Effect.gen(function* () { childCalls += 1 - const raw = yield* invokeChildTool({ + const result = yield* invokeChildTool({ plugin, entry, - args: input ?? {}, + args: (input ?? {}) as Record, callID: `${ctx.callID ?? entry.key}/${childCalls}`, - options: { toolCallId: ctx.callID ?? entry.key, abortSignal: ctx.abort, messages: [] }, ctx, - execute: entry.tool.execute!, }) - return projectMcpResult(raw, collect) + return projectMcpResult(result, collect) }).pipe( Effect.catchCause((cause) => { if (Cause.hasInterruptsOnly(cause)) return Effect.interrupt diff --git a/packages/opencode/test/tool/code-mode.test.ts b/packages/opencode/test/tool/code-mode.test.ts index 9c56040ef4..2418c76bf8 100644 --- a/packages/opencode/test/tool/code-mode.test.ts +++ b/packages/opencode/test/tool/code-mode.test.ts @@ -14,9 +14,7 @@ import { Plugin } from "@/plugin" import { Session } from "@/session/session" import { Tool } from "@/tool/tool" import * as Truncate from "@/tool/truncate" -import { McpCatalog } from "@/mcp/catalog" import { MessageID, SessionID } from "@/session/schema" -import type { Tool as AITool } from "ai" import { Effect, Layer, Schema } from "effect" const ctx: Tool.Context = { @@ -34,16 +32,18 @@ function mcpTool( name: string, handler: (args: Record) => unknown, inputSchema: Record = { type: "object", properties: {} }, -): AITool { - const client = { - callTool: async (params: { arguments?: Record }) => handler(params.arguments ?? {}), + outputSchema?: Record, +): MCP.McpTool { + return { + def: { name, description: name, inputSchema, ...(outputSchema ? { outputSchema } : {}) } as MCPToolDef, + client: { + callTool: async (params: { arguments?: Record }) => handler(params.arguments ?? {}), + } as unknown as MCP.McpTool["client"], } - return McpCatalog.convertTool({ name, description: name, inputSchema } as any, client as any) } function harness(input: { - mcpTools: Record - defs?: Record + mcpTools: Record servers: string[] permission?: PermissionV1.Rule[] trigger?: Plugin.Interface["trigger"] @@ -63,19 +63,17 @@ function harness(input: { }), Layer.mock(MCP.Service, { tools: () => Effect.succeed(input.mcpTools), - defs: () => Effect.succeed(input.defs ?? {}), clients: () => Effect.succeed(Object.fromEntries(input.servers.map((name) => [name, {} as any]))), }), ) } -function serverNames(mcpTools: Record, servers?: string[]) { +function serverNames(mcpTools: Record, servers?: string[]) { return servers ?? [...new Set(Object.keys(mcpTools).map((key) => key.split("_")[0]!))] } function build( - mcpTools: Record, - defs: Record = {}, + mcpTools: Record, servers?: string[], permission?: PermissionV1.Rule[], trigger?: Plugin.Interface["trigger"], @@ -84,18 +82,13 @@ function build( return Effect.runPromise( CodeModeTool.pipe( Effect.flatMap(Tool.init), - Effect.provide(harness({ mcpTools, defs, servers: names, permission, trigger })), + Effect.provide(harness({ mcpTools, servers: names, permission, trigger })), ), ) } -function describeFor( - mcpTools: Record, - defs: Record = {}, - servers?: string[], - permission: PermissionV1.Rule[] = [], -) { - return describeCatalog(Permission.visibleTools(mcpTools, permission), defs, serverNames(mcpTools, servers)) +function describeFor(mcpTools: Record, servers?: string[], permission: PermissionV1.Rule[] = []) { + return describeCatalog(Permission.visibleTools(mcpTools, permission), serverNames(mcpTools, servers)) } describe("code mode execute", () => { @@ -106,26 +99,29 @@ describe("code mode execute", () => { }) test("groups multi-underscore server names by longest matching prefix", () => { - const description = describeFor({ my_server_do_thing: mcpTool("do_thing", () => "") }, {}, ["my_server"]) + const description = describeFor({ my_server_do_thing: mcpTool("do_thing", () => "") }, ["my_server"]) expect(description).toContain("- my_server (1 tool)") expect(description).toContain("tools.my_server.do_thing(") }) test("groupByServer uses the whole key as the server name when it has no underscore", () => { - const description = describeFor({ standalone: mcpTool("standalone", () => "") }, {}, []) + const description = describeFor({ standalone: mcpTool("standalone", () => "") }, []) expect(description).toContain("- standalone (1 tool)") expect(description).toContain("tools.standalone.standalone(") }) test("describeCatalog carries the raw MCP schemas for rendering", () => { - const defs: Record = { - weather_current: { - name: "current", - inputSchema: { type: "object", properties: { city: { type: "string" } }, required: ["city"] }, - outputSchema: { type: "object", properties: { tempC: { type: "number" } }, required: ["tempC"] }, - } as any, - } - const description = describeFor({ weather_current: mcpTool("current", () => "") }, defs, ["weather"]) + const description = describeFor( + { + weather_current: mcpTool( + "current", + () => "", + { type: "object", properties: { city: { type: "string" } }, required: ["city"] }, + { type: "object", properties: { tempC: { type: "number" } }, required: ["tempC"] }, + ), + }, + ["weather"], + ) expect(description).toContain("tools.weather.current(input: { city: string }): Promise<{ tempC: number }>") }) @@ -167,43 +163,42 @@ describe("code mode execute", () => { }) test("signatures render the declared outputSchema as the return type", () => { - const defs: Record = { - weather_current: { - name: "current", - inputSchema: { type: "object", properties: { city: { type: "string" } }, required: ["city"] }, - outputSchema: { + const description = describeFor({ + weather_current: mcpTool( + "current", + () => "", + { type: "object", properties: { city: { type: "string" } }, required: ["city"] }, + { type: "object", properties: { tempC: { type: "number" }, summary: { type: "string" } }, required: ["tempC"], }, - } as any, - } - const description = describeFor({ weather_current: mcpTool("current", () => "") }, defs) + ), + }) expect(description).toContain( "tools.weather.current(input: { city: string }): Promise<{ tempC: number; summary?: string }>", ) }) test("large catalogs inline a budgeted PARTIAL list plus runtime search", async () => { - const tools: Record = {} + const tools: Record = {} const filler = "a searchable description of this operation that consumes catalog budget ".repeat(3) for (let i = 0; i < 150; i++) { - const client = { callTool: async () => ({ content: [] }) } - tools[`alpha_op_${i}`] = McpCatalog.convertTool( - { + tools[`alpha_op_${i}`] = { + def: { name: `op_${i}`, description: `${filler}${i}`, inputSchema: { type: "object", properties: { value: { type: "string" }, count: { type: "number" } } }, - } as any, - client as any, - ) + } as MCPToolDef, + client: { callTool: async () => ({ content: [] }) } as unknown as MCP.McpTool["client"], + } } tools["zeta_only_tool"] = mcpTool("only_tool", () => "", { type: "object", properties: { topic: { type: "string", description: "Subject to look up" } }, required: ["topic"], }) - const description = describeFor(tools, {}, ["alpha", "zeta"]) + const description = describeFor(tools, ["alpha", "zeta"]) expect(description).toContain("Available tools (PARTIAL - ") expect(description).toMatch(/- alpha \(150 tools, \d+ shown\)/) @@ -216,7 +211,7 @@ describe("code mode execute", () => { expect(description).toContain("tools.alpha.op_0(") expect(description).not.toContain("tools.alpha.op_99(") - const tool = await build(tools, {}, ["alpha", "zeta"]) + const tool = await build(tools, ["alpha", "zeta"]) const out = await Effect.runPromise( tool.execute({ code: "return await tools.$codemode.search({ query: 'only tool', limit: 3 })" }, ctx), ) @@ -385,7 +380,6 @@ describe("code mode execute", () => { a_tool: mcpTool("a", () => ({ content: [{ type: "text", text: "one" }] })), b_tool: mcpTool("b", () => ({ content: [{ type: "text", text: "two" }] })), }, - {}, undefined, undefined, trigger, @@ -421,7 +415,6 @@ describe("code mode execute", () => { } const tool = await build( { a_tool: mcpTool("a", record("a")), b_tool: mcpTool("b", record("b")) }, - {}, undefined, undefined, trigger, @@ -520,7 +513,7 @@ describe("code mode execute", () => { media_mixed: mcpTool("mixed", () => ({ content: [ { type: "image", data: "PNG3", mimeType: "image/png" }, - { type: "resource_link", uri: "file:///tmp/report.pdf", mimeType: "application/pdf" }, + { type: "resource_link", uri: "file:///tmp/report.pdf", name: "report.pdf", mimeType: "application/pdf" }, ], })), }) @@ -633,7 +626,7 @@ describe("code mode permission visibility", () => { github_create_issue: mcpTool("create_issue", ok), github_list_issues: mcpTool("list_issues", ok), } - const description = describeFor(mcpTools, {}, ["github"], [deny("github_create_issue")]) + const description = describeFor(mcpTools, ["github"], [deny("github_create_issue")]) expect(description).toContain("tools.github.list_issues(") expect(description).not.toContain("create_issue") expect(description).toContain("- github (1 tool)") @@ -644,7 +637,7 @@ describe("code mode permission visibility", () => { github_create_issue: mcpTool("create_issue", ok), github_list_issues: mcpTool("list_issues", ok), } - const description = describeFor(mcpTools, {}, ["github"], [askRule("github_create_issue")]) + const description = describeFor(mcpTools, ["github"], [askRule("github_create_issue")]) expect(description).toContain("tools.github.create_issue(") expect(description).toContain("tools.github.list_issues(") expect(description).toContain("- github (2 tools)") @@ -660,7 +653,6 @@ describe("code mode permission visibility", () => { }), github_list_issues: mcpTool("list_issues", ok), }, - {}, ["github"], [deny("github_create_issue")], ) @@ -685,7 +677,6 @@ describe("code mode permission visibility", () => { const askCtx: Tool.Context = { ...ctx, ask: (req) => Effect.sync(() => void asked.push(req.permission)) } const tool = await build( { github_list_issues: mcpTool("list_issues", ok) }, - {}, ["github"], [askRule("github_list_issues")], )