Ver Fonte

fix(core): defer permission rejection cleanup (#36067)

Kit Langton há 1 mês atrás
pai
commit
aa0e626b12

+ 13 - 9
packages/core/src/session/runner/llm.ts

@@ -256,8 +256,8 @@ const layer = Layer.effect(
         tools: toolMaterialization?.definitions ?? [],
         tools: toolMaterialization?.definitions ?? [],
         toolChoice: isLastStep ? "none" : undefined,
         toolChoice: isLastStep ? "none" : undefined,
       })
       })
-      const toolFibers = yield* FiberSet.make<void, ToolOutputStore.Error | UserInterruptedError>()
-      const ownedToolFibers: Array<Fiber.Fiber<void, ToolOutputStore.Error | UserInterruptedError>> = []
+      const toolFibers = yield* FiberSet.make<void, ToolOutputStore.Error | StepFailedError>()
+      const ownedToolFibers: Array<Fiber.Fiber<void, ToolOutputStore.Error | StepFailedError>> = []
       let needsContinuation = false
       let needsContinuation = false
       const availableTools = new Map(request.tools.map((tool) => [tool.name, tool]))
       const availableTools = new Map(request.tools.map((tool) => [tool.name, tool]))
       const requestEvent: SessionHooks["request"] = {
       const requestEvent: SessionHooks["request"] = {
@@ -364,11 +364,10 @@ const layer = Layer.effect(
                       }),
                       }),
                       settlement.error,
                       settlement.error,
                     ).pipe(
                     ).pipe(
+                      // Terminal cleanup must own failAssistant because the provider may still be streaming another tool input.
                       Effect.andThen(
                       Effect.andThen(
                         settlement.error?.type === "permission.rejected"
                         settlement.error?.type === "permission.rejected"
-                          ? serialized(publisher.failAssistant(settlement.error)).pipe(
-                              Effect.andThen(Effect.fail(new UserInterruptedError())),
-                            )
+                          ? Effect.fail(new StepFailedError({ error: settlement.error }))
                           : Effect.void,
                           : Effect.void,
                       ),
                       ),
                     ),
                     ),
@@ -464,14 +463,19 @@ const layer = Layer.effect(
               : settled.value.flatMap((exit) => (exit._tag === "Failure" ? [exit.cause] : []))
               : settled.value.flatMap((exit) => (exit._tag === "Failure" ? [exit.cause] : []))
           const toolsInterrupted = settledCauses.some(Cause.hasInterrupts)
           const toolsInterrupted = settledCauses.some(Cause.hasInterrupts)
           const userDeclined = settledCauses.some(isUserDeclined)
           const userDeclined = settledCauses.some(isUserDeclined)
-          const permissionRejected = settledCauses.some(
-            (cause) => Option.getOrUndefined(Cause.findErrorOption(cause)) instanceof UserInterruptedError,
-          )
+          const permissionRejected = settledCauses
+            .map((cause) => Option.getOrUndefined(Cause.findErrorOption(cause)))
+            .find(
+              (error): error is StepFailedError =>
+                error instanceof StepFailedError && error.error.type === "permission.rejected",
+            )
 
 
           if (userDeclined || permissionRejected || streamInterrupted || toolsInterrupted) {
           if (userDeclined || permissionRejected || streamInterrupted || toolsInterrupted) {
             yield* FiberSet.clear(toolFibers)
             yield* FiberSet.clear(toolFibers)
             yield* serialized(publisher.failUnsettledTools({ type: "aborted", message: "Tool execution interrupted" }))
             yield* serialized(publisher.failUnsettledTools({ type: "aborted", message: "Tool execution interrupted" }))
-            yield* serialized(publisher.failAssistant({ type: "aborted", message: "Step interrupted" }))
+            yield* serialized(
+              publisher.failAssistant(permissionRejected?.error ?? { type: "aborted", message: "Step interrupted" }),
+            )
           }
           }
           // A settled tool fiber failure is one of two things. A defect from a tool
           // A settled tool fiber failure is one of two things. A defect from a tool
           // 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

+ 116 - 17
packages/core/test/session-runner.test.ts

@@ -64,7 +64,7 @@ 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, Stream } from "effect"
 import { TestClock } from "effect/testing"
 import { TestClock } from "effect/testing"
-import { asc, eq } from "drizzle-orm"
+import { and, asc, eq } from "drizzle-orm"
 import { testEffect } from "./lib/effect"
 import { testEffect } from "./lib/effect"
 
 
 const requests: LLMRequest[] = []
 const requests: LLMRequest[] = []
