Skip to content

refactor(session): use onInterrupt finalizer for cancelled tool output#21751

Closed
kitlangton wants to merge 1 commit intodevfrom
kit/effect-native-tool-interrupt
Closed

refactor(session): use onInterrupt finalizer for cancelled tool output#21751
kitlangton wants to merge 1 commit intodevfrom
kit/effect-native-tool-interrupt

Conversation

@kitlangton
Copy link
Copy Markdown
Contributor

@kitlangton kitlangton commented Apr 9, 2026

Summary

Follow-up to #21724. Replaces the tool body's imperative if (options.abortSignal?.aborted) completeToolCall(...) tail-check with structural interruption handling via Effect.onInterrupt, then extracts the whole pattern into a small runToolExecute helper so the two call sites (built-in and MCP) read as declarative specs.

The core question we were poking at: the AI SDK hands us a detached Effect.runPromise orphan fiber — but can Effect's interruption machinery still handle cancel cleanly if we wire it up right? Answer: yes. This PR shows what that looks like.

The helper

function runToolExecute<Raw, Output>(options: {
  signal: AbortSignal | undefined
  before: Effect.Effect<unknown, any, any>
  execute: () => Promise<Raw>
  finalize: (result: Raw) => Effect.Effect<Output, any, any>
  onCancel: (output: Output) => Effect.Effect<unknown, any, any>
}): Promise<Output>

Four callbacks, four phases:

  • before — plugin hooks, permission asks, anything pre-execute
  • execute — the native Promise-returning call. Reference held outside Effect so a finalizer can re-await it.
  • finalize — builds the tool output from the raw result. Runs on both the happy path and the interrupt path.
  • onCancel — the side channel into the processor, only called if interrupted mid-flight.

Under the hood:

  • runPromise(..., { signal: options.abortSignal }) wires the SDK abort signal directly into fiber interrupt — the tool fiber is no longer orphan.
  • Effect.onInterrupt finalizer re-awaits the same in-flight pending uninterruptibly (finalizers always are), runs finalize again, posts via onCancel, stashes the output in a rescued closure var.
  • Effect.catchCause at the tail converts interrupt-with-rescued-output back into success, so the Promise returned to the SDK resolves with the finalized output instead of rejecting. SDK reports the tool as successful, not tool-error.
  • attach() (exported from @/effect/run-service) provides InstanceRef to the fiber so onInterrupt finalizers — which run outside the original ALS chain — can still resolve the instance context through the ServiceMap. Every makeRuntime-created runner already uses attach internally; now the tool fibers do too.

Built-in tool call site (the whole thing)

execute(args, options) {
  const ctx = context(args, options)
  const meta = { tool: item.id, sessionID: ctx.sessionID, callID: ctx.callID }
  return runToolExecute({
    signal: options.abortSignal,
    before: plugin.trigger("tool.execute.before", meta, { args }),
    execute: () => item.execute(args, ctx),
    finalize: (result) =>
      Effect.gen(function* () {
        const output = {
          ...result,
          attachments: result.attachments?.map((attachment) => ({
            ...attachment,
            id: PartID.ascending(),
            sessionID: ctx.sessionID,
            messageID: input.processor.message.id,
          })),
        }
        yield* plugin.trigger("tool.execute.after", { ...meta, args }, output)
        return output
      }),
    onCancel: (output) => input.processor.completeToolCall(options.toolCallId, output),
  })
}

No pending tracking, no captured tracking, no runPromiseExit + Exit.isSuccess / Cause.squash dance, no ALS gotcha at the call site, no imperative if (aborted) branch. Reads as "here's what to do before; here's the Promise; here's how to build the output; here's the cancel fallback."

Gotcha: AsyncLocalStorage loss inside onInterrupt finalizers

This took a while to track down — documenting it so the next person doesn't repeat it.

The original code worked because the tool body ran synchronously down the Promise chain from inside Instance.provide's AsyncLocalStorage.run call. Node preserves ALS through Promise microtasks, so Instance.current was always live all the way down, including inside bus-publish subscribers.

