fix(core): share one tool snapshot per request (#38596)

This commit is contained in:
Kit Langton 2026-07-23 22:54:22 -04:00 committed by GitHub
parent c228fc4886
commit 7456598cde
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 142 additions and 124 deletions

View file

@ -1,45 +1,24 @@
export * as CodeModeInstructions from "./instructions"
import { Context, Effect, Layer, Schema } from "effect"
import { makeLocationNode } from "@opencode-ai/util/effect/app-node"
import { AgentV2 } from "../agent"
import { CodeMode } from "../codemode"
import { Effect, Schema } from "effect"
import { Instructions } from "../instructions/index"
export interface Interface {
readonly load: (agent: AgentV2.Selection) => Effect.Effect<Instructions.Instructions>
const key = Instructions.Key.make("core/codemode")
const codec = Schema.toCodecJson(Schema.String)
const render = {
initial: (current: string) => current,
changed: (_previous: string, current: string) =>
[
"The Code Mode tool catalog has changed. This catalog supersedes the previous Code Mode tool catalog.",
current,
].join("\n\n"),
removed: () => "Code Mode tools are no longer available. Do not use any previously listed Code Mode tools.",
}
export class Service extends Context.Service<Service, Interface>()("@opencode/v2/CodeModeInstructions") {}
const layer = Layer.effect(
Service,
Effect.gen(function* () {
const codeMode = yield* CodeMode.Service
return Service.of({
load: Effect.fn("CodeModeInstructions.load")(function* (selection) {
const instructions = selection.info
? (yield* codeMode.materialize(selection.info.permissions)).instructions
: undefined
return Instructions.make({
key: Instructions.Key.make("core/codemode"),
codec: Schema.toCodecJson(Schema.String),
read: Effect.succeed(instructions ?? Instructions.removed),
render: {
initial: (current) => current,
changed: (_previous, current) =>
[
"The Code Mode tool catalog has changed. This catalog supersedes the previous Code Mode tool catalog.",
current,
].join("\n\n"),
removed: () =>
"Code Mode tools are no longer available. Do not use any previously listed Code Mode tools.",
},
})
}),
})
}),
)
export const node = makeLocationNode({ service: Service, layer, deps: [CodeMode.node] })
export const make = (content?: string): Instructions.Instructions =>
Instructions.make({
key,
codec,
read: Effect.succeed(content ?? Instructions.removed),
render,
})

View file

@ -3,7 +3,6 @@ import { AgentV2 } from "./agent"
import { AISDK } from "./aisdk"
import { Catalog } from "./catalog"
import { CodeMode } from "./codemode"
import { CodeModeInstructions } from "./codemode/instructions"
import { CommandV2 } from "./command"
import { Config } from "./config"
import { LayerNode } from "@opencode-ai/util/effect/layer-node"
@ -82,7 +81,6 @@ const locationServiceNodes = [
ToolRegistry.toolsNode,
Image.node,
SkillInstructions.node,
CodeModeInstructions.node,
ReferenceInstructions.node,
InstructionEntry.node,
Form.node,

View file

@ -2,6 +2,7 @@ export * as SessionContext from "./context"
import { Context, Effect, Layer } from "effect"
import { AgentV2 } from "../agent"
import { CodeModeInstructions } from "../codemode/instructions"
import { Database } from "../database/database"
import { makeLocationNode } from "@opencode-ai/util/effect/app-node"
import { InstructionDiscovery } from "../instruction-discovery"
@ -12,7 +13,7 @@ import { McpInstructions } from "../mcp/instructions"
import { PluginSupervisor } from "../plugin/supervisor"
import { ReferenceInstructions } from "../reference/instructions"
import { SkillInstructions } from "../skill/instructions"
import { CodeModeInstructions } from "../codemode/instructions"
import { ToolRegistry } from "../tool/registry"
import { AgentNotFoundError } from "./error"
import { SessionHistory } from "./history"
import { InstructionEntry } from "./instruction-entry"
@ -25,6 +26,7 @@ export interface Selection {
readonly session: SessionSchema.Info
readonly agent: AgentV2.Selection & { readonly info: AgentV2.Info }
readonly instructions: Instructions.Instructions
readonly toolSet: ToolRegistry.ToolSet
}
export interface Loaded {
@ -33,15 +35,17 @@ export interface Loaded {
readonly model: SessionRunnerModel.Resolved
readonly initial: string
readonly messages: ReadonlyArray<SessionMessage.Info>
readonly toolSet: ToolRegistry.ToolSet
}
/**
* Resolves model-request state in two phases: `select` fixes the Session,
* agent, and instruction sources; `load` adds the model and active history for
* that selection. This module does not build or execute the model request.
* agent, instruction sources, and tool snapshot; `load` adds the model and
* active history for that selection. This module does not build or execute the
* model request.
*/
export interface Interface {
/** Selects the Session, agent, and instruction sources used by subsequent work. */
/** Selects the Session, agent, instructions, and tools used by subsequent work. */
readonly select: (sessionID: SessionSchema.ID) => Effect.Effect<Selection, AgentNotFoundError>
/** Resolves the model and active history for that selection. */
readonly load: (selection: Selection) => Effect.Effect<Loaded, SessionRunnerModel.Error>
@ -55,7 +59,6 @@ const layer = Layer.effect(
Effect.gen(function* () {
const agents = yield* AgentV2.Service
const builtins = yield* InstructionBuiltIns.Service
const codeModeInstructions = yield* CodeModeInstructions.Service
const db = (yield* Database.Service).db
const discovery = yield* InstructionDiscovery.Service
const entries = yield* InstructionEntry.Service
@ -66,6 +69,7 @@ const layer = Layer.effect(
const referenceInstructions = yield* ReferenceInstructions.Service
const skillInstructions = yield* SkillInstructions.Service
const store = yield* SessionStore.Service
const registry = yield* ToolRegistry.Service
const select = Effect.fn("SessionContext.select")(function* (sessionID: SessionSchema.ID) {
const session = yield* store.get(sessionID)
@ -76,19 +80,32 @@ const layer = Layer.effect(
yield* plugins.flush
const agent = yield* agents.select(session.agent)
if (!agent.info) return yield* new AgentNotFoundError({ sessionID: session.id, agent: session.agent ?? agent.id })
const instructions = yield* Effect.all(
[
builtins.load(sessionID),
codeModeInstructions.load(agent),
discovery.load(),
skillInstructions.load(agent),
referenceInstructions.load(),
mcpInstructions.load(agent),
entries.load(sessionID),
],
const loaded = yield* Effect.all(
{
toolSet: registry.snapshot(agent.info.permissions),
builtins: builtins.load(sessionID),
discovery: discovery.load(),
skills: skillInstructions.load(agent),
references: referenceInstructions.load(),
mcp: mcpInstructions.load(agent),
entries: entries.load(sessionID),
},
{ concurrency: "unbounded" },
).pipe(Effect.map(Instructions.combine))
return { session, agent: { ...agent, info: agent.info }, instructions }
)
return {
session,
agent: { ...agent, info: agent.info },
instructions: Instructions.combine([
loaded.builtins,
CodeModeInstructions.make(loaded.toolSet.codeModeInstructions),
loaded.discovery,
loaded.skills,
loaded.references,
loaded.mcp,
loaded.entries,
]),
toolSet: loaded.toolSet,
}
})
const load = Effect.fn("SessionContext.load")(function* (selection: Selection) {
@ -100,6 +117,7 @@ const layer = Layer.effect(
model,
initial: history.initial,
messages: history.entries.map((entry) => entry.message),
toolSet: selection.toolSet,
}
})
@ -112,7 +130,6 @@ export const node = makeLocationNode({
layer,
deps: [
AgentV2.node,
CodeModeInstructions.node,
Database.node,
InstructionBuiltIns.node,
InstructionDiscovery.node,
@ -124,5 +141,6 @@ export const node = makeLocationNode({
SessionRunnerModel.node,
SessionStore.node,
SkillInstructions.node,
ToolRegistry.node,
],
})

View file

@ -12,7 +12,6 @@ import { SessionGenerate } from "./generate"
import { SessionHistory } from "./history"
import { SessionModelHeaders } from "./model-headers"
import { SessionRunnerModel } from "./runner/model"
import { ToolRegistry } from "../tool/registry"
import PROMPT_DEFAULT from "./runner/prompt/base.txt"
import { toLLMMessages } from "./runner/to-llm-message"
@ -24,7 +23,6 @@ export const layer = Layer.effect(
const hooks = yield* PluginHooks.Service
const llm = yield* LLMClient.Service
const models = yield* SessionRunnerModel.Service
const registry = yield* ToolRegistry.Service
const app = yield* App.Metadata
return SessionGenerate.Service.of({
@ -36,7 +34,7 @@ export const layer = Layer.effect(
const promptCacheKey = /^ses_[0-9a-f]{64}$/.test(selection.session.id)
? selection.session.id.slice(4)
: selection.session.id
const toolSet = yield* registry.snapshot(selection.agent.info.permissions)
const toolSet = selection.toolSet
const toolDefinitions = toolSet.definitions
const toolsByName = new Map(toolDefinitions.map((tool) => [tool.name, tool]))
const contextEvent = yield* hooks.trigger("session", "context", {
@ -89,13 +87,5 @@ export const layer = Layer.effect(
export const node = makeLocationNode({
service: SessionGenerate.Service,
layer,
deps: [
SessionContext.node,
Database.node,
PluginHooks.node,
SessionRunnerModel.node,
ToolRegistry.node,
App.node,
llmClient,
],
deps: [SessionContext.node, Database.node, PluginHooks.node, SessionRunnerModel.node, App.node, llmClient],
})

View file

@ -86,7 +86,6 @@ export const layer = Layer.effect(
Service,
Effect.gen(function* () {
const hooks = yield* PluginHooks.Service
const registry = yield* ToolRegistry.Service
const app = yield* App.Metadata
const prepare = Effect.fn("SessionModelRequest.prepare")(function* (input: PrepareInput) {
@ -98,7 +97,7 @@ export const layer = Layer.effect(
const stepLimitReached = agent.info.steps !== undefined && input.step >= agent.info.steps
// The final Step keeps definitions available to protocols with native "none",
// preserving their prompt cache prefix. Calls are still rejected at execution.
const toolSet = yield* registry.snapshot(agent.info.permissions)
const toolSet = input.context.toolSet
const promptCacheKey = /^ses_[0-9a-f]{64}$/.test(session.id) ? session.id.slice(4) : session.id
const system = [agent.info.system ? agent.info.system : PROMPT_DEFAULT, input.context.initial]
.filter((part) => part.length > 0)
@ -162,5 +161,5 @@ export const layer = Layer.effect(
export const node = makeLocationNode({
service: Service,
layer,
deps: [PluginHooks.node, ToolRegistry.node, App.node],
deps: [PluginHooks.node, App.node],
})

View file

@ -44,12 +44,13 @@ export interface Interface {
}
/**
* One request-scoped snapshot pairing advertised definitions with captured
* tools. A model request executes exactly the tool values it advertised
* even if registration changes while the request is in flight.
* One request-scoped snapshot pairing Code Mode instructions and advertised
* definitions with captured tools. A model request executes exactly the tool
* values it advertised even if registration changes while it is in flight.
*/
export interface ToolSet {
readonly definitions: ReadonlyArray<ToolDefinition>
readonly codeModeInstructions?: string
readonly execute: (input: ExecuteInput) => Effect.Effect<ToolOutcome, ToolOutputStore.Error>
}
@ -320,8 +321,12 @@ const registryLayer = Layer.effect(
if (whollyDisabled(registration.permission, rules)) continue
direct.set(name, registration)
}
const codemodeTool = (yield* codeMode.materialize(permissions)).tool
const codeModeMaterialization = yield* codeMode.materialize(permissions)
const codemodeTool = codeModeMaterialization.tool
return {
...(codeModeMaterialization.instructions === undefined
? {}
: { codeModeInstructions: codeModeMaterialization.instructions }),
definitions: [
// Definitions are prompt-cache prefix bytes, so order only after effective registrations settle.
...Array.from(direct)

View file

@ -1,15 +1,12 @@
import { describe, expect } from "bun:test"
import { AgentV2 } from "@opencode-ai/core/agent"
import { CodeMode } from "@opencode-ai/core/codemode"
import { CodeModeInstructions } from "@opencode-ai/core/codemode/instructions"
import { AppNodeBuilder } from "@opencode-ai/core/effect/app-node-builder"
import { Tool } from "@opencode-ai/core/tool/tool"
import { Effect, Layer, Schema } from "effect"
import { Effect, Schema } from "effect"
import { it } from "../lib/effect"
import { readInitial, readUpdate } from "../lib/instructions"
const agent = AgentV2.Info.make(AgentV2.Info.empty(AgentV2.ID.make("build")))
describe("CodeModeInstructions", () => {
it.effect("treats equivalent registration orders as an instruction no-op", () => {
const alpha = Tool.make({
@ -24,70 +21,46 @@ describe("CodeModeInstructions", () => {
output: Schema.String,
execute: () => Effect.succeed({ output: "zeta" }),
})
const codeModeLayer = AppNodeBuilder.build(CodeMode.node)
const layer = Layer.merge(
codeModeLayer,
AppNodeBuilder.build(CodeModeInstructions.node, [[CodeMode.node, codeModeLayer]]),
)
return Effect.gen(function* () {
const codeMode = yield* CodeMode.Service
const instructions = yield* CodeModeInstructions.Service
const initialized = yield* Effect.scoped(
Effect.gen(function* () {
yield* codeMode.register(Tool.registrationEntries({ zeta, alpha }, { namespace: "tools" }))
return yield* instructions.load({ id: agent.id, info: agent }).pipe(Effect.flatMap(readInitial))
const materialization = yield* codeMode.materialize()
return yield* readInitial(CodeModeInstructions.make(materialization.instructions))
}),
)
const reordered = yield* Effect.scoped(
Effect.gen(function* () {
yield* codeMode.register(Tool.registrationEntries({ alpha, zeta }, { namespace: "tools" }))
return yield* instructions
.load({ id: agent.id, info: agent })
.pipe(Effect.flatMap((context) => readUpdate(context, initialized)))
const materialization = yield* codeMode.materialize()
return yield* readUpdate(CodeModeInstructions.make(materialization.instructions), initialized)
}),
)
expect(reordered.changed).toBe(false)
expect(reordered.text).toBe("")
}).pipe(Effect.provide(layer))
}).pipe(Effect.provide(codeModeLayer))
})
it.effect("renders catalog changes and removal", () => {
let catalog: string | undefined = "Initial Code Mode catalog"
const layer = AppNodeBuilder.build(CodeModeInstructions.node, [
[
CodeMode.node,
Layer.mock(CodeMode.Service, {
materialize: () => Effect.succeed({ ...(catalog === undefined ? {} : { instructions: catalog }) }),
register: () => Effect.void,
}),
],
])
return Effect.gen(function* () {
const instructions = yield* CodeModeInstructions.Service
const initialized = yield* instructions.load({ id: agent.id, info: agent }).pipe(Effect.flatMap(readInitial))
const initialized = yield* readInitial(CodeModeInstructions.make(catalog))
expect(initialized.text).toBe("Initial Code Mode catalog")
catalog = "Updated Code Mode catalog"
expect(
yield* instructions
.load({ id: agent.id, info: agent })
.pipe(Effect.flatMap((context) => readUpdate(context, initialized))),
).toMatchObject({
expect(yield* readUpdate(CodeModeInstructions.make(catalog), initialized)).toMatchObject({
text: "The Code Mode tool catalog has changed. This catalog supersedes the previous Code Mode tool catalog.\n\nUpdated Code Mode catalog",
})
catalog = undefined
expect(
yield* instructions
.load({ id: agent.id, info: agent })
.pipe(Effect.flatMap((context) => readUpdate(context, initialized))),
).toMatchObject({
expect(yield* readUpdate(CodeModeInstructions.make(catalog), initialized)).toMatchObject({
text: "Code Mode tools are no longer available. Do not use any previously listed Code Mode tools.",
})
}).pipe(Effect.provide(layer))
})
})
})

View file

@ -578,7 +578,8 @@ describe("LocationServiceMap", () => {
const blockedState = yield* update(blocked.path, blockedID)
expect(blockedState.providers.some((provider) => provider.id === blockedID)).toBe(true)
expect(blockedState.providers.some((provider) => provider.id === allowedID)).toBe(false)
expect(blockedState.tools.map((tool) => tool.name).sort()).toEqual([
const blockedTools = blockedState.tools.map((tool) => tool.name)
expect(blockedTools.filter((name) => name !== "execute").sort()).toEqual([
"edit",
"glob",
"grep",
@ -595,7 +596,9 @@ describe("LocationServiceMap", () => {
const allowedState = yield* update(allowed.path, allowedID)
expect(allowedState.providers.some((provider) => provider.id === allowedID)).toBe(true)
expect(allowedState.providers.some((provider) => provider.id === blockedID)).toBe(false)
expect(allowedState.tools.map((tool) => tool.name).sort()).toEqual([
const allowedTools = allowedState.tools.map((tool) => tool.name)
expect(blockedTools.includes("execute")).toBe(allowedTools.includes("execute"))
expect(allowedTools.filter((name) => name !== "execute").sort()).toEqual([
"edit",
"glob",
"grep",

View file

@ -97,6 +97,7 @@ const plugins = Layer.mock(PluginSupervisor.Service, { flush: Effect.void })
const tools = Layer.mock(ToolRegistry.Service, {
snapshot: () =>
Effect.succeed({
codeModeInstructions: "Captured Code Mode catalog",
definitions: [ToolDefinition.make({ name: "lookup", description: "Lookup", inputSchema: { type: "object" } })],
execute: () => Effect.die(new Error("unused")),
}),
@ -285,13 +286,14 @@ it.effect("generates from fresh settled Session context without durable mutation
expect(requests[0]?.system.map((part) => part.text)).toContain("Initial context")
expect(requests[0]?.http?.headers).toMatchObject({ "X-Session-Id": sessionID })
expect(requests[0]?.providerOptions).toMatchObject({ openai: { promptCacheKey: sessionID } })
expect(
requests[0]?.messages.flatMap((message) =>
message.role === "system"
? message.content.flatMap((content) => (content.type === "text" ? [content.text] : []))
: [],
),
).toEqual(["Changed context"])
const instructionUpdates = requests[0]?.messages.flatMap((message) =>
message.role === "system"
? message.content.flatMap((content) => (content.type === "text" ? [content.text] : []))
: [],
)
expect(instructionUpdates).toHaveLength(1)
expect(instructionUpdates?.[0]).toContain("Changed context")
expect(instructionUpdates?.[0]).toContain("Captured Code Mode catalog")
expect(userTexts(requests[0])).toEqual(["Existing durable context", "Summarize privately"])
expect(
requests[0]?.messages.flatMap((message) =>

View file

@ -533,6 +533,7 @@ describe("ToolRegistry", () => {
.pipe(Scope.provide(scope))
const toolSet = yield* service.snapshot()
const execute = toolSet.definitions.find((tool) => tool.name === "execute")
expect(toolSet.codeModeInstructions).toContain("tools.echo")
expect(execute?.description).toContain("confined Code Mode runtime")
expect(execute?.description).not.toContain("Echo text")
yield* Scope.close(scope, Exit.void)

View file

@ -43,6 +43,7 @@ import * as SessionRunnerLLM from "@opencode-ai/core/session/runner/llm"
import { SessionRunnerModel } from "@opencode-ai/core/session/runner/model"
import { SessionUsage } from "@opencode-ai/core/session/usage"
import { ToolRegistry } from "@opencode-ai/core/tool/registry"
import { CodeMode } from "@opencode-ai/core/codemode"
import { PluginSupervisor } from "@opencode-ai/core/plugin/supervisor"
import { PluginHooks } from "@opencode-ai/core/plugin/hooks"
import { SystemPromptPlugin } from "@opencode-ai/core/plugin/system-prompt"
@ -368,6 +369,12 @@ const pluginSupervisor = Layer.succeed(
flush: Effect.suspend(() => pluginFlushHook),
}),
)
let codeModeMaterializations: ReadonlyArray<CodeMode.Materialization> = []
let codeModeMaterializationCount = 0
const codeMode = Layer.mock(CodeMode.Service, {
register: () => Effect.void,
materialize: () => Effect.sync(() => codeModeMaterializations[codeModeMaterializationCount++] ?? {}),
})
const promptCatalog = Layer.mock(Catalog.Service, {
provider: {
get: () => Effect.succeed(undefined),
@ -405,6 +412,7 @@ const runnerLayer = AppNodeBuilder.build(SessionRunnerLLM.node, [
[McpInstructions.node, mcpInstructions],
[ToolOutputStore.node, toolOutputStore],
[PluginSupervisor.node, pluginSupervisor],
[CodeMode.node, codeMode],
])
const execution = Layer.effect(
SessionExecution.Service,
@ -464,6 +472,7 @@ const it = testEffect(
[Config.node, config],
[ToolOutputStore.node, toolOutputStore],
[PluginSupervisor.node, pluginSupervisor],
[CodeMode.node, codeMode],
],
),
)
@ -512,6 +521,8 @@ const setup = Effect.gen(function* () {
systemLoadHook = Effect.void
modelResolveHook = Effect.void
pluginFlushHook = Effect.void
codeModeMaterializations = []
codeModeMaterializationCount = 0
currentModel = model
skillBaselines.clear()
responses = undefined
@ -823,6 +834,45 @@ const verifyPartialFlushOnInterruption = (kind: FragmentKind) =>
})
describe("SessionRunnerLLM", () => {
it.effect("uses one Code Mode materialization per request for instructions and execution", () =>
Effect.gen(function* () {
const executed: string[] = []
const execute = (name: string) =>
Tool.make({
description: `Execute ${name}`,
input: Schema.Struct({}),
output: Schema.String,
execute: () => Effect.sync(() => executed.push(name)).pipe(Effect.as({ output: name })),
})
const session = yield* setup
codeModeMaterializations = [
{ instructions: "Code Mode catalog A", tool: execute("A") },
{ instructions: "Code Mode catalog B", tool: execute("B") },
{ instructions: "Code Mode catalog C", tool: execute("C") },
{ instructions: "Code Mode catalog D", tool: execute("D") },
]
yield* admit(session, "Use Code Mode")
responses = [reply.tool("call-execute", "execute", {}), reply.stop()]
yield* session.resume(sessionID)
expect(requests).toHaveLength(2)
expect(codeModeMaterializationCount).toBe(2)
expect(requests[0]?.system.some((part) => part.text.includes("Code Mode catalog A"))).toBe(true)
expect(requests[0]?.system.some((part) => part.text.includes("Code Mode catalog B"))).toBe(false)
expect(requests[0]?.tools.find((tool) => tool.name === "execute")?.description).toBe("Execute A")
expect(executed).toEqual(["A"])
expect(requests[1]?.tools.find((tool) => tool.name === "execute")?.description).toBe("Execute B")
expect(
requests[1]?.messages.some(
(message) =>
message.role === "system" &&
message.content.some((part) => part.type === "text" && part.text.includes("Code Mode catalog B")),
),
).toBe(true)
}),
)
it.effect("applies session context hooks without exposing unavailable tools", () =>
Effect.gen(function* () {
const session = yield* setup