diff --git a/packages/codemode/src/codemode.ts b/packages/codemode/src/codemode.ts index 14782d7c8aae..c2569b573dd5 100644 --- a/packages/codemode/src/codemode.ts +++ b/packages/codemode/src/codemode.ts @@ -17,7 +17,7 @@ export type ExecutionLimits = { } /** Controls how much of the tool catalog is inlined in agent instructions. */ -export type DiscoveryOptions = { +export type DiscoveryOptions = ToolRuntime.CatalogSelection & { /** Approximate token budget (chars/4, default 2000) for full catalog entries. */ readonly catalogBudget?: number } @@ -149,7 +149,7 @@ export const make = = {}>( const tools = (options.tools ?? {}) as HostTools> ToolRuntime.assertValidTools(tools) const limits = resolveExecutionLimits(options.limits) - const prepared = ToolRuntime.prepare(tools, options.discovery?.catalogBudget) + const prepared = ToolRuntime.prepare(tools, options.discovery?.catalogBudget, options.discovery) return { catalog: () => prepared.catalog, diff --git a/packages/codemode/src/tool-runtime.ts b/packages/codemode/src/tool-runtime.ts index f4ccc61d4c49..78b27960e8c5 100644 --- a/packages/codemode/src/tool-runtime.ts +++ b/packages/codemode/src/tool-runtime.ts @@ -342,6 +342,27 @@ const visibleDefinitions = (tools: HostTools) => export const catalog = (tools: HostTools): ReadonlyArray => visibleDefinitions(tools).map(({ description }) => description) +/** + * A one-line index entry for tools that are not inlined: `match` is a RegExp source tested + * against the tool name within `namespace`. An entry is listed only while at least one + * matching tool is present and not inlined, so it never advertises an absent tool. + */ +export type CatalogCategory = { + readonly namespace: string + readonly name: string + readonly match: string + readonly description: string +} + +/** Caller-ranked selection for the inlined catalog. */ +export type CatalogSelection = { + /** Tool paths (`namespace.tool`), most important first; order is kept within each namespace. When set, only these are inlined. */ + readonly pinned?: ReadonlyArray + /** Maximum number of pinned tools inlined. Default: as many as fit the budget. */ + readonly pinnedLimit?: number + readonly categories?: ReadonlyArray +} + export type DiscoveryPlan = { readonly catalog: ReadonlyArray readonly instructions: string @@ -481,6 +502,72 @@ export const assertValidTools = (tools: HostTools): void => { } } +/** + * Round-robin fairness: in each round (namespaces alphabetical), every namespace still + * holding un-inlined tools tries to place its next-cheapest line against the shared budget; + * a namespace whose next line does not fit is done - the others keep going - so every + * namespace gets some representation before any namespace gets everything. + */ +const selectRoundRobin = ( + ordered: ReadonlyArray]>, + catalogBudget: number, +): ReadonlyMap> => { + const selections = ordered.map(([namespace, group]) => ({ + namespace, + group, + picked: new Set(), + queue: [...group].sort( + (left, right) => + estimateTokens(catalogLine(left)) - estimateTokens(catalogLine(right)) || left.path.localeCompare(right.path), + ), + })) + let used = 0 + let active = selections.filter((selection) => selection.queue.length > 0) + while (active.length > 0) { + const stillActive: typeof active = [] + for (const selection of active) { + const tool = selection.queue[0]! + const cost = estimateTokens(catalogLine(tool)) + if (used + cost > catalogBudget) continue + selection.queue.shift() + selection.picked.add(tool) + used += cost + if (selection.queue.length > 0) stillActive.push(selection) + } + active = stillActive + } + return new Map( + selections.map(({ namespace, group, picked }) => [namespace, group.filter((tool) => picked.has(tool))]), + ) +} + +/** Caller-ranked selection: pinned paths in order, skipping absent tools and lines that do not fit. */ +const selectPinned = ( + described: ReadonlyArray, + pinned: ReadonlyArray, + pinnedLimit: number, + catalogBudget: number, +): ReadonlyMap> => { + const byPath = new Map(described.map((tool) => [tool.path, tool])) + const picked = new Set() + let used = 0 + for (const path of pinned) { + if (picked.size >= pinnedLimit) break + const tool = byPath.get(path) + if (tool === undefined || picked.has(tool)) continue + const cost = estimateTokens(catalogLine(tool)) + if (used + cost > catalogBudget) continue + picked.add(tool) + used += cost + } + const shown = new Map>() + for (const tool of picked) { + const [namespace = tool.path] = tool.path.split(".") + shown.set(namespace, [...(shown.get(namespace) ?? []), tool]) + } + return shown +} + /** * Budgeted catalog: every namespace is always listed with its tool count; full call * signatures are inlined against the `catalogBudget` (estimated tokens, @@ -491,11 +578,28 @@ export const assertValidTools = (tools: HostTools): void => { * states exactly how comprehensive it is - overall (COMPLETE vs PARTIAL) and per * namespace. Namespace stub lines are never budgeted: every namespace appears with its * tool count even at budget 0. + * + * With `selection.pinned`, the caller's ranking replaces the round-robin: pinned tools + * that exist are inlined in that order until `pinnedLimit` or the budget is reached, and + * nothing else is inlined. `selection.categories` adds an unbudgeted one-line index per + * category of tools that are present but not inlined, so every category stays discoverable by search. */ -export const prepare = (tools: HostTools, catalogBudget = defaultCatalogBudget): DiscoveryPlan => { +export const prepare = ( + tools: HostTools, + catalogBudget = defaultCatalogBudget, + selection: CatalogSelection = {}, +): DiscoveryPlan => { if (!Number.isSafeInteger(catalogBudget) || catalogBudget < 0) { throw new RangeError("discovery.catalogBudget must be a non-negative safe integer") } + const pinnedLimit = selection.pinnedLimit ?? Number.MAX_SAFE_INTEGER + if (!Number.isSafeInteger(pinnedLimit) || pinnedLimit < 0) { + throw new RangeError("discovery.pinnedLimit must be a non-negative safe integer") + } + const categories = (selection.categories ?? []).map((category) => ({ + ...category, + pattern: new RegExp(category.match), + })) const visible = visibleDefinitions(tools) const described = visible.map(({ description }) => description) @@ -509,38 +613,12 @@ export const prepare = (tools: HostTools, catalogBudget = defaultCatalogBu const ordered = [...namespaces].sort(([left], [right]) => left.localeCompare(right)) // Select which signatures fit the budget before emitting, so the list can state - // exactly how comprehensive it is. Round-robin fairness: in each round (namespaces - // alphabetical), every namespace still holding un-inlined tools tries to place its - // next-cheapest line against the shared budget; a namespace whose next line does not - // fit is done - the others keep going - so every namespace gets some representation - // before any namespace gets everything. - const selections = ordered.map(([namespace, group]) => ({ - namespace, - picked: new Set(), - queue: [...group].sort( - (left, right) => - estimateTokens(catalogLine(left)) - estimateTokens(catalogLine(right)) || left.path.localeCompare(right.path), - ), - })) - let used = 0 - let active = selections.filter((selection) => selection.queue.length > 0) - while (active.length > 0) { - const stillActive: typeof active = [] - for (const selection of active) { - const tool = selection.queue[0]! - const cost = estimateTokens(catalogLine(tool)) - if (used + cost > catalogBudget) continue - selection.queue.shift() - selection.picked.add(tool) - used += cost - if (selection.queue.length > 0) stillActive.push(selection) - } - active = stillActive - } - const shown = new Map>( - selections.map(({ namespace, picked }) => [namespace, picked]), - ) - const totalShown = selections.reduce((total, { picked }) => total + picked.size, 0) + // exactly how comprehensive it is. + const shown = + selection.pinned === undefined + ? selectRoundRobin(ordered, catalogBudget) + : selectPinned(described, selection.pinned, pinnedLimit, catalogBudget) + const totalShown = [...shown.values()].reduce((total, picked) => total + picked.length, 0) const complete = totalShown === described.length const empty = described.length === 0 @@ -621,18 +699,33 @@ export const prepare = (tools: HostTools, catalogBudget = defaultCatalogBu "", ) for (const [namespace, group] of ordered) { - const picked = shown.get(namespace)! + const picked = shown.get(namespace) ?? [] const count = `${group.length} tool${group.length === 1 ? "" : "s"}` // Annotate only when a namespace is not fully shown, so a comprehensive // namespace reads cleanly and a truncated one is unambiguous. const label = - picked.size === group.length + picked.length === group.length ? count - : picked.size === 0 + : picked.length === 0 ? `${count}, none shown` - : `${count}, ${picked.size} shown` + : `${count}, ${picked.length} shown` toolSection.push(`- ${namespace} (${label})`) - for (const tool of group) if (picked.has(tool)) toolSection.push(catalogLine(tool)) + for (const tool of picked) toolSection.push(catalogLine(tool)) + const hiddenNames = group + .filter((tool) => !picked.includes(tool)) + .map((tool) => tool.path.slice(namespace.length + 1)) + const categoryLines = categories + .filter((category) => category.namespace === namespace) + .flatMap((category) => { + const matching = hiddenNames.filter((name) => category.pattern.test(name)).length + return matching === 0 ? [] : [` - ${category.name} (${matching}): ${category.description}`] + }) + if (categoryLines.length > 0) { + toolSection.push( + ` Not shown, by category (get exact signatures with tools.$codemode.search({ query, namespace: ${JSON.stringify(namespace)} })):`, + ...categoryLines, + ) + } } if (!complete) { toolSection.push("", "Search returns complete callable signatures:", `- ${searchDescription.signature}`) diff --git a/packages/codemode/test/codemode.test.ts b/packages/codemode/test/codemode.test.ts index e9ecab4b21b2..c1a78d2a4701 100644 --- a/packages/codemode/test/codemode.test.ts +++ b/packages/codemode/test/codemode.test.ts @@ -1200,3 +1200,102 @@ describe("CodeMode public contract", () => { expect(() => CodeMode.make({ tools: { $codemode: { lookup } } })).toThrow(/reserved for CodeMode discovery tools/) }) }) + +describe("CodeMode pinned catalog", () => { + const make = (description: string) => + Tool.make({ + description, + input: Schema.Struct({ q: Schema.String }), + output: Schema.String, + run: () => Effect.succeed("ok"), + }) + const tools = { + store: { + list_sites: make("List sites"), + shop_products_get: make("A much longer description that the length-ordered selection would place last"), + shop_products_search: make("Search products"), + shop_collections_get: make("Get a collection"), + get_account: make("Read the account"), + }, + web: { fetch: make("Fetch a page") }, + } + + test("inlines only pinned tools, in the caller's order, skipping absent paths", () => { + const instructions = CodeMode.make({ + tools, + discovery: { pinned: ["store.shop_products_get", "store.missing", "store.list_sites"] }, + }).instructions() + + expect(instructions).toContain( + "Available tools (PARTIAL - 2 of 6 shown; find the rest with tools.$codemode.search)", + ) + expect(instructions).toContain("- store (5 tools, 2 shown)") + expect(instructions.indexOf("tools.store.shop_products_get(")).toBeLessThan( + instructions.indexOf("tools.store.list_sites("), + ) + expect(instructions).not.toContain("tools.store.get_account(") + expect(instructions).toContain("- web (1 tool, none shown)") + expect(instructions).not.toContain("tools.web.fetch(") + }) + + test("caps pinned tools at pinnedLimit and the budget", () => { + const limited = CodeMode.make({ + tools, + discovery: { pinned: ["store.get_account", "store.list_sites", "web.fetch"], pinnedLimit: 2 }, + }).instructions() + expect(limited).toContain("PARTIAL - 2 of 6 shown") + expect(limited).not.toContain("tools.web.fetch(") + + const budgeted = CodeMode.make({ + tools, + discovery: { pinned: ["store.shop_products_get", "store.list_sites"], catalogBudget: 25 }, + }).instructions() + expect(budgeted).not.toContain("tools.store.shop_products_get(") + expect(budgeted).toContain("tools.store.list_sites(") + }) + + test("indexes categories of tools that are present but not inlined", () => { + const instructions = CodeMode.make({ + tools, + discovery: { + pinned: ["store.shop_products_get"], + categories: [ + { namespace: "store", name: "shop_*", match: "^shop_", description: "products, collections" }, + { namespace: "store", name: "figma_*", match: "^figma_", description: "files" }, + { namespace: "store", name: "other", match: "^(?!shop_)", description: "sites, account" }, + ], + }, + }).instructions() + + expect(instructions).toContain( + ' Not shown, by category (get exact signatures with tools.$codemode.search({ query, namespace: "store" })):\n - shop_* (2): products, collections\n - other (2): sites, account', + ) + expect(instructions).not.toContain("figma_*") + }) + + test("omits a category whose tools are all inlined", () => { + const instructions = CodeMode.make({ + tools, + discovery: { + pinned: ["web.fetch"], + categories: [{ namespace: "web", name: "fetch", match: "^fetch$", description: "pages" }], + }, + }).instructions() + expect(instructions).not.toContain("Not shown, by category") + }) + + test("rejects an invalid pinnedLimit or category pattern", () => { + expect(() => CodeMode.make({ tools, discovery: { pinned: [], pinnedLimit: -1 } })).toThrow(RangeError) + expect(() => + CodeMode.make({ + tools, + discovery: { categories: [{ namespace: "store", name: "x", match: "(", description: "" }] }, + }), + ).toThrow(SyntaxError) + }) + + test("matches the default catalog when discovery options are empty", () => { + expect(CodeMode.make({ tools }).instructions()).toBe(CodeMode.make({ tools, discovery: {} }).instructions()) + expect(CodeMode.make({ tools }).instructions()).toContain("COMPLETE list") + }) +}) diff --git a/packages/opencode/src/tool/code-mode.ts b/packages/opencode/src/tool/code-mode.ts index 931618c568d5..6741552dff98 100644 --- a/packages/opencode/src/tool/code-mode.ts +++ b/packages/opencode/src/tool/code-mode.ts @@ -55,12 +55,17 @@ function groupByServer(mcpTools: Record, servers: readonly return groups } -export function describeCatalog(mcpTools: Record, servers: readonly string[]): string { +export function describeCatalog( + mcpTools: Record, + servers: readonly string[], + discovery?: CodeMode.DiscoveryOptions, +): string { return CodeMode.make({ tools: toolTree( [...groupByServer(mcpTools, servers).values()].flat(), () => () => Effect.fail(toolError("Tool preview is not executable.")), ), + discovery, }).instructions() } diff --git a/packages/opencode/src/tool/registry.ts b/packages/opencode/src/tool/registry.ts index d4a313c604ef..75c663647946 100644 --- a/packages/opencode/src/tool/registry.ts +++ b/packages/opencode/src/tool/registry.ts @@ -54,6 +54,7 @@ import { ModelV2 } from "@opencode-ai/core/model" import { MCP } from "@/mcp" import { PermissionV1 } from "@opencode-ai/core/v1/permission" import { McpCatalog } from "@/mcp/catalog" +import type { CodeMode } from "@opencode-ai/codemode" export function webSearchEnabled(providerID: ProviderV2.ID, flags = { exa: false, parallel: false }) { return ( @@ -301,12 +302,23 @@ const layer = Layer.effect( const describeCodeMode = Effect.fn("ToolRegistry.describeCodeMode")(function* (input: { agent: Agent.Info permission?: PermissionV1.Ruleset + discovery?: CodeMode.DiscoveryOptions }) { if (!codeMode) return const ruleset = Permission.merge(input.agent.permission, input.permission ?? []) const tools = Permission.visibleTools(yield* mcp.tools(), ruleset) if (Object.keys(tools).length === 0) return - return codeMode.describeCatalog(tools, Object.keys(yield* mcp.clients()).map(McpCatalog.sanitize)) + const servers = Object.keys(yield* mcp.clients()).map(McpCatalog.sanitize) + const discovery = input.discovery + if (!discovery) return codeMode.describeCatalog(tools, servers) + // Plugin-supplied options are data from outside opencode; a bad pattern must not remove execute. + return yield* Effect.try({ + try: () => codeMode.describeCatalog(tools, servers, discovery), + catch: (error) => String(error), + }).pipe( + Effect.tapError((error) => Effect.logWarning("ignoring invalid execute discovery options", { error })), + Effect.orElseSucceed(() => codeMode.describeCatalog(tools, servers)), + ) }) const tools: Interface["tools"] = Effect.fn("ToolRegistry.tools")(function* (input) { @@ -331,12 +343,21 @@ const layer = Layer.effect( return yield* Effect.forEach( visible, Effect.fnUntraced(function* (tool: Tool.Def) { - const output = { + const output: { + description: string + parameters: Tool.Def["parameters"] + jsonSchema: Tool.Def["jsonSchema"] + discovery?: CodeMode.DiscoveryOptions + } = { description: tool.description, parameters: tool.parameters, jsonSchema: tool.jsonSchema, } yield* plugin.trigger("tool.definition", { toolID: tool.id }, output) + const executeDescription = + tool.id === "execute" && output.discovery + ? ((yield* describeCodeMode({ ...input, discovery: output.discovery })) ?? codeModeDescription) + : codeModeDescription const jsonSchema = output.parameters === tool.parameters || output.jsonSchema !== tool.jsonSchema ? output.jsonSchema @@ -346,7 +367,7 @@ const layer = Layer.effect( description: [ output.description, tool.id === TaskTool.id ? yield* describeTask(input.agent) : undefined, - tool.id === "execute" ? codeModeDescription : undefined, + tool.id === "execute" ? executeDescription : undefined, ] .filter(Boolean) .join("\n"), diff --git a/packages/opencode/test/tool/registry.test.ts b/packages/opencode/test/tool/registry.test.ts index 4b33e0f1148a..c69ed9ce9062 100644 --- a/packages/opencode/test/tool/registry.test.ts +++ b/packages/opencode/test/tool/registry.test.ts @@ -96,6 +96,76 @@ const withEmptyCodeMode = testEffect( ) const withBrokenPlugin = testEffect(LayerNode.compile(root, [...replacements, [Plugin.node, brokenPluginLayer]])) +// A plugin that steers execute's catalog through the tool.definition hook. +const discoveryPluginLayer = (discovery: unknown) => + Layer.succeed( + Plugin.Service, + Plugin.Service.of({ + init: () => Effect.void, + reload: () => Effect.void, + trigger: ((name: string, input: { toolID?: string }, output: Record) => + Effect.sync(() => { + if (name === "tool.definition" && input.toolID === "execute") { + output.description = `Guidance first.\n${output.description}` + output.discovery = discovery + } + return output + })) as unknown as Plugin.Interface["trigger"], + list: () => Effect.succeed([]), + }), + ) +const withCodeModeDiscovery = (discovery: unknown) => + testEffect( + LayerNode.compile(root, [ + [Config.node, configLayer], + [RuntimeFlags.node, RuntimeFlags.layer({ experimentalCodeMode: true })], + [ + MCP.node, + Layer.mock(MCP.Service, { + tools: () => + Effect.succeed({ + weather_current: { + def: { + name: "current", + description: "current weather", + inputSchema: { type: "object", properties: { city: { type: "string" } }, required: ["city"] }, + } as MCPToolDef, + client: {} as MCP.McpTool["client"], + }, + weather_forecast_hourly: { + def: { + name: "forecast_hourly", + description: "hourly forecast for the next several days, with a long description", + inputSchema: { type: "object", properties: { city: { type: "string" } }, required: ["city"] }, + } as MCPToolDef, + client: {} as MCP.McpTool["client"], + }, + }), + clients: () => Effect.succeed({ weather: {} as any }), + }), + ], + [Plugin.node, discoveryPluginLayer(discovery)], + ]), + ) +const withPinnedCatalog = withCodeModeDiscovery({ + pinned: ["weather.forecast_hourly"], + categories: [{ namespace: "weather", name: "current*", match: "^current", description: "conditions now" }], +}) +const withInvalidDiscovery = withCodeModeDiscovery({ + pinned: ["weather.forecast_hourly"], + categories: [{ namespace: "weather", name: "x", match: "(", description: "" }], +}) +const executeDescription = Effect.gen(function* () { + const registry = yield* ToolRegistry.Service + const agents = yield* Agent.Service + const tools = yield* registry.tools({ + providerID: ProviderV2.ID.opencode, + modelID: ModelV2.ID.make("test"), + agent: yield* agents.defaultInfo(), + }) + return tools.find((tool) => tool.id === "execute")?.description ?? "" +}) + afterEach(async () => { await disposeAllInstances() }) @@ -137,6 +207,27 @@ describe("tool.registry", () => { }), ) + withPinnedCatalog.instance("lets a tool.definition hook rank execute's catalog", () => + Effect.gen(function* () { + const description = yield* executeDescription + + expect(description.startsWith("Guidance first.\n")).toBe(true) + expect(description).toContain("PARTIAL - 1 of 2 shown") + expect(description).toContain("tools.weather.forecast_hourly(input: {") + expect(description).not.toContain("tools.weather.current(input:") + expect(description).toContain(" - current* (1): conditions now") + }), + ) + + withInvalidDiscovery.instance("falls back to the default catalog when hook discovery options are invalid", () => + Effect.gen(function* () { + const description = yield* executeDescription + + expect(description).toContain("COMPLETE list") + expect(description).toContain("tools.weather.current(input: {") + }), + ) + withEmptyCodeMode.instance("does not expose execute when code mode has no visible tools", () => Effect.gen(function* () { const registry = yield* ToolRegistry.Service diff --git a/packages/plugin/src/index.ts b/packages/plugin/src/index.ts index edfa0139dfca..ca6121edac9d 100644 --- a/packages/plugin/src/index.ts +++ b/packages/plugin/src/index.ts @@ -331,5 +331,23 @@ export interface Hooks { /** * Modify tool definitions (description and parameters) sent to LLM */ - "tool.definition"?: (input: { toolID: string }, output: { description: string; parameters: any }) => Promise + "tool.definition"?: ( + input: { toolID: string }, + output: { + description: string + parameters: any + /** + * `execute` only: how its MCP tool catalog is inlined. `pinned` tool paths + * (`namespace.tool`, most important first; order is kept within each namespace) replace + * the default selection of cheapest signatures round-robin across namespaces; + * `categories` index tools that are present but not inlined. + */ + discovery?: { + catalogBudget?: number + pinned?: string[] + pinnedLimit?: number + categories?: { namespace: string; name: string; match: string; description: string }[] + } + }, + ) => Promise }