diff --git a/packages/opencode/src/session/code-mode.ts b/packages/opencode/src/session/code-mode.ts index 74a1a37632..f4894c1361 100644 --- a/packages/opencode/src/session/code-mode.ts +++ b/packages/opencode/src/session/code-mode.ts @@ -1,13 +1,13 @@ import { Tool } from "@/tool/tool" import { EffectBridge } from "@/effect/bridge" -import type { Tool as AITool } from "ai" +import { asSchema, type Tool as AITool, type JSONSchema7 } from "ai" import { Effect, Schema } from "effect" export const CODE_MODE_TOOL = "execute" export const Parameters = Schema.Struct({ code: Schema.String.annotate({ - description: "JavaScript to run. Call tools as `await tools.(args)` and `return` the final value.", + description: "JavaScript to run. Call tools as `await tools..(input)` and `return` the final value.", }), }) @@ -22,19 +22,90 @@ const AsyncFunction = Object.getPrototypeOf(async function () {}).constructor as new (...args: string[]): (...args: unknown[]) => Promise } -function describe(mcpTools: Record) { - const names = Object.keys(mcpTools).sort((a, b) => a.localeCompare(b)) - return [ +const IDENTIFIER = /^[A-Za-z_$][A-Za-z0-9_$]*$/ + +type NamespacedTool = { local: string; key: string; tool: AITool } + +/** + * Group the flat `server_tool` catalog into per-server namespaces for display. + * `servers` are the sanitized MCP client names; the longest matching prefix wins + * so a server named `a_b` is preferred over `a` for the key `a_b_tool`. Routing + * never depends on this split — it re-joins `${server}_${local}` back to the key. + */ +export 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.slice(0, key.indexOf("_")) + const local = server && key.startsWith(server + "_") ? key.slice(server.length + 1) : key + const entry = groups.get(server) ?? [] + entry.push({ local, key, tool: mcpTools[key]! }) + groups.set(server, entry) + } + return groups +} + +const access = (segment: string) => (IDENTIFIER.test(segment) ? `.${segment}` : `[${JSON.stringify(segment)}]`) + +function jsonType(def: JSONSchema7 | boolean | undefined): string { + if (!def || typeof def === "boolean") return "any" + if (Array.isArray(def.enum)) return def.enum.map((value) => JSON.stringify(value)).join(" | ") + const type = Array.isArray(def.type) ? def.type[0] : def.type + switch (type) { + case "integer": + return "number" + case "array": + return "any[]" + case undefined: + return "any" + default: + return type + } +} + +function inputHint(tool: AITool): string { + try { + const schema = asSchema(tool.inputSchema).jsonSchema as JSONSchema7 | undefined + const props = schema?.properties + if (!props || typeof props !== "object") return "input" + const required = new Set(Array.isArray(schema?.required) ? schema.required : []) + const fields = Object.entries(props).map( + ([name, def]) => `${name}${required.has(name) ? "" : "?"}: ${jsonType(def as JSONSchema7)}`, + ) + return fields.length > 0 ? `{ ${fields.join("; ")} }` : "{}" + } catch { + return "input" + } +} + +const firstLine = (text: string | undefined) => (text ?? "").split("\n", 1)[0]!.trim() + +export function describe(groups: Map): string { + const lines = [ "Execute JavaScript with access to connected MCP tools.", - "Each tool is callable as `await tools.(args)`; `return` the final value.", - names.length > 0 ? `Available tools: ${names.join(", ")}` : "No MCP tools are currently connected.", - ].join("\n") + "Every connected MCP server is a namespace on `tools`. Call a tool with `await tools..(input)`; each returns a Promise.", + "Compose multiple calls in one program and `return` the final value — intermediate results stay in the sandbox and never re-enter the conversation.", + ] + if (groups.size === 0) { + lines.push("", "No MCP servers are currently connected.") + return lines.join("\n") + } + lines.push("", "Available namespaces:") + for (const [server, tools] of [...groups].sort(([a], [b]) => a.localeCompare(b))) { + lines.push("", `// ${server}`) + for (const { local, tool } of tools) { + const signature = `tools${access(server)}${access(local)}(${inputHint(tool)})` + const summary = firstLine(tool.description) + lines.push(summary ? `${signature} // ${summary}` : signature) + } + } + return lines.join("\n") } /** * Reduce an MCP tool result to the value the program should see: structured * content when present, otherwise the joined text blocks, otherwise the raw - * result. Mirrors how the model-facing output is derived elsewhere. + * result. */ export function toolResultValue(result: unknown): unknown { if (result === null || typeof result !== "object") return result @@ -72,25 +143,26 @@ function errorMessage(error: unknown): string { } } -export function define(mcpTools: Record) { +export function define(mcpTools: Record, servers: readonly string[]) { + const groups = groupByServer(mcpTools, servers) return Tool.define( CODE_MODE_TOOL, Effect.succeed>({ - description: describe(mcpTools), + description: describe(groups), parameters: Parameters, execute: Effect.fn("CodeMode.execute")(function* (params, ctx) { const run = yield* EffectBridge.make() const calls: string[] = [] - // Each `tools.(args)` call runs the native MCP tool through the - // permission gate, so approving `execute` does not approve every child call. - const invoke = (name: string, tool: AITool, args: unknown) => + // Each tool call runs the native MCP tool through the permission gate, so + // approving `execute` does not approve every child call. + const invoke = (key: string, tool: AITool, args: unknown) => Effect.gen(function* () { - yield* ctx.ask({ permission: name, metadata: {}, patterns: ["*"], always: ["*"] }) + yield* ctx.ask({ permission: key, metadata: {}, patterns: ["*"], always: ["*"] }) const result = yield* Effect.promise(() => Promise.resolve( tool.execute!(args ?? {}, { - toolCallId: ctx.callID ?? name, + toolCallId: ctx.callID ?? key, abortSignal: ctx.abort, messages: [], }), @@ -99,21 +171,30 @@ export function define(mcpTools: Record) { return toolResultValue(result) }) + // `tools..(args)` — the server/tool split is cosmetic; routing + // re-joins `${server}_${tool}` back into the original flat catalog key. + const namespace = (server: string) => + new Proxy(Object.create(null) as Record, { + get(_target, prop) { + if (typeof prop !== "string" || prop === "then") return undefined + const key = `${server}_${prop}` + const tool = mcpTools[key] + if (!tool || !tool.execute) { + return () => { + throw new Error(`Unknown tool 'tools.${server}.${prop}'. Available: ${Object.keys(mcpTools).join(", ")}`) + } + } + return (args: unknown) => { + calls.push(key) + return run.promise(invoke(key, tool, args)) + } + }, + }) + const tools = new Proxy(Object.create(null) as Record, { get(_target, prop) { if (typeof prop !== "string" || prop === "then") return undefined - const tool = mcpTools[prop] - if (!tool || !tool.execute) { - return () => { - throw new Error( - `Unknown tool '${prop}'. Available tools: ${Object.keys(mcpTools).join(", ") || "(none)"}`, - ) - } - } - return (args: unknown) => { - calls.push(prop) - return run.promise(invoke(prop, tool, args)) - } + return namespace(prop) }, }) diff --git a/packages/opencode/src/session/tools.ts b/packages/opencode/src/session/tools.ts index e7eb1ff9b0..115c1bd314 100644 --- a/packages/opencode/src/session/tools.ts +++ b/packages/opencode/src/session/tools.ts @@ -22,6 +22,7 @@ import { ProviderV2 } from "@opencode-ai/core/provider" import { ModelV2 } from "@opencode-ai/core/model" import { isRecord } from "@/util/record" import { RuntimeFlags } from "@/effect/runtime-flags" +import { McpCatalog } from "@/mcp/catalog" import * as CodeModeTool from "./code-mode" const MCP_RESOURCE_TOOLS = { @@ -93,9 +94,16 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: { // When code mode is enabled and MCP tools are present, expose them through the // single code-mode `execute` tool instead of registering each MCP tool directly // (see the early return below). Code mode is experimental and off by default. + // Sanitized client names give code mode the per-server namespaces; tool keys are + // `sanitize(server)_sanitize(tool)`, so the names match the catalog key prefixes. const codeModeTool = flags.experimentalCodeMode && Object.keys(mcpTools).length > 0 - ? yield* Tool.init(yield* CodeModeTool.define(mcpTools)) + ? yield* Tool.init( + yield* CodeModeTool.define( + mcpTools, + Object.keys(yield* mcp.clients()).map(McpCatalog.sanitize), + ), + ) : undefined const registryTools = yield* registry.tools({ modelID: ModelV2.ID.make(input.model.api.id), diff --git a/packages/opencode/test/session/code-mode.test.ts b/packages/opencode/test/session/code-mode.test.ts index 623c2881e8..02b8524971 100644 --- a/packages/opencode/test/session/code-mode.test.ts +++ b/packages/opencode/test/session/code-mode.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test" -import { Parameters, define, formatValue, toolResultValue } from "@/session/code-mode" +import { Parameters, define, describe as describeTools, formatValue, groupByServer, toolResultValue } from "@/session/code-mode" import { Agent } from "@/agent/agent" import { Tool } from "@/tool/tool" import * as Truncate from "@/tool/truncate" @@ -21,14 +21,15 @@ const ctx: Tool.Context = { // Build a real MCP-derived AI SDK tool over a fake transport, so the proxy exercises // the same `convertTool` execution path that `mcp.tools()` produces at runtime. -function mcpTool(name: string, handler: (args: Record) => unknown): AITool { +function mcpTool( + name: string, + handler: (args: Record) => unknown, + inputSchema: Record = { type: "object", properties: {} }, +): AITool { const client = { callTool: async (params: { arguments?: Record }) => handler(params.arguments ?? {}), } - return McpCatalog.convertTool( - { name, description: name, inputSchema: { type: "object", properties: {} } } as any, - client as any, - ) + return McpCatalog.convertTool({ name, description: name, inputSchema } as any, client as any) } // Truncate echoes its input so assertions read the exact program output. Agent.get is @@ -40,8 +41,11 @@ const layer = Layer.mergeAll( Layer.succeed(Agent.Service, Agent.Service.of({ get: () => Effect.succeed({ name: "build" } as any) } as any)), ) -function build(mcpTools: Record) { - return Effect.runPromise(define(mcpTools).pipe(Effect.flatMap(Tool.init), Effect.provide(layer))) +// Derive sanitized server namespaces from the catalog keys, mirroring how +// session/tools.ts passes `Object.keys(mcp.clients()).map(sanitize)`. +function build(mcpTools: Record, servers?: string[]) { + const names = servers ?? [...new Set(Object.keys(mcpTools).map((key) => key.split("_")[0]!))] + return Effect.runPromise(define(mcpTools, names).pipe(Effect.flatMap(Tool.init), Effect.provide(layer))) } describe("code mode execute", () => { @@ -51,9 +55,32 @@ describe("code mode execute", () => { await expect(Effect.runPromise(decode({}))).rejects.toThrow() }) - test("lists available tools in the description", async () => { - const tool = await build({ beta_b: mcpTool("b", () => "b"), alpha_a: mcpTool("a", () => "a") }) - expect(tool.description).toContain("Available tools: alpha_a, beta_b") + test("describes tools grouped into per-server namespaces with signatures", () => { + const description = describeTools( + groupByServer( + { + github_create_issue: mcpTool("create_issue", () => "", { + type: "object", + properties: { title: { type: "string" }, body: { type: "string" } }, + required: ["title"], + }), + linear_search: mcpTool("search", () => ""), + }, + ["github", "linear"], + ), + ) + + expect(description).toContain("await tools..(input)") + expect(description).toContain("// github") + expect(description).toContain("tools.github.create_issue({ title: string; body?: string })") + expect(description).toContain("// linear") + expect(description).toContain("tools.linear.search") + }) + + test("groups multi-underscore server names by longest matching prefix", () => { + const groups = groupByServer({ my_server_do_thing: mcpTool("do_thing", () => "") }, ["my_server"]) + expect([...groups.keys()]).toEqual(["my_server"]) + expect(groups.get("my_server")![0]).toMatchObject({ local: "do_thing", key: "my_server_do_thing" }) }) test("runs plain JavaScript and returns the value as text", async () => { @@ -63,7 +90,7 @@ describe("code mode execute", () => { expect(output.metadata.toolCalls).toEqual([]) }) - test("calls an MCP tool and flows its text result back into the program", async () => { + test("calls a namespaced MCP tool and flows its text result back into the program", async () => { const seen: Record[] = [] const tool = await build({ greeter_hello: mcpTool("hello", (args) => { @@ -73,7 +100,7 @@ describe("code mode execute", () => { }) const output = await Effect.runPromise( - tool.execute({ code: "const r = await tools.greeter_hello({ name: 'world' }); return r.toUpperCase()" }, ctx), + tool.execute({ code: "const r = await tools.greeter.hello({ name: 'world' }); return r.toUpperCase()" }, ctx), ) expect(seen).toEqual([{ name: "world" }]) @@ -93,8 +120,8 @@ describe("code mode execute", () => { tool.execute( { code: ` - const first = await tools.math_add({ a: 1, b: 2 }) - const second = await tools.math_add({ a: first.sum, b: 10 }) + const first = await tools.math.add({ a: 1, b: 2 }) + const second = await tools.math.add({ a: first.sum, b: 10 }) return { total: second.sum } `, }, @@ -114,7 +141,7 @@ describe("code mode execute", () => { const output = await Effect.runPromise( tool.execute( - { code: "const [a, b] = await Promise.all([tools.echo_one({}), tools.echo_two({})]); return a + b" }, + { code: "const [a, b] = await Promise.all([tools.echo.one({}), tools.echo.two({})]); return a + b" }, ctx, ), ) @@ -132,9 +159,9 @@ describe("code mode execute", () => { test("reports an unknown tool with the available names", async () => { const tool = await build({ known_tool: mcpTool("tool", () => "ok") }) - const output = await Effect.runPromise(tool.execute({ code: "return await tools.missing({})" }, ctx)) + const output = await Effect.runPromise(tool.execute({ code: "return await tools.known.missing({})" }, ctx)) expect(output.metadata.error).toBe(true) - expect(output.output).toContain("Unknown tool 'missing'") + expect(output.output).toContain("Unknown tool 'tools.known.missing'") expect(output.output).toContain("known_tool") }) @@ -144,7 +171,7 @@ describe("code mode execute", () => { }) const output = await Effect.runPromise( tool.execute( - { code: "try { await tools.bad_tool({}) } catch (e) { return 'caught: ' + e.message }" }, + { code: "try { await tools.bad.tool({}) } catch (e) { return 'caught: ' + e.message }" }, ctx, ), ) @@ -158,7 +185,7 @@ describe("code mode execute", () => { const tool = await build({ a_tool: mcpTool("a", ok), b_tool: mcpTool("b", ok) }) await Effect.runPromise( - tool.execute({ code: "await tools.a_tool({}); await tools.b_tool({}); return 'done'" }, permissionCtx), + tool.execute({ code: "await tools.a.tool({}); await tools.b.tool({}); return 'done'" }, permissionCtx), ) expect(asked.map((req: any) => req.permission)).toEqual(["a_tool", "b_tool"])