@@ -182,6 +182,20 @@ test("does not apply an ineligible tier without base pricing", () => {
 
 
 const authorizations: Tool.Context[] = []
 const authorizations: Tool.Context[] = []
 const executions: string[] = []
 const executions: string[] = []
+const permissionFail = Tool.make({
+  description: "Reject a permission",
+  input: Schema.Struct({}),
+  output: Schema.Struct({}),
+  execute: () =>
+    new ToolFailure({
+      message: "Permission denied: edit",
+      error: new PermissionV2.BlockedError({
+        rules: [],
+        permission: "edit",
+        resources: ["src/index.ts"],
+      }),
+    }),
+})
 const permission = Layer.succeed(
 const permission = Layer.succeed(
   PermissionV2.Service,
   PermissionV2.Service,
   PermissionV2.Service.of({
   PermissionV2.Service.of({
@@ -504,6 +518,22 @@ const recordedEventTypes = (id: SessionV2.ID) =>
       )
       )
   })
   })
 
 
+const recordedToolInputEnds = (id: SessionV2.ID, callID: string) =>
+  Effect.gen(function* () {
+    const { db } = yield* Database.Service
+    return (yield* db
+      .select({ data: EventTable.data })
+      .from(EventTable)
+      .where(
+        and(
+          eq(EventTable.aggregate_id, id),
+          eq(EventTable.type, EventV2.versionedType(SessionEvent.Tool.Input.Ended.type, 1)),
+        ),
+      )
+      .all()
+      .pipe(Effect.orDie)).filter((event) => event.data.callID === callID)
+  })
+
 const recordedStepSettlementEvents = (id: SessionV2.ID, assistantMessageID: SessionMessage.ID) =>
 const recordedStepSettlementEvents = (id: SessionV2.ID, assistantMessageID: SessionMessage.ID) =>
   Effect.gen(function* () {
   Effect.gen(function* () {
     const { db } = yield* Database.Service
     const { db } = yield* Database.Service
@@ -2980,22 +3010,7 @@ describe("SessionRunnerLLM", () => {
     Effect.gen(function* () {
     Effect.gen(function* () {
       const session = yield* setup
       const session = yield* setup
       const registry = yield* ToolRegistry.Service
       const registry = yield* ToolRegistry.Service
-      yield* registry.register({
-        permissionfail: Tool.make({
-          description: "Reject a permission",
-          input: Schema.Struct({}),
-          output: Schema.Struct({}),
-          execute: () =>
-            new ToolFailure({
-              message: "Permission denied: edit",
-              error: new PermissionV2.BlockedError({
-                rules: [],
-                permission: "edit",
-                resources: ["src/index.ts"],
-              }),
-            }),
-        }),
-      })
+      yield* registry.register({ permissionfail: permissionFail })
       yield* admit(session, "Reject permission")
       yield* admit(session, "Reject permission")
       responses = [
       responses = [
         reply.tool("call-permission", "permissionfail", {}),
         reply.tool("call-permission", "permissionfail", {}),
@@ -3034,6 +3049,90 @@ describe("SessionRunnerLLM", () => {
     }),
     }),
   )
   )
 
 
+  const rejectPermissionWhileToolInputStreams = (lateEvent: LLMEvent) =>
+    Effect.gen(function* () {
+      const session = yield* setup
+      const registry = yield* ToolRegistry.Service
+      const releaseLateEvent = yield* Deferred.make<void>()
+      yield* registry.register({ permissionfail: permissionFail })
+      const events = yield* EventV2.Service
+      const permissionFailed = yield* events
+        .subscribe(SessionEvent.Tool.Failed)
+        .pipe(
+          Stream.filter((event) => event.data.sessionID === sessionID && event.data.callID === "call-permission"),
+          Stream.runHead,
+          Effect.forkScoped({ startImmediately: true }),
+        )
+      yield* admit(session, "Reject permission while another tool input streams")
+      responseStream = Stream.concat(
+        Stream.fromIterable([
+          LLMEvent.stepStart({ index: 0 }),
+          LLMEvent.toolInputStart({ id: "call-streaming", name: "echo" }),
+          LLMEvent.toolCall({ id: "call-permission", name: "permissionfail", input: {} }),
+        ]),
+        Stream.fromEffect(Deferred.await(releaseLateEvent)).pipe(Stream.flatMap(() => Stream.make(lateEvent))),
+      )
+
+      const run = yield* session.resume(sessionID).pipe(Effect.forkChild)
+      yield* Fiber.join(permissionFailed).pipe(Effect.timeout("1 second"))
+      yield* Effect.yieldNow
+      const inputEndsBeforeRelease = yield* recordedToolInputEnds(sessionID, "call-streaming")
+      yield* Deferred.succeed(releaseLateEvent, undefined)
+      const exit = yield* Fiber.await(run)
+      return {
+        exit,
+        inputEndsBeforeRelease,
+        context: yield* session.context(sessionID),
+        inputEnds: yield* recordedToolInputEnds(sessionID, "call-streaming"),
+      }
+    })
+
+  for (const testCase of [
+    {
+      name: "does not end concurrent tool input when permission is rejected",
+      event: LLMEvent.toolInputDelta({
+        id: "call-streaming",
+        name: "echo",
+        text: '{"text":"still streaming"}',
+      }),
+      defect: "Tool input delta after end: call-streaming",
+    },
+    {
+      name: "does not duplicate concurrent tool input end when permission is rejected",
+      event: LLMEvent.toolInputEnd({ id: "call-streaming", name: "echo" }),
+      defect: "Duplicate tool input end: call-streaming",
+    },
+  ]) {
+    it.effect(testCase.name, () =>
+      Effect.gen(function* () {
+        const result = yield* rejectPermissionWhileToolInputStreams(testCase.event)
+
+        expect(result.inputEndsBeforeRelease).toHaveLength(0)
+        expect(Exit.isFailure(result.exit)).toBe(true)
+        if (Exit.isFailure(result.exit)) {
+          expect(Cause.pretty(result.exit.cause)).not.toContain(testCase.defect)
+          expect(Cause.hasDies(result.exit.cause)).toBe(false)
+        }
+        expect(result.context).toMatchObject([
+          { type: "user" },
+          {
+            type: "assistant",
+            error: { type: "permission.rejected", message: "Permission denied: edit" },
+            content: [
+              {
+                type: "tool",
+                id: "call-streaming",
+                state: { status: "error", error: { type: "aborted", message: "Tool execution interrupted" } },
+              },
+              { type: "tool", id: "call-permission", state: { status: "error" } },
+            ],
+          },
+        ])
+        expect(result.inputEnds).toHaveLength(1)
+      }),
+    )
+  }
+
   it.effect("interrupts runner continuation when a question is cancelled", () =>
   it.effect("interrupts runner continuation when a question is cancelled", () =>
     Effect.gen(function* () {
     Effect.gen(function* () {
       const session = yield* setup
       const session = yield* setup