From 3785eddfa0c045ec5aaadbeaf2d13042fcec484d Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Thu, 9 Jul 2026 21:16:00 -0400 Subject: [PATCH] fix(core): preserve admitted tool generations (#36177) --- packages/core/src/session/runner/llm.ts | 6 +- packages/core/src/tool/execute.ts | 11 +-- packages/core/src/tool/registry.ts | 17 +---- .../test/session-runner-tool-registry.test.ts | 73 ++++++++++++++----- packages/core/test/session-runner.test.ts | 71 +++++++++++++++++- 5 files changed, 127 insertions(+), 51 deletions(-) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 2d31338b90..556c1168dc 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -173,7 +173,7 @@ const layer = Layer.effect( sessionID, assistantMessageID: message.id, callID: tool.id, - error: { type: "tool.stale", message: `Tool execution interrupted: ${tool.name}` }, + error: { type: "aborted", message: `Tool execution interrupted: ${tool.name}` }, executed: tool.executed === true, }) } @@ -466,9 +466,7 @@ const layer = Layer.effect( // implementation becomes a failed tool call the model can read, and the step still // settles so the model may recover. A typed infrastructure failure (tool output // could not be persisted) also fails the assistant and then fails the drain. - const settledFailure = settledCauses.find( - (cause) => !Cause.hasInterrupts(cause) && !isUserDeclined(cause), - ) + const settledFailure = settledCauses.find((cause) => !Cause.hasInterrupts(cause) && !isUserDeclined(cause)) const infraError = settledFailure === undefined ? undefined : Option.getOrUndefined(Cause.findErrorOption(settledFailure)) if (settledFailure !== undefined) { diff --git a/packages/core/src/tool/execute.ts b/packages/core/src/tool/execute.ts index f5bd776e5a..d9769c396e 100644 --- a/packages/core/src/tool/execute.ts +++ b/packages/core/src/tool/execute.ts @@ -37,16 +37,12 @@ type CollectedFiles = { } export interface Registration { - readonly identity: object readonly tool: AnyTool readonly name: string readonly group?: string } -export const create = (options: { - readonly registrations: ReadonlyMap - readonly current: (name: string) => Registration | undefined -}) => { +export const create = (options: { readonly registrations: ReadonlyMap }) => { const runtime = ( invoke: (name: string, registration: Registration, input: unknown) => Effect.Effect, hooks?: CodeMode.ToolCallHooks, @@ -115,11 +111,8 @@ export const create = (options: { (name, registration, input) => Effect.gen(function* () { const index = yield* Ref.getAndUpdate(callIndex, (index) => index + 1) - const current = options.current(name) - if (!current || current.identity !== registration.identity) - return yield* Effect.fail(toolError(`Stale tool call: ${name}`)) const output = yield* settle( - current.tool, + registration.tool, { type: "tool-call", id: context.toolCallID, name, input }, { sessionID: context.sessionID, diff --git a/packages/core/src/tool/registry.ts b/packages/core/src/tool/registry.ts index cef6f3c071..1588c81005 100644 --- a/packages/core/src/tool/registry.ts +++ b/packages/core/src/tool/registry.ts @@ -58,7 +58,6 @@ const registryLayer = Layer.effect( const resources = yield* ToolOutputStore.Service const toolHooks = yield* ToolHooks.Service type Registration = { - readonly identity: object readonly tool: AnyTool readonly name: string readonly group?: string @@ -134,18 +133,6 @@ const registryLayer = Layer.effect( } }) - const settleWith = Effect.fn("ToolRegistry.settle")(function* (input: ExecuteInput, advertised: object) { - const registration = local.get(input.call.name)?.at(-1)?.registration - if (!registration || registration.identity !== advertised) { - const message = `Stale tool call: ${input.call.name}` - return { - result: { type: "error" as const, value: message }, - error: { type: "tool.stale" as const, message }, - } - } - return yield* settleTool(input, registration.tool) - }) - return Service.of({ register: Effect.fn("ToolRegistry.register")(function* (tools, options) { const entries = registrationEntries(tools, options?.group) @@ -164,7 +151,6 @@ const registryLayer = Layer.effect( { token, registration: { - identity: {}, tool: entry.tool, name: entry.name, group: entry.group, @@ -204,7 +190,6 @@ const registryLayer = Layer.effect( deferred.size > 0 && !whollyDisabled("execute", input.permissions ?? []) ? ExecuteTool.create({ registrations: deferred, - current: (name) => local.get(name)?.at(-1)?.registration, }) : undefined return { @@ -215,7 +200,7 @@ const registryLayer = Layer.effect( settle: (input) => { if (input.call.name === "execute" && execute) return settleTool(input, execute) const registration = direct.get(input.call.name) - if (registration) return settleWith(input, registration.identity) + if (registration) return settleTool(input, registration.tool) return Effect.succeed({ result: { type: "error", value: `Unknown tool: ${input.call.name}` }, error: { type: "tool.unknown", message: `Unknown tool: ${input.call.name}` }, diff --git a/packages/core/test/session-runner-tool-registry.test.ts b/packages/core/test/session-runner-tool-registry.test.ts index 2944f07972..a1c11583a4 100644 --- a/packages/core/test/session-runner-tool-registry.test.ts +++ b/packages/core/test/session-runner-tool-registry.test.ts @@ -97,12 +97,7 @@ describe("ToolRegistry", () => { .pipe(Effect.map((materialized) => materialized.definitions.map((tool) => tool.name))) expect(yield* names({ id: "gpt-5", provider: "openai" })).toEqual(["read", "edit", "write", "patch"]) - expect(yield* names({ id: "claude-sonnet-4", provider: "anthropic" })).toEqual([ - "read", - "edit", - "write", - "patch", - ]) + expect(yield* names({ id: "claude-sonnet-4", provider: "anthropic" })).toEqual(["read", "edit", "write", "patch"]) }), ) @@ -355,7 +350,7 @@ describe("ToolRegistry", () => { }), ) - it.effect("rejects a call when its advertised registration was removed", () => + it.effect("executes the advertised registration after it is removed", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service const scope = yield* Scope.make() @@ -363,29 +358,23 @@ describe("ToolRegistry", () => { const materialized = yield* service.materialize({ model: testModel }) yield* Scope.close(scope, Exit.void) - expect((yield* materialized.settle(call("echo"))).result).toEqual({ - type: "error", - value: "Stale tool call: echo", - }) + expect((yield* materialized.settle(call("echo"))).result).toEqual({ type: "text", value: "echo" }) }), ) - it.effect("rejects only the replaced name from a multi-tool provider turn", () => + it.effect("executes each registration advertised for a provider turn after replacement", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service yield* service.register({ first: make(), second: make() }) const materialized = yield* service.materialize({ model: testModel }) yield* service.register({ first: make() }) - expect((yield* materialized.settle(call("first"))).result).toEqual({ - type: "error", - value: "Stale tool call: first", - }) + expect((yield* materialized.settle(call("first"))).result).toEqual({ type: "text", value: "first" }) expect((yield* materialized.settle(call("second"))).result).toEqual({ type: "text", value: "second" }) }), ) - it.effect("treats revealing a previous overlay as stale", () => + it.effect("executes an advertised overlay after the previous registration is revealed", () => Effect.gen(function* () { const service = yield* ToolRegistry.Service yield* service.register({ echo: make() }) @@ -394,10 +383,54 @@ describe("ToolRegistry", () => { const materialized = yield* service.materialize({ model: testModel }) yield* Scope.close(overlay, Exit.void) - expect((yield* materialized.settle(call("echo"))).result).toEqual({ - type: "error", - value: "Stale tool call: echo", + expect((yield* materialized.settle(call("echo"))).result).toEqual({ type: "text", value: "echo" }) + }), + ) + + it.effect("executes deferred registrations from the advertised generation", () => + Effect.gen(function* () { + const service = yield* ToolRegistry.Service + const executed: string[] = [] + const scope = yield* Scope.make() + yield* service + .register( + { + echo: Tool.make({ + description: "Echo text", + input: Schema.Struct({ text: Schema.String }), + output: Schema.Struct({ text: Schema.String }), + execute: ({ text }) => Effect.sync(() => executed.push(`old:${text}`)).pipe(Effect.as({ text })), + }), + }, + { deferred: true }, + ) + .pipe(Scope.provide(scope)) + const materialized = yield* service.materialize({ model: testModel }) + yield* Scope.close(scope, Exit.void) + yield* service.register( + { + echo: Tool.make({ + description: "Echo text", + input: Schema.Struct({ text: Schema.String }), + output: Schema.Struct({ text: Schema.String }), + execute: ({ text }) => Effect.sync(() => executed.push(`new:${text}`)).pipe(Effect.as({ text })), + }), + }, + { deferred: true }, + ) + + const settlement = yield* materialized.settle({ + ...call("execute"), + call: { + type: "tool-call", + id: "call-execute", + name: "execute", + input: { code: 'return await tools.echo({ text: "admitted" })' }, + }, }) + + expect(settlement.result).toMatchObject({ type: "text" }) + expect(executed).toEqual(["old:admitted"]) }), ) diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 73e1cbf461..c1aae6333d 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -62,7 +62,7 @@ import { McpGuidance } from "@opencode-ai/core/mcp/guidance" import { ModelV2 } from "@opencode-ai/core/model" import { Location } from "@opencode-ai/core/location" import { ProviderV2 } from "@opencode-ai/core/provider" -import { Cause, DateTime, Deferred, Effect, Exit, Fiber, Layer, Schema, Stream } from "effect" +import { Cause, DateTime, Deferred, Effect, Exit, Fiber, Layer, Schema, Scope, Stream } from "effect" import { TestClock } from "effect/testing" import { asc, eq } from "drizzle-orm" import { testEffect } from "./lib/effect" @@ -788,6 +788,73 @@ describe("SessionRunnerLLM", () => { }), ) + it.effect("executes parallel tool calls against the generation admitted before a registry reload", () => + Effect.gen(function* () { + const session = yield* setup + const registry = yield* ToolRegistry.Service + const scope = yield* Scope.make() + const generation: string[] = [] + yield* registry + .register({ + reloaded: Tool.make({ + description: "Record the admitted generation", + input: Schema.Struct({}), + output: Schema.Struct({ generation: Schema.String }), + execute: () => Effect.sync(() => generation.push("admitted")).pipe(Effect.as({ generation: "admitted" })), + }), + }) + .pipe(Scope.provide(scope)) + yield* admit(session, "Use the reloaded tool") + responses = [ + [ + LLMEvent.stepStart({ index: 0 }), + LLMEvent.toolCall({ id: "call-reloaded-1", name: "reloaded", input: {} }), + LLMEvent.toolCall({ id: "call-reloaded-2", name: "reloaded", input: {} }), + LLMEvent.stepFinish({ index: 0, reason: "tool-calls" }), + LLMEvent.finish({ reason: "tool-calls" }), + ], + [], + ] + streamGate = yield* Deferred.make() + streamStarted = yield* Deferred.make() + + const run = yield* session.resume(sessionID).pipe(Effect.forkChild) + yield* Deferred.await(streamStarted) + yield* Scope.close(scope, Exit.void) + yield* registry.register({ + reloaded: Tool.make({ + description: "Record the replacement generation", + input: Schema.Struct({}), + output: Schema.Struct({ generation: Schema.String }), + execute: () => + Effect.sync(() => generation.push("replacement")).pipe(Effect.as({ generation: "replacement" })), + }), + }) + yield* Deferred.succeed(streamGate, undefined) + yield* Fiber.join(run) + + expect(generation).toEqual(["admitted", "admitted"]) + expect(yield* session.context(sessionID)).toMatchObject([ + { type: "user", text: "Use the reloaded tool" }, + { + type: "assistant", + content: [ + { + type: "tool", + id: "call-reloaded-1", + state: { status: "completed", structured: { generation: "admitted" } }, + }, + { + type: "tool", + id: "call-reloaded-2", + state: { status: "completed", structured: { generation: "admitted" } }, + }, + ], + }, + ]) + }), + ) + it.effect("starts a real runner turn after default prompt recording", () => Effect.gen(function* () { const session = yield* setup @@ -2659,7 +2726,7 @@ describe("SessionRunnerLLM", () => { id: "call-interrupted", state: { status: "error", - error: { type: "tool.stale", message: "Tool execution interrupted: echo" }, + error: { type: "aborted", message: "Tool execution interrupted: echo" }, }, }, ],