diff --git a/packages/opencode/src/mcp/catalog.ts b/packages/opencode/src/mcp/catalog.ts index 3f985eeb94dc..3b8584562f96 100644 --- a/packages/opencode/src/mcp/catalog.ts +++ b/packages/opencode/src/mcp/catalog.ts @@ -3,9 +3,10 @@ import { CallToolResultSchema, ListToolsResultSchema, ToolSchema, + type CallToolResult, type Tool as MCPToolDef, } from "@modelcontextprotocol/sdk/types.js" -import { dynamicTool, jsonSchema, type JSONSchema7, type Tool } from "ai" +import { dynamicTool, jsonSchema, type JSONSchema7, type Tool, type ToolExecutionOptions } from "ai" import { Effect } from "effect" const DEFAULT_TIMEOUT = 30_000 @@ -50,38 +51,48 @@ export function convertTool(mcpTool: MCPToolDef, client: Client, timeout?: numbe return dynamicTool({ description: mcpTool.description ?? "", inputSchema: jsonSchema(inputSchema), - execute: async (args: unknown, options) => { - const result = await client.callTool( - { - name: mcpTool.name, - arguments: (args || {}) as Record, - }, - CallToolResultSchema, - { - resetTimeoutOnProgress: true, - signal: options.abortSignal, - timeout, - // The MCP SDK only sends a progress token when this hook is present, enabling timeout resets. - onprogress: () => {}, - }, - ) - if (result.isError) - throw new Error( - result.content - .flatMap((item) => (item.type === "text" ? [item.text] : [])) - .filter((text) => text.trim()) - .join("\n\n") || "MCP tool returned an error", - ) - if (result.content.length > 0 || result.structuredContent === undefined || result.structuredContent === null) - return result - return { - ...result, - content: [{ type: "text" as const, text: JSON.stringify(result.structuredContent) }], - } - }, + execute: (args, options) => callTool(mcpTool, client, args, options, timeout), }) } +export async function callTool( + mcpTool: MCPToolDef, + client: Client, + args: unknown, + options: Pick, + timeout?: number, + meta?: Record, +): Promise { + const result = await client.callTool( + { + name: mcpTool.name, + arguments: (args || {}) as Record, + ...(meta === undefined ? {} : { _meta: meta }), + }, + CallToolResultSchema, + { + resetTimeoutOnProgress: true, + signal: options.abortSignal, + timeout, + // The MCP SDK only sends a progress token when this hook is present, enabling timeout resets. + onprogress: () => {}, + }, + ) + if (result.isError) + throw new Error( + result.content + .flatMap((item) => (item.type === "text" ? [item.text] : [])) + .filter((text) => text.trim()) + .join("\n\n") || "MCP tool returned an error", + ) + if (result.content.length > 0 || result.structuredContent === undefined || result.structuredContent === null) + return result + return { + ...result, + content: [{ type: "text" as const, text: JSON.stringify(result.structuredContent) }], + } +} + export function fetch( clientName: string, client: Client, diff --git a/packages/opencode/src/session/tools.ts b/packages/opencode/src/session/tools.ts index 0f401c7562fa..72489981f5e2 100644 --- a/packages/opencode/src/session/tools.ts +++ b/packages/opencode/src/session/tools.ts @@ -399,14 +399,17 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: { run.promise( Effect.gen(function* () { const ctx = context(args, opts) + const hook: { args: typeof args; _meta?: Record } = { args } yield* plugin.trigger( "tool.execute.before", { tool: key, sessionID: ctx.sessionID, callID: opts.toolCallId }, - { args }, + hook, ) const result: Awaited>> = yield* Effect.gen(function* () { yield* ctx.ask({ permission: key, metadata: {}, patterns: ["*"], always: ["*"] }) - return yield* Effect.promise(() => execute(args, opts)) + return yield* Effect.promise(() => + McpCatalog.callTool(entry.def, entry.client, args, opts, entry.timeout, hook._meta), + ) }).pipe( Effect.withSpan("Tool.execute", { attributes: { diff --git a/packages/opencode/src/tool/code-mode.ts b/packages/opencode/src/tool/code-mode.ts index 332d4b43f150..4903eea8e4c7 100644 --- a/packages/opencode/src/tool/code-mode.ts +++ b/packages/opencode/src/tool/code-mode.ts @@ -1,5 +1,5 @@ import * as Tool from "./tool" -import { CallToolResultSchema, type CallToolResult } from "@modelcontextprotocol/sdk/types.js" +import { type CallToolResult } from "@modelcontextprotocol/sdk/types.js" import { Cause, Effect, Schema } from "effect" import { CodeMode, Tool as SandboxTool, toolError } from "@opencode-ai/codemode" import { MCP } from "@/mcp" @@ -138,35 +138,24 @@ const invokeChildTool = Effect.fn("CodeMode.invokeChildTool")(function* (input: callID: string ctx: Tool.Context }) { + const hook: { args: typeof input.args; _meta?: Record } = { args: input.args } yield* input.plugin.trigger( "tool.execute.before", { tool: input.entry.key, sessionID: input.ctx.sessionID, callID: input.callID }, - { args: input.args }, + hook, ) const result: CallToolResult = yield* Effect.gen(function* () { yield* input.ctx.ask({ permission: input.entry.key, metadata: {}, patterns: ["*"], always: ["*"] }) - // 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 - }) + return yield* Effect.promise(() => + McpCatalog.callTool( + input.entry.tool.def, + input.entry.tool.client, + input.args, + { abortSignal: input.ctx.abort }, + input.entry.tool.timeout, + hook._meta, + ), + ) }).pipe( Effect.withSpan("Tool.execute", { attributes: { diff --git a/packages/opencode/test/mcp/catalog.test.ts b/packages/opencode/test/mcp/catalog.test.ts index 7b0d6403bb16..290d5bbf3a07 100644 --- a/packages/opencode/test/mcp/catalog.test.ts +++ b/packages/opencode/test/mcp/catalog.test.ts @@ -105,3 +105,32 @@ test("preserves output schema validation across paginated tool discovery", async await Promise.all([client.close(), server.close()]) } }) + +test("forwards request metadata through the MCP transport", async () => { + const traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + const server = new Server({ name: "metadata", version: "1.0.0" }, { capabilities: { tools: {} } }) + let request: { arguments?: Record; _meta?: Record } | undefined + server.setRequestHandler(CallToolRequestSchema, ({ params }) => { + request = params + return Promise.resolve({ content: [{ type: "text", text: "ok" }] }) + }) + + const client = new Client({ name: "metadata-test", version: "1.0.0" }) + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair() + await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]) + + try { + await McpCatalog.callTool(mcpTool(), client, { target: "screen" }, options, undefined, { + traceparent, + "com.example/correlation-id": "request-1", + }) + expect(request?.arguments).toEqual({ target: "screen" }) + expect(request?._meta).toMatchObject({ + traceparent, + "com.example/correlation-id": "request-1", + }) + expect(request?._meta?.progressToken).toBeDefined() + } finally { + await Promise.all([client.close(), server.close()]) + } +}) diff --git a/packages/opencode/test/session/tools.test.ts b/packages/opencode/test/session/tools.test.ts new file mode 100644 index 000000000000..88a7598eb8a9 --- /dev/null +++ b/packages/opencode/test/session/tools.test.ts @@ -0,0 +1,82 @@ +import { expect } from "bun:test" +import { Agent } from "@/agent/agent" +import { RuntimeFlags } from "@/effect/runtime-flags" +import { MCP } from "@/mcp" +import { Permission } from "@/permission" +import { Plugin } from "@/plugin" +import { Provider } from "@/provider/provider" +import { Session } from "@/session/session" +import { SessionTools } from "@/session/tools" +import { MessageID, SessionID } from "@/session/schema" +import { ToolRegistry } from "@/tool/registry" +import { Truncate } from "@/tool/truncate" +import type { TaskPromptOps } from "@/tool/task" +import type { SessionV1 } from "@opencode-ai/core/v1/session" +import { CallToolResultSchema, type CallToolRequest, type Tool as MCPToolDef } from "@modelcontextprotocol/sdk/types.js" +import type { Client } from "@modelcontextprotocol/sdk/client/index.js" +import { Effect, Layer } from "effect" +import { it } from "../lib/effect" + +it.effect("forwards plugin metadata to normal MCP tool calls", () => { + const traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + let request: CallToolRequest["params"] | undefined + const def = { + name: "tool", + description: "Test tool", + inputSchema: { type: "object", properties: {} }, + } as MCPToolDef + const client = { + getServerCapabilities: () => undefined, + callTool: async (params: CallToolRequest["params"], schema: typeof CallToolResultSchema) => { + request = params + return schema.parse({ content: [{ type: "text", text: "ok" }] }) + }, + } as unknown as Client + const trigger = ((name: unknown, _input: unknown, output: { _meta?: Record }) => + Effect.sync(() => { + if (name === "tool.execute.before") output._meta = { traceparent } + return output + })) as Plugin.Interface["trigger"] + const layer = Layer.mergeAll( + Layer.mock(Plugin.Service, { trigger }), + Layer.mock(Permission.Service, { ask: () => Effect.void }), + Layer.mock(ToolRegistry.Service, { tools: () => Effect.succeed([]) }), + Layer.mock(MCP.Service, { + tools: () => Effect.succeed({ server_tool: { def, client } }), + clients: () => Effect.succeed({}), + }), + Layer.mock(Truncate.Service, { + output: (text: string) => Effect.succeed({ content: text, truncated: false as const }), + }), + RuntimeFlags.layer({ experimentalCodeMode: false }), + ) + + return Effect.gen(function* () { + const tools = yield* SessionTools.resolve({ + agent: { name: "build", permission: [] } as unknown as Agent.Info, + model: { providerID: "test", api: { id: "test", npm: "test" } } as Provider.Model, + session: { id: SessionID.make("ses_meta"), permission: [] } as unknown as Session.Info, + processor: { + message: { id: MessageID.make("msg_meta") } as unknown as SessionV1.Assistant, + updateToolCall: () => Effect.succeed(undefined), + completeToolCall: () => Effect.void, + }, + bypassAgentCheck: false, + messages: [], + promptOps: {} as TaskPromptOps, + }) + + yield* Effect.promise(() => + tools.server_tool!.execute!( + {}, + { + toolCallId: "call_meta", + abortSignal: new AbortController().signal, + messages: [], + }, + ), + ) + + expect(request?._meta).toMatchObject({ traceparent }) + }).pipe(Effect.provide(layer)) +}) diff --git a/packages/opencode/test/tool/code-mode.test.ts b/packages/opencode/test/tool/code-mode.test.ts index 34b3faa610d7..39604c004768 100644 --- a/packages/opencode/test/tool/code-mode.test.ts +++ b/packages/opencode/test/tool/code-mode.test.ts @@ -25,14 +25,15 @@ const ctx: Tool.Context = { function mcpTool( name: string, - handler: (args: Record) => unknown, + handler: (args: Record, meta?: Record) => unknown, inputSchema: Record = { type: "object", properties: {} }, outputSchema?: Record, ): MCP.McpTool { return { def: { name, description: name, inputSchema, ...(outputSchema ? { outputSchema } : {}) } as MCPToolDef, client: { - callTool: async (params: { arguments?: Record }) => handler(params.arguments ?? {}), + callTool: async (params: { arguments?: Record; _meta?: Record }) => + handler(params.arguments ?? {}, params._meta), } as unknown as MCP.McpTool["client"], } } @@ -426,6 +427,39 @@ describe("code mode execute", () => { expect(after!.output).toEqual({ content: [{ type: "text", text: "one" }] }) }) + test("child calls forward plugin metadata independently", async () => { + const received: Array | undefined> = [] + const trigger = ((name: unknown, input: { callID: string }, output: { _meta?: Record }) => + Effect.sync(() => { + if (name === "tool.execute.before") output._meta = { "com.example/call-id": input.callID } + return output + })) as Plugin.Interface["trigger"] + const tool = await build( + { + a_tool: mcpTool("a", (_args, meta) => { + received.push(meta) + return { content: [{ type: "text", text: "one" }] } + }), + b_tool: mcpTool("b", (_args, meta) => { + received.push(meta) + return { content: [{ type: "text", text: "two" }] } + }), + }, + undefined, + undefined, + trigger, + ) + + await Effect.runPromise( + tool.execute({ code: "await tools.a.tool({}); await tools.b.tool({}); return 'done'" }, ctx), + ) + + expect(received).toEqual([ + { "com.example/call-id": "call_code_mode/1" }, + { "com.example/call-id": "call_code_mode/2" }, + ]) + }) + test("a failing before hook fails only that child call as a catchable in-program error", async () => { const trigger = ((name: unknown, input: any, output: unknown) => { if (name === "tool.execute.before" && input.tool === "a_tool") return Effect.die(new Error("hook exploded")) diff --git a/packages/plugin/src/index.ts b/packages/plugin/src/index.ts index edfa0139dfca..3530192c658a 100644 --- a/packages/plugin/src/index.ts +++ b/packages/plugin/src/index.ts @@ -265,7 +265,7 @@ export interface Hooks { ) => Promise "tool.execute.before"?: ( input: { tool: string; sessionID: string; callID: string }, - output: { args: any }, + output: { args: any; _meta?: Record }, ) => Promise "shell.env"?: ( input: { cwd: string; sessionID?: string; callID?: string },