fix(core): preserve admitted tool generations (#36177)

This commit is contained in:
Kit Langton 2026-07-09 21:16:00 -04:00 committed by GitHub
commit 3785eddfa0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 127 additions and 51 deletions

View file

@ -173,7 +173,7 @@ const layer = Layer.effect(
sessionID, sessionID,
assistantMessageID: message.id, assistantMessageID: message.id,
callID: tool.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, 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 // 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 // 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. // could not be persisted) also fails the assistant and then fails the drain.
const settledFailure = settledCauses.find( const settledFailure = settledCauses.find((cause) => !Cause.hasInterrupts(cause) && !isUserDeclined(cause))
(cause) => !Cause.hasInterrupts(cause) && !isUserDeclined(cause),
)
const infraError = const infraError =
settledFailure === undefined ? undefined : Option.getOrUndefined(Cause.findErrorOption(settledFailure)) settledFailure === undefined ? undefined : Option.getOrUndefined(Cause.findErrorOption(settledFailure))
if (settledFailure !== undefined) { if (settledFailure !== undefined) {

View file

@ -37,16 +37,12 @@ type CollectedFiles = {
} }
export interface Registration { export interface Registration {
readonly identity: object
readonly tool: AnyTool readonly tool: AnyTool
readonly name: string readonly name: string
readonly group?: string readonly group?: string
} }
export const create = (options: { export const create = (options: { readonly registrations: ReadonlyMap<string, Registration> }) => {
readonly registrations: ReadonlyMap<string, Registration>
readonly current: (name: string) => Registration | undefined
}) => {
const runtime = ( const runtime = (
invoke: (name: string, registration: Registration, input: unknown) => Effect.Effect<unknown, unknown>, invoke: (name: string, registration: Registration, input: unknown) => Effect.Effect<unknown, unknown>,
hooks?: CodeMode.ToolCallHooks, hooks?: CodeMode.ToolCallHooks,
@ -115,11 +111,8 @@ export const create = (options: {
(name, registration, input) => (name, registration, input) =>
Effect.gen(function* () { Effect.gen(function* () {
const index = yield* Ref.getAndUpdate(callIndex, (index) => index + 1) 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( const output = yield* settle(
current.tool, registration.tool,
{ type: "tool-call", id: context.toolCallID, name, input }, { type: "tool-call", id: context.toolCallID, name, input },
{ {
sessionID: context.sessionID, sessionID: context.sessionID,

View file

@ -58,7 +58,6 @@ const registryLayer = Layer.effect(
const resources = yield* ToolOutputStore.Service const resources = yield* ToolOutputStore.Service
const toolHooks = yield* ToolHooks.Service const toolHooks = yield* ToolHooks.Service
type Registration = { type Registration = {
readonly identity: object
readonly tool: AnyTool readonly tool: AnyTool
readonly name: string readonly name: string
readonly group?: 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({ return Service.of({
register: Effect.fn("ToolRegistry.register")(function* (tools, options) { register: Effect.fn("ToolRegistry.register")(function* (tools, options) {
const entries = registrationEntries(tools, options?.group) const entries = registrationEntries(tools, options?.group)
@ -164,7 +151,6 @@ const registryLayer = Layer.effect(
{ {
token, token,
registration: { registration: {
identity: {},
tool: entry.tool, tool: entry.tool,
name: entry.name, name: entry.name,
group: entry.group, group: entry.group,
@ -204,7 +190,6 @@ const registryLayer = Layer.effect(
deferred.size > 0 && !whollyDisabled("execute", input.permissions ?? []) deferred.size > 0 && !whollyDisabled("execute", input.permissions ?? [])
? ExecuteTool.create({ ? ExecuteTool.create({
registrations: deferred, registrations: deferred,
current: (name) => local.get(name)?.at(-1)?.registration,
}) })
: undefined : undefined
return { return {
@ -215,7 +200,7 @@ const registryLayer = Layer.effect(
settle: (input) => { settle: (input) => {
if (input.call.name === "execute" && execute) return settleTool(input, execute) if (input.call.name === "execute" && execute) return settleTool(input, execute)
const registration = direct.get(input.call.name) const registration = direct.get(input.call.name)
if (registration) return settleWith(input, registration.identity) if (registration) return settleTool(input, registration.tool)
return Effect.succeed({ return Effect.succeed({
result: { type: "error", value: `Unknown tool: ${input.call.name}` }, result: { type: "error", value: `Unknown tool: ${input.call.name}` },
error: { type: "tool.unknown", message: `Unknown tool: ${input.call.name}` }, error: { type: "tool.unknown", message: `Unknown tool: ${input.call.name}` },

View file

@ -97,12 +97,7 @@ describe("ToolRegistry", () => {
.pipe(Effect.map((materialized) => materialized.definitions.map((tool) => tool.name))) .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: "gpt-5", provider: "openai" })).toEqual(["read", "edit", "write", "patch"])
expect(yield* names({ id: "claude-sonnet-4", provider: "anthropic" })).toEqual([ expect(yield* names({ id: "claude-sonnet-4", provider: "anthropic" })).toEqual(["read", "edit", "write", "patch"])
"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* () { Effect.gen(function* () {
const service = yield* ToolRegistry.Service const service = yield* ToolRegistry.Service
const scope = yield* Scope.make() const scope = yield* Scope.make()
@ -363,29 +358,23 @@ describe("ToolRegistry", () => {
const materialized = yield* service.materialize({ model: testModel }) const materialized = yield* service.materialize({ model: testModel })
yield* Scope.close(scope, Exit.void) yield* Scope.close(scope, Exit.void)
expect((yield* materialized.settle(call("echo"))).result).toEqual({ expect((yield* materialized.settle(call("echo"))).result).toEqual({ type: "text", value: "echo" })
type: "error",
value: "Stale tool call: 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* () { Effect.gen(function* () {
const service = yield* ToolRegistry.Service const service = yield* ToolRegistry.Service
yield* service.register({ first: make(), second: make() }) yield* service.register({ first: make(), second: make() })
const materialized = yield* service.materialize({ model: testModel }) const materialized = yield* service.materialize({ model: testModel })
yield* service.register({ first: make() }) yield* service.register({ first: make() })
expect((yield* materialized.settle(call("first"))).result).toEqual({ expect((yield* materialized.settle(call("first"))).result).toEqual({ type: "text", value: "first" })
type: "error",
value: "Stale tool call: first",
})
expect((yield* materialized.settle(call("second"))).result).toEqual({ type: "text", value: "second" }) 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* () { Effect.gen(function* () {
const service = yield* ToolRegistry.Service const service = yield* ToolRegistry.Service
yield* service.register({ echo: make() }) yield* service.register({ echo: make() })
@ -394,10 +383,54 @@ describe("ToolRegistry", () => {
const materialized = yield* service.materialize({ model: testModel }) const materialized = yield* service.materialize({ model: testModel })
yield* Scope.close(overlay, Exit.void) yield* Scope.close(overlay, Exit.void)
expect((yield* materialized.settle(call("echo"))).result).toEqual({ expect((yield* materialized.settle(call("echo"))).result).toEqual({ type: "text", value: "echo" })
type: "error", }),
value: "Stale tool call: 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"])
}), }),
) )

View file

@ -62,7 +62,7 @@ import { McpGuidance } from "@opencode-ai/core/mcp/guidance"
import { ModelV2 } from "@opencode-ai/core/model" import { ModelV2 } from "@opencode-ai/core/model"
import { Location } from "@opencode-ai/core/location" import { Location } from "@opencode-ai/core/location"
import { ProviderV2 } from "@opencode-ai/core/provider" 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 { TestClock } from "effect/testing"
import { asc, eq } from "drizzle-orm" import { asc, eq } from "drizzle-orm"
import { testEffect } from "./lib/effect" 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<void>()
streamStarted = yield* Deferred.make<void>()
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", () => it.effect("starts a real runner turn after default prompt recording", () =>
Effect.gen(function* () { Effect.gen(function* () {
const session = yield* setup const session = yield* setup
@ -2659,7 +2726,7 @@ describe("SessionRunnerLLM", () => {
id: "call-interrupted", id: "call-interrupted",
state: { state: {
status: "error", status: "error",
error: { type: "tool.stale", message: "Tool execution interrupted: echo" }, error: { type: "aborted", message: "Tool execution interrupted: echo" },
}, },
}, },
], ],