refactor(opencode): reuse deferred MCP resolution

This commit is contained in:
Aiden Cline 2026-06-29 07:59:55 -05:00
commit c2d619741b
4 changed files with 30 additions and 46 deletions

View file

@ -61,6 +61,11 @@ interface DeferredToolDescriptor {
searchText: string searchText: string
} }
interface Output {
tools: Record<string, Tool>
deferredSystemPrompt?: string
}
const SearchDeferredToolsParameters = Schema.Struct({ const SearchDeferredToolsParameters = Schema.Struct({
query: Schema.String.annotate({ query: Schema.String.annotate({
description: description:
@ -229,9 +234,8 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In
return output return output
}) })
const hasMcpResourceServer = Object.values(yield* mcp.clients()).some( const clients = yield* mcp.clients()
(client) => !!client.getServerCapabilities()?.resources, const hasMcpResourceServer = Object.values(clients).some((client) => !!client.getServerCapabilities()?.resources)
)
if (hasMcpResourceServer) addResourceTools(tools, { input, run, mcp, plugin, truncate, context }) if (hasMcpResourceServer) addResourceTools(tools, { input, run, mcp, plugin, truncate, context })
const mcpTools = yield* mcp.tools() const mcpTools = yield* mcp.tools()
@ -261,7 +265,10 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In
allowedMcpTools, allowedMcpTools,
executeMcpTool, executeMcpTool,
}) })
return tools return {
tools,
deferredSystemPrompt: deferredSystemPrompt(allowedMcpTools, clients),
} satisfies Output
} }
for (const [key, item] of Object.entries(mcpTools)) { for (const [key, item] of Object.entries(mcpTools)) {
@ -275,35 +282,11 @@ export const resolve = Effect.fn("SessionMcpTools.resolve")(function* (input: In
tools[key] = item tools[key] = item
} }
return tools return { tools } satisfies Output
}) })
export const systemPrompt = Effect.fn("SessionMcpTools.systemPrompt")(function* (input: { function deferredSystemPrompt(allowedMcpTools: Record<string, Tool>, clients: Record<string, unknown>) {
agent: Agent.Info const servers = deferredServerSummaries(Object.keys(allowedMcpTools), Object.keys(clients))
session: Session.Info
messages: SessionV1.WithParts[]
}) {
const mcp = yield* MCP.Service
const flags = yield* RuntimeFlags.Service
if (!flags.experimentalToolSearch) return undefined
const mcpTools = yield* mcp.tools()
const mcpDisabled = Permission.disabled(
Object.keys(mcpTools),
Permission.merge(input.agent.permission, input.session.permission ?? []),
)
const userTools = currentUserToolOverrides(input.messages)
const allowedTools = Object.fromEntries(
Object.entries(mcpTools).filter(([key]) => userTools?.[key] !== false && !mcpDisabled.has(key)),
)
if (Object.keys(allowedTools).length === 0) return undefined
const deferredDescriptors = yield* deferredToolDescriptors(allowedTools)
if (deferredToolSchemaTokens(deferredDescriptors) < MIN_DEFERRED_MCP_SCHEMA_TOKENS) {
return undefined
}
const servers = deferredServerSummaries(Object.keys(allowedTools), Object.keys(yield* mcp.clients()))
return [ return [
DEFERRED_TOOL_SYSTEM_PROMPT, DEFERRED_TOOL_SYSTEM_PROMPT,
servers.length === 0 servers.length === 0
@ -315,7 +298,7 @@ export const systemPrompt = Effect.fn("SessionMcpTools.systemPrompt")(function*
] ]
.filter((part): part is string => part !== undefined) .filter((part): part is string => part !== undefined)
.join("\n\n") .join("\n\n")
}) }
function addResourceTools( function addResourceTools(
tools: Record<string, Tool>, tools: Record<string, Tool>,

View file

@ -1223,7 +1223,7 @@ export const layer = Layer.effect(
const bypassAgentCheck = lastUserMsg?.parts.some((p) => p.type === "agent") ?? false const bypassAgentCheck = lastUserMsg?.parts.some((p) => p.type === "agent") ?? false
const promptOps = yield* ops() const promptOps = yield* ops()
const tools = yield* SessionTools.resolve({ const resolvedTools = yield* SessionTools.resolve({
agent, agent,
session, session,
model, model,
@ -1240,6 +1240,7 @@ export const layer = Layer.effect(
Effect.provideService(Truncate.Service, truncate), Effect.provideService(Truncate.Service, truncate),
Effect.provideService(RuntimeFlags.Service, flags), Effect.provideService(RuntimeFlags.Service, flags),
) )
const tools = resolvedTools.tools
if (lastUser.format?.type === "json_schema") { if (lastUser.format?.type === "json_schema") {
tools["StructuredOutput"] = createStructuredOutputTool({ tools["StructuredOutput"] = createStructuredOutputTool({
@ -1260,12 +1261,7 @@ export const layer = Layer.effect(
sys.environment(model), sys.environment(model),
instruction.system().pipe(Effect.orDie), instruction.system().pipe(Effect.orDie),
sys.mcp(agent, session.permission), sys.mcp(agent, session.permission),
tools.search_deferred_tools && tools.call_deferred_tool Effect.succeed(resolvedTools.deferredSystemPrompt),
? SessionTools.deferredSystemPrompt({ agent, session, messages: msgs }).pipe(
Effect.provideService(MCP.Service, mcp),
Effect.provideService(RuntimeFlags.Service, flags),
)
: Effect.succeed(undefined),
MessageV2.toModelMessagesEffect(msgs, model), MessageV2.toModelMessagesEffect(msgs, model),
]) ])
const system = [ const system = [

View file

@ -17,9 +17,6 @@ import { SessionProcessor } from "./processor"
import { Session } from "./session" import { Session } from "./session"
import { PartID } from "./schema" import { PartID } from "./schema"
export const DEFERRED_TOOL_SYSTEM_PROMPT = SessionMcpTools.DEFERRED_TOOL_SYSTEM_PROMPT
export const deferredSystemPrompt = SessionMcpTools.systemPrompt
export const resolve = Effect.fn("SessionTools.resolve")(function* (input: { export const resolve = Effect.fn("SessionTools.resolve")(function* (input: {
agent: Agent.Info agent: Agent.Info
model: Provider.Model model: Provider.Model
@ -112,8 +109,12 @@ export const resolve = Effect.fn("SessionTools.resolve")(function* (input: {
}) })
} }
Object.assign(tools, yield* SessionMcpTools.resolve(input)) const mcp = yield* SessionMcpTools.resolve(input)
return tools Object.assign(tools, mcp.tools)
return {
tools,
deferredSystemPrompt: mcp.deferredSystemPrompt,
}
}) })
export * as SessionTools from "./tools" export * as SessionTools from "./tools"

View file

@ -132,7 +132,7 @@ const belowThresholdIt = makeIt({
queryDescription: "Natural language analytics query", queryDescription: "Natural language analytics query",
}) })
function resolveTools(input: { messages?: SessionV1.WithParts[]; canDeferMcpTools?: boolean } = {}) { function resolveToolResult(input: { messages?: SessionV1.WithParts[]; canDeferMcpTools?: boolean } = {}) {
return SessionTools.resolve({ return SessionTools.resolve({
agent, agent,
model, model,
@ -145,6 +145,10 @@ function resolveTools(input: { messages?: SessionV1.WithParts[]; canDeferMcpTool
}) })
} }
function resolveTools(input: { messages?: SessionV1.WithParts[]; canDeferMcpTools?: boolean } = {}) {
return resolveToolResult(input).pipe(Effect.map((result) => result.tools))
}
describe("session.tools", () => { describe("session.tools", () => {
deferredIt.instance("defers MCP tools behind fixed search and call tools", () => deferredIt.instance("defers MCP tools behind fixed search and call tools", () =>
Effect.gen(function* () { Effect.gen(function* () {
@ -181,7 +185,7 @@ describe("session.tools", () => {
deferredIt.instance("lists deferred MCP servers in the system prompt", () => deferredIt.instance("lists deferred MCP servers in the system prompt", () =>
Effect.gen(function* () { Effect.gen(function* () {
const prompt = yield* SessionTools.deferredSystemPrompt({ agent, session, messages: [] }) const prompt = (yield* resolveToolResult()).deferredSystemPrompt
expect(prompt).toContain("Deferred MCP servers available through `search_deferred_tools`:") expect(prompt).toContain("Deferred MCP servers available through `search_deferred_tools`:")
expect(prompt).toContain("- posthog: 2 tools") expect(prompt).toContain("- posthog: 2 tools")