With onInterrupt, the finalizer is invoked from the interrupt propagation path — specifically, the signal.addEventListener("abort", () => fiber.interruptUnsafe()) callback that runPromise registers. That listener fires outside the original ALS chain. When the finalizer's downstream code (bus.publish → sync subscriber in sync/index.ts:155result.then(...)Database.transaction) eventually calls Instance.current, ALS is empty and it throws NotFound.

Fix: the existing attach helper from @/effect/run-service (previously private, now exported) captures Instance.current synchronously and provides it as InstanceRef. InstanceState.context already prefers InstanceRef over the ALS fallback, so everything resolves through ServiceMap. Every makeRuntime runner in the codebase already goes through attach; the tool fiber was the odd one out.

Debug path that led here:

  1. Finalizer runs, completeToolCall completes, part is marked completed
  2. But the outer process() fiber's exit is Failure, test asserts Exit.isSuccess and fails
  3. Stack trace → instance-state.ts:28 via bus/index.ts:172sync/index.ts:155
  4. That's a .then() in the bus publish sync subscriber, running in a microtask whose parent chain originated from the interrupt listener rather than the original execute() call — ALS dropped

What this buys you

  • Structural interrupt handling instead of an imperative tail-check.
  • Tool fibers are no longer orphans — cancel flows through Effect's interruption, not a data bit the body happens to observe.
  • Each tool's execute is a declarative spec. No knowledge of completeToolCall or runPromise plumbing at the call site.
  • Future tool additions get the cancel-finalize behavior for free by using the helper.
  • Consistent with how the rest of the codebase bridges ALS → ServiceMap via attach.

What this doesn't buy you

  • The Deferred in processor.ts is still required. The processor-cleanup-loop vs tool-finalizer race is a separate concern. Eliminating it needs making tool fibers scoped children of the processor fiber so its cleanup naturally awaits them. Out of scope here.
  • Slightly larger diff than fix: finalize interrupted bash via tool result path #21724 because of the helper (+134 / -93 vs what's in dev). The call sites are smaller; the helper makes up the difference. If you'd rather not extract the helper, the inline version is roughly the same size as the current PR.

Draft status / open questions

  1. Helper naming and location. Currently module-private runToolExecute at the top of prompt.ts. Could live in a separate file (src/session/tool-execute.ts) if we expect to grow it.
  2. any types on the callback Effect generics. Using Effect.Effect<unknown, any, any> for before / finalize / onCancel to sidestep E and R unification. Could tighten with proper generics but gets noisy; all call sites have R = never (services are closure-captured) so it doesn't matter at runtime.
  3. Worth dropping the Deferred in processor.ts as a follow-up by restructuring tool fibers to be scoped children of the processor fiber? Separate PR if so.

Testing

  • bun typecheck clean
  • bun run test --timeout 60000 test/session/prompt-effect.test.ts -t "cancel finalizes interrupted bash tool output through normal truncation" passes
  • bun run test --timeout 60000 test/session/ — 241 pass / 4 skip / 1 todo / 0 fail

@kitlangton kitlangton force-pushed the kit/effect-native-tool-interrupt branch from c1f0673 to f37b4b5 Compare April 9, 2026 20:19
Wire the AI SDK's abortSignal into the tool fiber via runPromiseExit's
signal option so interruption is first-class, and move the "finalize on
cancel" path into an Effect.onInterrupt finalizer that re-awaits the
still-running native Promise uninterruptibly, builds the output, and
posts it through completeToolCall.

Replaces the imperative `if (options.abortSignal?.aborted)` tail check
with structural interruption handling. When the fiber is interrupted,
the finalizer captures the truncated bash output (or MCP tool result)
and the .then on runPromiseExit resolves the SDK's Promise with the
captured value instead of propagating the interrupt cause as a
rejection, so the tool is reported as successfully completed rather
than as a tool-error.

InstanceRef is provided on the tool fiber so InstanceState.context
resolves through ServiceMap rather than falling through to the
AsyncLocalStorage, which the onInterrupt finalizer runs outside of.
@kitlangton kitlangton force-pushed the kit/effect-native-tool-interrupt branch from f37b4b5 to ea19ee7 Compare April 9, 2026 20:41
@kitlangton kitlangton closed this Apr 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant