فهرست منبع

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

Kit Langton 1 ماه پیش
والد
کامیت
3785eddfa0

+ 2 - 4
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) {

+ 2 - 9
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<string, Registration>
-  readonly current: (name: string) => Registration | undefined
-}) => {
+export const create = (options: { readonly registrations: ReadonlyMap<string, Registration> }) => {
   const runtime = (
     invoke: (name: string, registration: Registration, input: unknown) => Effect.Effect<unknown, unknown>,
     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,

+ 1 - 16
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}` },

+ 53 - 20
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"])
     }),
   )
 

+ 69 - 2
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<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", () =>
     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" },
               },
             },
           ],