refactor(opencode): activate MCP search by catalog size
This commit is contained in:
parent
94b7d450c5
commit
bc189de080
6 changed files with 125 additions and 78 deletions
|
|
@ -48,8 +48,6 @@ export const Flag = {
|
||||||
|
|
||||||
OPENCODE_WORKSPACE_ID: process.env["OPENCODE_WORKSPACE_ID"],
|
OPENCODE_WORKSPACE_ID: process.env["OPENCODE_WORKSPACE_ID"],
|
||||||
OPENCODE_EXPERIMENTAL_WORKSPACES: enabledByExperimental("OPENCODE_EXPERIMENTAL_WORKSPACES"),
|
OPENCODE_EXPERIMENTAL_WORKSPACES: enabledByExperimental("OPENCODE_EXPERIMENTAL_WORKSPACES"),
|
||||||
OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH: enabledByExperimental("OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH"),
|
|
||||||
|
|
||||||
// Evaluated at access time (not module load) because tests, the CLI, and
|
// Evaluated at access time (not module load) because tests, the CLI, and
|
||||||
// external tooling set these env vars at runtime.
|
// external tooling set these env vars at runtime.
|
||||||
get OPENCODE_DISABLE_PROJECT_CONFIG() {
|
get OPENCODE_DISABLE_PROJECT_CONFIG() {
|
||||||
|
|
|
||||||
|
|
@ -1,10 +1,13 @@
|
||||||
import { jsonSchema, tool, type JSONSchema7, type Tool, type ToolExecutionOptions } from "ai"
|
import { jsonSchema, tool, type JSONSchema7, type Tool, type ToolExecutionOptions } from "ai"
|
||||||
import fuzzysort from "fuzzysort"
|
import fuzzysort from "fuzzysort"
|
||||||
|
import { Token } from "@opencode-ai/core/util/token"
|
||||||
|
|
||||||
// Match Hermes defaults. OpenClaw independently uses the same maximum.
|
// Match Hermes defaults. OpenClaw independently uses the same maximum.
|
||||||
const DEFAULT_SEARCH_LIMIT = 5
|
const DEFAULT_SEARCH_LIMIT = 5
|
||||||
const MAX_SEARCH_LIMIT = 20
|
const MAX_SEARCH_LIMIT = 20
|
||||||
const MAX_SEARCH_DESCRIPTION = 400
|
const MAX_SEARCH_DESCRIPTION = 400
|
||||||
|
const SEARCH_THRESHOLD_TOKENS = 15_000
|
||||||
|
const controls = new WeakSet<Tool>()
|
||||||
|
|
||||||
type Entry = {
|
type Entry = {
|
||||||
id: string
|
id: string
|
||||||
|
|
@ -14,6 +17,18 @@ type Entry = {
|
||||||
tool: Tool
|
tool: Tool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export function shouldUse(tools: Record<string, Tool>, schemas: Record<string, JSONSchema7>) {
|
||||||
|
const catalog = Object.entries(tools)
|
||||||
|
.toSorted(([a], [b]) => a.localeCompare(b))
|
||||||
|
.map(([name, item]) => JSON.stringify({ name, description: item.description, inputSchema: schemas[name] }))
|
||||||
|
.join("\n")
|
||||||
|
return Token.estimate(catalog) > SEARCH_THRESHOLD_TOKENS
|
||||||
|
}
|
||||||
|
|
||||||
|
export function isControl(item: Tool) {
|
||||||
|
return controls.has(item)
|
||||||
|
}
|
||||||
|
|
||||||
export function create(input: {
|
export function create(input: {
|
||||||
tools: Record<string, Tool>
|
tools: Record<string, Tool>
|
||||||
schemas: Record<string, JSONSchema7>
|
schemas: Record<string, JSONSchema7>
|
||||||
|
|
@ -39,7 +54,7 @@ export function create(input: {
|
||||||
|
|
||||||
if (entries.size === 0) return {}
|
if (entries.size === 0) return {}
|
||||||
|
|
||||||
return {
|
const result = {
|
||||||
mcp_search: tool({
|
mcp_search: tool({
|
||||||
description:
|
description:
|
||||||
"Search connected MCP tools by capability. Returns matching tool IDs and short descriptions. Use mcp_describe to inspect a tool before calling it.",
|
"Search connected MCP tools by capability. Returns matching tool IDs and short descriptions. Use mcp_describe to inspect a tool before calling it.",
|
||||||
|
|
@ -128,6 +143,8 @@ export function create(input: {
|
||||||
},
|
},
|
||||||
}),
|
}),
|
||||||
}
|
}
|
||||||
|
for (const item of Object.values(result)) controls.add(item)
|
||||||
|
return result
|
||||||
}
|
}
|
||||||
|
|
||||||
function resolve(entries: Map<string, Entry>, id: string) {
|
function resolve(entries: Map<string, Entry>, id: string) {
|
||||||
|
|
|
||||||
|
|
@ -15,7 +15,7 @@ import { jsonSchema, tool as aiTool, type ModelMessage, type Tool } from "ai"
|
||||||
import type { Plugin } from "@/plugin"
|
import type { Plugin } from "@/plugin"
|
||||||
import { mergeDeep } from "remeda"
|
import { mergeDeep } from "remeda"
|
||||||
import { Wildcard } from "@opencode-ai/core/util/wildcard"
|
import { Wildcard } from "@opencode-ai/core/util/wildcard"
|
||||||
import { Flag } from "@opencode-ai/core/flag/flag"
|
import { McpToolSearch } from "@/mcp/tool-search"
|
||||||
|
|
||||||
const USER_AGENT = `opencode/${InstallationVersion}`
|
const USER_AGENT = `opencode/${InstallationVersion}`
|
||||||
|
|
||||||
|
|
@ -200,11 +200,10 @@ export const prepare = Effect.fn("LLMRequestPrep.prepare")(function* (input: Pre
|
||||||
export function resolveTools(input: Pick<PrepareInput, "tools" | "agent" | "permission" | "user">) {
|
export function resolveTools(input: Pick<PrepareInput, "tools" | "agent" | "permission" | "user">) {
|
||||||
const ruleset = Permission.merge(input.agent.permission, input.permission ?? [])
|
const ruleset = Permission.merge(input.agent.permission, input.permission ?? [])
|
||||||
const disabled = Permission.disabled(Object.keys(input.tools), ruleset)
|
const disabled = Permission.disabled(Object.keys(input.tools), ruleset)
|
||||||
const controls = new Set(["mcp_search", "mcp_describe", "mcp_call"])
|
return Record.filter(input.tools, (item, key) => {
|
||||||
return Record.filter(input.tools, (_, key) => {
|
|
||||||
if (input.user.tools?.[key] === false) return false
|
if (input.user.tools?.[key] === false) return false
|
||||||
if (!disabled.has(key)) return true
|
if (!disabled.has(key)) return true
|
||||||
if (!Flag.OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH || !controls.has(key)) return false
|
if (!McpToolSearch.isControl(item)) return false
|
||||||
return ruleset.findLast((rule) => Wildcard.match(key, rule.permission))?.permission === "*"
|
return ruleset.findLast((rule) => Wildcard.match(key, rule.permission))?.permission === "*"
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -21,7 +21,7 @@ import { EffectBridge } from "@/effect/bridge"
|
||||||
import { ProviderV2 } from "@opencode-ai/core/provider"
|
import { ProviderV2 } from "@opencode-ai/core/provider"
|
||||||
import { ModelV2 } from "@opencode-ai/core/model"
|
import { ModelV2 } from "@opencode-ai/core/model"
|
||||||
import { isRecord } from "@/util/record"
|
import { isRecord } from "@/util/record"
|
||||||
import { Flag } from "@opencode-ai/core/flag/flag"
|
import { McpToolSearch } from "@/mcp/tool-search"
|
||||||
|
|
||||||
const MCP_RESOURCE_TOOLS = {
|
const MCP_RESOURCE_TOOLS = {
|
||||||
list: "list_mcp_resources",
|
list: "list_mcp_resources",
|
||||||
|
|
@ -384,6 +384,7 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: {
|
||||||
|
|
||||||
const mcpTools: Record<string, AITool> = {}
|
const mcpTools: Record<string, AITool> = {}
|
||||||
const mcpSchemas: Record<string, JSONSchema7> = {}
|
const mcpSchemas: Record<string, JSONSchema7> = {}
|
||||||
|
const mcpProviderSchemas: Record<string, JSONSchema7> = {}
|
||||||
for (const [key, item] of Object.entries(yield* mcp.tools())) {
|
for (const [key, item] of Object.entries(yield* mcp.tools())) {
|
||||||
const execute = item.execute
|
const execute = item.execute
|
||||||
if (!execute) continue
|
if (!execute) continue
|
||||||
|
|
@ -391,6 +392,7 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: {
|
||||||
const schema = yield* Effect.promise(() => Promise.resolve(asSchema(item.inputSchema).jsonSchema))
|
const schema = yield* Effect.promise(() => Promise.resolve(asSchema(item.inputSchema).jsonSchema))
|
||||||
mcpSchemas[key] = { ...schema, properties: schema.properties ?? {} }
|
mcpSchemas[key] = { ...schema, properties: schema.properties ?? {} }
|
||||||
const transformed = ProviderTransform.schema(input.model, { ...schema, properties: schema.properties ?? {} })
|
const transformed = ProviderTransform.schema(input.model, { ...schema, properties: schema.properties ?? {} })
|
||||||
|
mcpProviderSchemas[key] = transformed
|
||||||
item.inputSchema = jsonSchema(transformed)
|
item.inputSchema = jsonSchema(transformed)
|
||||||
item.execute = (args, opts) =>
|
item.execute = (args, opts) =>
|
||||||
run.promise(
|
run.promise(
|
||||||
|
|
@ -486,12 +488,6 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: {
|
||||||
mcpTools[key] = item
|
mcpTools[key] = item
|
||||||
}
|
}
|
||||||
|
|
||||||
if (!Flag.OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH) {
|
|
||||||
Object.assign(tools, mcpTools)
|
|
||||||
return tools
|
|
||||||
}
|
|
||||||
|
|
||||||
const { McpToolSearch } = yield* Effect.promise(() => import("@/mcp/tool-search"))
|
|
||||||
const disabled = Permission.disabled(
|
const disabled = Permission.disabled(
|
||||||
Object.keys(mcpTools),
|
Object.keys(mcpTools),
|
||||||
Permission.merge(input.agent.permission, input.session.permission ?? []),
|
Permission.merge(input.agent.permission, input.session.permission ?? []),
|
||||||
|
|
@ -501,6 +497,12 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: {
|
||||||
const searchable = Object.fromEntries(
|
const searchable = Object.fromEntries(
|
||||||
Object.entries(mcpTools).filter(([key]) => overrides?.[key] !== false && !disabled.has(key)),
|
Object.entries(mcpTools).filter(([key]) => overrides?.[key] !== false && !disabled.has(key)),
|
||||||
)
|
)
|
||||||
|
const providerSchemas = Object.fromEntries(Object.keys(searchable).map((key) => [key, mcpProviderSchemas[key]]))
|
||||||
|
if (!McpToolSearch.shouldUse(searchable, providerSchemas)) {
|
||||||
|
Object.assign(tools, mcpTools)
|
||||||
|
return tools
|
||||||
|
}
|
||||||
|
|
||||||
const schemas = Object.fromEntries(Object.keys(searchable).map((key) => [key, mcpSchemas[key]]))
|
const schemas = Object.fromEntries(Object.keys(searchable).map((key) => [key, mcpSchemas[key]]))
|
||||||
const controls = McpToolSearch.create({
|
const controls = McpToolSearch.create({
|
||||||
tools: searchable,
|
tools: searchable,
|
||||||
|
|
|
||||||
|
|
@ -44,6 +44,22 @@ async function catalog() {
|
||||||
}
|
}
|
||||||
|
|
||||||
describe("MCP tool search", () => {
|
describe("MCP tool search", () => {
|
||||||
|
test("uses direct tools below the token threshold", async () => {
|
||||||
|
const tools = {
|
||||||
|
github_create_issue: target("github_create_issue", "Create an issue"),
|
||||||
|
}
|
||||||
|
const schemas = { github_create_issue: { type: "object", properties: {} } } satisfies Record<string, JSONSchema7>
|
||||||
|
expect(McpToolSearch.shouldUse(tools, schemas)).toBe(false)
|
||||||
|
})
|
||||||
|
|
||||||
|
test("uses search above the token threshold", async () => {
|
||||||
|
const tools = {
|
||||||
|
large_catalog: target("large_catalog", "x".repeat(60_000)),
|
||||||
|
}
|
||||||
|
const schemas = { large_catalog: { type: "object", properties: {} } } satisfies Record<string, JSONSchema7>
|
||||||
|
expect(McpToolSearch.shouldUse(tools, schemas)).toBe(true)
|
||||||
|
})
|
||||||
|
|
||||||
test("exposes only the three stable control tools", async () => {
|
test("exposes only the three stable control tools", async () => {
|
||||||
expect(Object.keys(await catalog())).toEqual(["mcp_search", "mcp_describe", "mcp_call"])
|
expect(Object.keys(await catalog())).toEqual(["mcp_search", "mcp_describe", "mcp_call"])
|
||||||
})
|
})
|
||||||
|
|
|
||||||
|
|
@ -28,7 +28,7 @@ import { Session as SessionNs } from "@/session/session"
|
||||||
import { LLMRequestPrep } from "@/session/llm/request"
|
import { LLMRequestPrep } from "@/session/llm/request"
|
||||||
import { ProviderV2 } from "@opencode-ai/core/provider"
|
import { ProviderV2 } from "@opencode-ai/core/provider"
|
||||||
import { ModelV2 } from "@opencode-ai/core/model"
|
import { ModelV2 } from "@opencode-ai/core/model"
|
||||||
import { Flag } from "@opencode-ai/core/flag/flag"
|
import { McpToolSearch } from "@/mcp/tool-search"
|
||||||
|
|
||||||
type ConfigModel = NonNullable<NonNullable<ConfigV1.Info["provider"]>[string]["models"]>[string]
|
type ConfigModel = NonNullable<NonNullable<ConfigV1.Info["provider"]>[string]["models"]>[string]
|
||||||
|
|
||||||
|
|
@ -177,73 +177,88 @@ describe("session.llm.hasToolCalls", () => {
|
||||||
|
|
||||||
describe("session.llm.resolveTools", () => {
|
describe("session.llm.resolveTools", () => {
|
||||||
test("keeps MCP search controls under a generic deny rule", () => {
|
test("keeps MCP search controls under a generic deny rule", () => {
|
||||||
const previous = Flag.OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH
|
const controls = McpToolSearch.create({
|
||||||
Flag.OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH = true
|
tools: { target: tool({ inputSchema: z.object({}), execute: async () => ({ output: "" }) }) },
|
||||||
try {
|
schemas: { target: { type: "object", properties: {} } },
|
||||||
const agent = {
|
transformSchema: (schema) => schema,
|
||||||
name: "test",
|
})
|
||||||
mode: "primary",
|
const agent = {
|
||||||
options: {},
|
name: "test",
|
||||||
permission: [
|
mode: "primary",
|
||||||
{ permission: "*", pattern: "*", action: "deny" },
|
options: {},
|
||||||
{ permission: "github_create_issue", pattern: "*", action: "allow" },
|
permission: [
|
||||||
],
|
{ permission: "*", pattern: "*", action: "deny" },
|
||||||
} satisfies Agent.Info
|
{ permission: "github_create_issue", pattern: "*", action: "allow" },
|
||||||
const user = {
|
],
|
||||||
id: MessageID.make("msg_mcp_tool_search"),
|
} satisfies Agent.Info
|
||||||
sessionID: SessionID.make("ses_mcp_tool_search"),
|
const user = {
|
||||||
role: "user",
|
id: MessageID.make("msg_mcp_tool_search"),
|
||||||
time: { created: Date.now() },
|
sessionID: SessionID.make("ses_mcp_tool_search"),
|
||||||
agent: agent.name,
|
role: "user",
|
||||||
model: { providerID: ProviderV2.ID.make("test"), modelID: ModelV2.ID.make("test") },
|
time: { created: Date.now() },
|
||||||
} satisfies SessionV1.User
|
agent: agent.name,
|
||||||
const control = tool({ inputSchema: z.object({}), execute: async () => ({ output: "" }) })
|
model: { providerID: ProviderV2.ID.make("test"), modelID: ModelV2.ID.make("test") },
|
||||||
expect(
|
} satisfies SessionV1.User
|
||||||
Object.keys(
|
expect(
|
||||||
LLMRequestPrep.resolveTools({
|
Object.keys(
|
||||||
tools: { mcp_search: control, mcp_describe: control, mcp_call: control },
|
|
||||||
agent,
|
|
||||||
user,
|
|
||||||
}),
|
|
||||||
),
|
|
||||||
).toEqual(["mcp_search", "mcp_describe", "mcp_call"])
|
|
||||||
} finally {
|
|
||||||
Flag.OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH = previous
|
|
||||||
}
|
|
||||||
})
|
|
||||||
|
|
||||||
test("honors explicit MCP search control denies", () => {
|
|
||||||
const previous = Flag.OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH
|
|
||||||
Flag.OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH = true
|
|
||||||
try {
|
|
||||||
const agent = {
|
|
||||||
name: "test",
|
|
||||||
mode: "primary",
|
|
||||||
options: {},
|
|
||||||
permission: [
|
|
||||||
{ permission: "*", pattern: "*", action: "deny" },
|
|
||||||
{ permission: "mcp_*", pattern: "*", action: "deny" },
|
|
||||||
],
|
|
||||||
} satisfies Agent.Info
|
|
||||||
const user = {
|
|
||||||
id: MessageID.make("msg_mcp_tool_search_deny"),
|
|
||||||
sessionID: SessionID.make("ses_mcp_tool_search_deny"),
|
|
||||||
role: "user",
|
|
||||||
time: { created: Date.now() },
|
|
||||||
agent: agent.name,
|
|
||||||
model: { providerID: ProviderV2.ID.make("test"), modelID: ModelV2.ID.make("test") },
|
|
||||||
} satisfies SessionV1.User
|
|
||||||
const control = tool({ inputSchema: z.object({}), execute: async () => ({ output: "" }) })
|
|
||||||
expect(
|
|
||||||
LLMRequestPrep.resolveTools({
|
LLMRequestPrep.resolveTools({
|
||||||
tools: { mcp_search: control, mcp_describe: control, mcp_call: control },
|
tools: controls,
|
||||||
agent,
|
agent,
|
||||||
user,
|
user,
|
||||||
}),
|
}),
|
||||||
).toEqual({})
|
),
|
||||||
} finally {
|
).toEqual(["mcp_search", "mcp_describe", "mcp_call"])
|
||||||
Flag.OPENCODE_EXPERIMENTAL_MCP_TOOL_SEARCH = previous
|
})
|
||||||
}
|
|
||||||
|
test("honors explicit MCP search control denies", () => {
|
||||||
|
const controls = McpToolSearch.create({
|
||||||
|
tools: { target: tool({ inputSchema: z.object({}), execute: async () => ({ output: "" }) }) },
|
||||||
|
schemas: { target: { type: "object", properties: {} } },
|
||||||
|
transformSchema: (schema) => schema,
|
||||||
|
})
|
||||||
|
const agent = {
|
||||||
|
name: "test",
|
||||||
|
mode: "primary",
|
||||||
|
options: {},
|
||||||
|
permission: [
|
||||||
|
{ permission: "*", pattern: "*", action: "deny" },
|
||||||
|
{ permission: "mcp_*", pattern: "*", action: "deny" },
|
||||||
|
],
|
||||||
|
} satisfies Agent.Info
|
||||||
|
const user = {
|
||||||
|
id: MessageID.make("msg_mcp_tool_search_deny"),
|
||||||
|
sessionID: SessionID.make("ses_mcp_tool_search_deny"),
|
||||||
|
role: "user",
|
||||||
|
time: { created: Date.now() },
|
||||||
|
agent: agent.name,
|
||||||
|
model: { providerID: ProviderV2.ID.make("test"), modelID: ModelV2.ID.make("test") },
|
||||||
|
} satisfies SessionV1.User
|
||||||
|
expect(
|
||||||
|
LLMRequestPrep.resolveTools({
|
||||||
|
tools: controls,
|
||||||
|
agent,
|
||||||
|
user,
|
||||||
|
}),
|
||||||
|
).toEqual({})
|
||||||
|
})
|
||||||
|
|
||||||
|
test("does not exempt plugin tools that use MCP control names", () => {
|
||||||
|
const agent = {
|
||||||
|
name: "test",
|
||||||
|
mode: "primary",
|
||||||
|
options: {},
|
||||||
|
permission: [{ permission: "*", pattern: "*", action: "deny" }],
|
||||||
|
} satisfies Agent.Info
|
||||||
|
const user = {
|
||||||
|
id: MessageID.make("msg_mcp_plugin_tool"),
|
||||||
|
sessionID: SessionID.make("ses_mcp_plugin_tool"),
|
||||||
|
role: "user",
|
||||||
|
time: { created: Date.now() },
|
||||||
|
agent: agent.name,
|
||||||
|
model: { providerID: ProviderV2.ID.make("test"), modelID: ModelV2.ID.make("test") },
|
||||||
|
} satisfies SessionV1.User
|
||||||
|
const pluginTool = tool({ inputSchema: z.object({}), execute: async () => ({ output: "" }) })
|
||||||
|
expect(LLMRequestPrep.resolveTools({ tools: { mcp_call: pluginTool }, agent, user })).toEqual({})
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue