refactor(core): use tagged constructors and exhaustive matching

packages/core/src/effect/state-machine.ts

@@ -0,0 +1,236 @@
+export * as StateMachine from "./state-machine.js"
+
+import { Cause, Effect, Exit, Fiber, Queue, type Scope } from "effect"
+
+export type Command<Operation> =
+  | {
+      readonly _tag: "Invoke"
+      readonly id: string
+      readonly operation: Operation
+    }
+  | {
+      readonly _tag: "Stop"
+      readonly id: string
+    }
+  | {
+      readonly _tag: "StopAndJoin"
+      readonly id: string
+      readonly ids: ReadonlyArray<string>
+      readonly waitFor: ReadonlyArray<string>
+    }
+
+export type InvocationExited<Event, Operation, Error> = {
+  readonly _tag: "InvocationExited"
+  readonly id: string
+  readonly generation: number
+  readonly operation: Operation
+  readonly exit: Exit.Exit<Event, Error>
+}
+
+export type RuntimeEvent<Event, Operation, Error> =
+  | {
+      readonly _tag: "Input"
+      readonly input: Event
+      readonly cause?: Cause.Cause<never>
+    }
+  | InvocationExited<Event, Operation, Error>
+  | {
+      readonly _tag: "InvocationsStopped"
+      readonly id: string
+      readonly exits: ReadonlyArray<InvocationExited<Event, Operation, Error>>
+    }
+
+export type Continue<State, Operation> = {
+  readonly _tag: "Continue"
+  readonly state: State
+  readonly commands: ReadonlyArray<Command<Operation>>
+}
+
+export type Decision<State, Operation, Output> =
+  | Continue<State, Operation>
+  | {
+      readonly _tag: "Done"
+      readonly output: Output
+    }
+
+export type Definition<State, Event, Operation, Error, Output> = {
+  readonly initial: Continue<State, Operation>
+  readonly transition: (
+    state: State,
+    event: RuntimeEvent<Event, Operation, Error>,
+  ) => Decision<State, Operation, Output>
+  readonly interruption?: Event
+}
+
+export type Executor<Event, Operation, Error, Requirements> = (
+  operation: Operation,
+) => Effect.Effect<Event, Error, Requirements>
+
+export function define<State, Event, Operation, Error, Output>(
+  definition: Definition<State, Event, Operation, Error, Output>,
+) {
+  return definition
+}
+
+export function next<State, Operation = never>(state: State, ...commands: ReadonlyArray<Command<Operation>>) {
+  return { _tag: "Continue", state, commands } as const
+}
+
+export function done<Output>(output: Output) {
+  return { _tag: "Done", output } as const
+}
+
+export function invoke<Operation>(id: string, operation: Operation): Command<Operation> {
+  return { _tag: "Invoke", id, operation }
+}
+
+export function stop(id: string): Command<never> {
+  return { _tag: "Stop", id }
+}
+
+/** Stops `ids`, awaits `waitFor` without interruption, and delivers their exits as one batch. */
+export function stopAndJoin(
+  id: string,
+  ids: ReadonlyArray<string>,
+  waitFor: ReadonlyArray<string> = [],
+): Command<never> {
+  return { _tag: "StopAndJoin", id, ids, waitFor }
+}
+
+export const run = Effect.fn("StateMachine.run")(function* <State, Event, Operation, Error, Output, Requirements>(
+  definition: Definition<State, Event, Operation, Error, Output>,
+  execute: Executor<Event, Operation, Error, Requirements>,
+) {
+  return yield* Effect.uninterruptibleMask((restore) =>
+    Effect.scoped(
+      Effect.gen(function* () {
+        const queue = yield* Queue.unbounded<RuntimeEvent<Event, Operation, Error>>()
+        const invocations = new Map<
+          string,
+          {
+            readonly generation: number
+            readonly operation: Operation
+            readonly fiber: Fiber.Fiber<Event, Error>
+          }
+        >()
+        let generation = 0
+
+        const executeCommands = Effect.fnUntraced(function* (
+          commands: ReadonlyArray<Command<Operation>>,
+          interruptibleExecution: boolean,
+        ) {
+          yield* Effect.forEach(
+            commands,
+            (command) =>
+              Effect.gen(function* () {
+                if (command._tag === "Stop") {
+                  const invocation = invocations.get(command.id)
+                  yield* invocation
+                    ? Fiber.interrupt(invocation.fiber)
+                    : Effect.die(new Error(`Unknown state machine invocation: ${command.id}`))
+                  return
+                }
+
+                if (command._tag === "StopAndJoin") {
+                  const captured = [...command.ids, ...command.waitFor].flatMap((id) => {
+                    const invocation = invocations.get(id)
+                    return invocation ? [{ id, ...invocation }] : []
+                  })
+                  if (captured.length !== command.ids.length + command.waitFor.length)
+                    yield* Effect.die(new Error("Unknown state machine invocation in StopAndJoin"))
+
+                  // Invalidate individual exits, including ones already queued, before interrupting.
+                  captured.forEach((invocation) => invocations.delete(invocation.id))
+                  yield* Fiber.interruptAll(captured.slice(0, command.ids.length).map((invocation) => invocation.fiber))
+                  const exits = yield* Effect.forEach(captured, (invocation) =>
+                    Fiber.await(invocation.fiber).pipe(
+                      Effect.map((exit) => ({
+                        _tag: "InvocationExited" as const,
+                        id: invocation.id,
+                        generation: invocation.generation,
+                        operation: invocation.operation,
+                        exit,
+                      })),
+                    ),
+                  )
+                  yield* Queue.offer(queue, { _tag: "InvocationsStopped", id: command.id, exits })
+                  return
+                }
+
+                const previous = invocations.get(command.id)
+                if (previous) yield* Fiber.interrupt(previous.fiber)
+
+                generation += 1
+                const current = generation
+                const execution = interruptibleExecution
+                  ? restore(execute(command.operation))
+                  : execute(command.operation)
+                const fiber = yield* execution.pipe(Effect.forkScoped({ startImmediately: false }))
+                invocations.set(command.id, { generation: current, operation: command.operation, fiber })
+                // A deferred child may be interrupted before an Effect.onExit observer starts.
+                fiber.addObserver((exit) => {
+                  Queue.offerUnsafe(queue, {
+                    _tag: "InvocationExited",
+                    id: command.id,
+                    generation: current,
+                    operation: command.operation,
+                    exit,
+                  })
+                })
+              }),
+            { discard: true },
+          )
+        })
+
+        const handleInterruption = (
+          state: State,
+          cause: Cause.Cause<never>,
+        ): Effect.Effect<Output, never, Requirements | Scope.Scope> =>
+          Effect.gen(function* () {
+            if (!Cause.hasInterruptsOnly(cause) || definition.interruption === undefined)
+              return yield* Effect.failCause(cause)
+            return yield* dispatch(
+              definition.transition(state, {
+                _tag: "Input",
+                input: definition.interruption,
+                cause,
+              }),
+              true,
+            )
+          })
+
+        const dispatch = (
+          decision: Decision<State, Operation, Output>,
+          interrupted: boolean,
+        ): Effect.Effect<Output, never, Requirements | Scope.Scope> =>
+          Effect.gen(function* () {
+            if (decision._tag === "Done") return decision.output
+            yield* executeCommands(decision.commands, !interrupted)
+            if (interrupted) return yield* Effect.suspend(() => loop(decision.state, true))
+
+            const boundary = yield* restore(Effect.void).pipe(Effect.exit)
+            if (Exit.isFailure(boundary)) return yield* handleInterruption(decision.state, boundary.cause)
+            return yield* Effect.suspend(() => loop(decision.state, false))
+          })
+
+        const loop = (state: State, interrupted: boolean): Effect.Effect<Output, never, Requirements | Scope.Scope> =>
+          Effect.gen(function* () {
+            const received = yield* (interrupted ? Queue.take(queue) : restore(Queue.take(queue))).pipe(Effect.exit)
+            if (Exit.isFailure(received)) return yield* handleInterruption(state, received.cause)
+
+            if (received.value._tag === "InvocationExited") {
+              const invocation = invocations.get(received.value.id)
+              if (!invocation || invocation.generation !== received.value.generation) {
+                return yield* Effect.suspend(() => loop(state, interrupted))
+              }
+              invocations.delete(received.value.id)
+            }
+
+            return yield* dispatch(definition.transition(state, received.value), interrupted)
+          })
+
+        return yield* dispatch(definition.initial, false)
+      }),
+    ),
+  )
+})

packages/core/src/session/runner/llm.ts

@@ -22,6 +22,7 @@ import { llmClient } from "../../effect/app-node-platform.js"
 import { StepFailedError } from "../error.js"
 import { SessionRunnerRetry } from "./retry.js"
 import { SessionStep } from "./step.js"
+import { SessionStepMachine } from "./step-machine.js"
 import { ToolOutput } from "../../tool-output.js"
 import { PluginSupervisor } from "../../plugin/supervisor.js"
 import { MAX_STEPS_PROMPT } from "./max-steps.js"
@@ -167,91 +168,83 @@ const layer = Layer.effect(
       return selected
     })
 
-    /** Owns logical Step policy; each attempt owns its streaming, tools, and durable settlement. */
+    /** Owns logical Step policy; each attempt owns provider observation, tools, and durable settlement. */
     const runStep = Effect.fn("SessionRunner.runStep")(function* (first: SessionContext.Loaded, step: number) {
       const sessionID = first.session.id
-      let assistantMessageID = SessionMessage.ID.create()
       const retry = yield* Schedule.toStepWithSleep(SessionRunnerRetry.schedule(bus, sessionID))
       let initial: SessionContext.Loaded | undefined = first
-      let recoverOverflow = true
-      let recoverContinuation = true
-      while (true) {
-        // Reuse boundary preparation once; retries refresh context without delivering more input.
-        const loaded = initial ?? (yield* prepareContext(sessionID).pipe(Effect.flatMap(context.load)))
-        initial = undefined
-        const compactionInput = {
-          session: loaded.session,
-          messages: loaded.messages,
-          resolved: loaded.model,
-          prepare: context.prepare,
-        }
-        if (compaction.required(compactionInput)) {
-          const compacted = yield* compaction.compact(compactionInput)
-          if (compacted.status !== "completed") return yield* new StepFailedError({ error: compacted.error })
-          assistantMessageID = SessionMessage.ID.create()
-          continue
-        }
-        const stepLimitReached = loaded.agent.info.steps !== undefined && step >= loaded.agent.info.steps
-        const transcript = SessionModelRequest.baseTranscript({
-          agent: loaded.agent.info,
-          model: loaded.model,
-          tools: loaded.tools,
-          initial: loaded.initial,
-          messages: loaded.messages,
-        })
-        const prepared = yield* context.prepare({
-          scope: { session: loaded.session, agentID: loaded.agent.id, model: loaded.model, tools: loaded.tools },
-          transcript: {
-            system: transcript.system,
-            messages: stepLimitReached
-              ? [...transcript.messages, Message.assistant(MAX_STEPS_PROMPT)]
-              : transcript.messages,
-          },
-          // Keep tool definitions on the final Step to preserve the provider's cached prefix.
-          toolChoice: stepLimitReached ? "none" : undefined,
-          webSocket: "session",
-        })
-        const outcome = yield* steps.attempt({
-          sessionID,
-          assistantMessageID,
-          agent: loaded.agent.id,
-          model: loaded.model,
-          prepared,
-          recoverContinuation,
-          recoverOverflow: Effect.suspend(() =>
-            recoverOverflow && compaction.enabled()
-              ? compaction.compact(compactionInput).pipe(Effect.map((result) => result.status === "completed"))
-              : Effect.succeed(false),
-          ),
-        })
-        const completed = yield* SessionStep.Outcome.$match(outcome, {
-          Completed: (outcome) => Effect.succeed(outcome.needsContinuation),
-          Retry: (outcome) =>
-            retry({ cause: outcome.cause, error: outcome.error, assistantMessageID }).pipe(
-              Pull.catchDone(() =>
-                bus
-                  .publish(SessionEvent.Step.Failed, { sessionID, assistantMessageID, error: outcome.error })
-                  .pipe(Effect.andThen(outcome.cause)),
+      return yield* SessionStepMachine.run(SessionMessage.ID.create(), {
+        prepare: Effect.fnUntraced(function* (state) {
+          // Reuse boundary preparation once; retries refresh context without delivering more input.
+          const loaded = initial ?? (yield* prepareContext(sessionID).pipe(Effect.flatMap(context.load)))
+          initial = undefined
+          const compactionInput = {
+            session: loaded.session,
+            messages: loaded.messages,
+            resolved: loaded.model,
+            prepare: context.prepare,
+          }
+          if (compaction.required(compactionInput)) {
+            const compacted = yield* compaction.compact(compactionInput)
+            if (compacted.status !== "completed") return yield* new StepFailedError({ error: compacted.error })
+            return SessionStepMachine.Preparation.Rebuilt()
+          }
+          const stepLimitReached = loaded.agent.info.steps !== undefined && step >= loaded.agent.info.steps
+          const transcript = SessionModelRequest.baseTranscript({
+            agent: loaded.agent.info,
+            model: loaded.model,
+            tools: loaded.tools,
+            initial: loaded.initial,
+            messages: loaded.messages,
+          })
+          const prepared = yield* context.prepare({
+            scope: { session: loaded.session, agentID: loaded.agent.id, model: loaded.model, tools: loaded.tools },
+            transcript: {
+              system: transcript.system,
+              messages: stepLimitReached
+                ? [...transcript.messages, Message.assistant(MAX_STEPS_PROMPT)]
+                : transcript.messages,
+            },
+            // Keep tool definitions on the final Step to preserve the provider's cached prefix.
+            toolChoice: stepLimitReached ? "none" : undefined,
+            webSocket: "session",
+          })
+          return SessionStepMachine.Preparation.Ready({
+            attempt: yield* steps.open({
+              sessionID,
+              assistantMessageID: state.assistantMessageID,
+              agent: loaded.agent.id,
+              model: loaded.model,
+              prepared,
+              recoverContinuation: state.recoverContinuation,
+              recoverOverflow: Effect.suspend(() =>
+                compaction.enabled()
+                  ? compaction.compact(compactionInput).pipe(Effect.map((result) => result.status === "completed"))
+                  : Effect.succeed(false),
               ),
-              Effect.asVoid,
+            }),
+          })
+        }),
+        retry: (state, outcome) =>
+          retry({ cause: outcome.cause, error: outcome.error, assistantMessageID: state.assistantMessageID }).pipe(
+            Pull.catchDone(() =>
+              outcome._tag === "Retry"
+                ? bus
+                    .publish(SessionEvent.Step.Failed, {
+                      sessionID,
+                      assistantMessageID: state.assistantMessageID,
+                      error: outcome.error,
+                    })
+                    .pipe(Effect.andThen(outcome.cause))
+                : outcome.cause,
             ),
-          Continue: Effect.fnUntraced(function* (outcome) {
-            yield* retry({ cause: outcome.cause, error: outcome.error, assistantMessageID }).pipe(
-              Pull.catchDone(() => outcome.cause),
-            )
-            yield* bus.publish(SessionEvent.Synthetic, { sessionID, text: CONTINUE_AFTER_INCOMPLETE_STREAM })
-            assistantMessageID = SessionMessage.ID.create()
-          }),
-          Compacted: Effect.fnUntraced(function* () {
-            recoverOverflow = false
-            assistantMessageID = SessionMessage.ID.create()
-          }),
-          RecoverFull: Effect.fnUntraced(function* () {
-            recoverContinuation = false
-          }),
-        })
-        if (completed !== undefined) return completed
-      }
+            Effect.asVoid,
+          ),
+        publishSynthetic: bus.publish(SessionEvent.Synthetic, {
+          sessionID,
+          text: CONTINUE_AFTER_INCOMPLETE_STREAM,
+        }),
+      })
     })
 
     const settleStaleToolCalls = Effect.fn("SessionRunner.settleStaleToolCalls")(function* (

packages/core/src/session/runner/step-machine.ts

@@ -0,0 +1,402 @@
+export * as SessionStepMachine from "./step-machine.js"
+
+import { AIError, type ToolCall } from "@opencode-ai/ai"
+import { Cause, Data, Effect, Exit } from "effect"
+import { StateMachine } from "../../effect/state-machine.js"
+import { StepFailedError } from "../error.js"
+import { SessionMessage } from "../message.js"
+import { SessionStep } from "./step.js"
+
+const PREPARATION = "preparation"
+const PROVIDER = "provider"
+const COMPACTION = "compaction"
+const SETTLEMENT = "settlement"
+const RETRY = "retry"
+
+export type Context = {
+  readonly assistantMessageID: SessionMessage.ID
+  readonly recoverOverflow: boolean
+  readonly recoverContinuation: boolean
+}
+
+export type Preparation = Data.TaggedEnum<{
+  Rebuilt: {}
+  Ready: { readonly attempt: SessionStep.Attempt }
+}>
+export const Preparation = Data.taggedEnum<Preparation>()
+
+type AttemptFailure = AIError | StepFailedError
+type BackoffOutcome = Data.TaggedEnum.Value<SessionStep.Outcome, "Retry" | "Continue">
+
+type ToolRun = {
+  readonly call: ToolCall
+  readonly exit?: SessionStep.ToolExit
+}
+
+type ActiveAttempt = {
+  readonly context: Context
+  readonly attempt: SessionStep.Attempt
+  readonly tools: ReadonlyMap<string, ToolRun>
+}
+
+type AttemptState = Data.TaggedEnum<{
+  ObservingProvider: { readonly active: ActiveAttempt }
+  FinalizingProvider: {
+    readonly active: ActiveAttempt
+    readonly stream: Exit.Exit<void, AIError>
+    readonly stopping?: Cause.Cause<never>
+  }
+  AwaitingTools: { readonly active: ActiveAttempt; readonly stream: Exit.Exit<void, AIError> }
+  RecoveringOverflow: { readonly active: ActiveAttempt; readonly stream: Exit.Exit<void, AIError> }
+}>
+
+export type State =
+  | AttemptState
+  | Data.TaggedEnum<{
+      PreparingAttempt: { readonly context: Context }
+      SettlingAttempt: { readonly active: ActiveAttempt; readonly stopping?: Cause.Cause<never> }
+      BackingOff: {
+        readonly context: Context
+        readonly outcome: BackoffOutcome
+      }
+      Stopping: { readonly from?: AttemptState; readonly cause: Cause.Cause<never> }
+    }>
+export const State = Data.taggedEnum<State>()
+
+export type Event<Failure> = Data.TaggedEnum<{
+  Prepared: { readonly exit: Exit.Exit<{ readonly context: Context; readonly preparation: Preparation }, Failure> }
+  ProviderObserved: { readonly exit: Exit.Exit<SessionStep.ProviderObservation, AIError> }
+  ToolFinished: { readonly call: ToolCall; readonly exit: SessionStep.ToolExit }
+  ProviderFinished: { readonly exit: Exit.Exit<void> }
+  OverflowRecovered: { readonly exit: Exit.Exit<boolean> }
+  AttemptSettled: { readonly exit: Exit.Exit<SessionStep.Outcome, AttemptFailure> }
+  RetryFinished: { readonly exit: Exit.Exit<void, Failure> }
+  CancelRequested: {}
+}>
+interface EventDefinition extends Data.TaggedEnum.WithGenerics<1> {
+  readonly taggedEnum: Event<this["A"]>
+}
+export const Event = Data.taggedEnum<EventDefinition>()
+
+export type Operation = Data.TaggedEnum<{
+  PrepareAttempt: { readonly context: Context; readonly freshAssistant: boolean }
+  ObserveProvider: { readonly attempt: SessionStep.Attempt }
+  RunTool: { readonly attempt: SessionStep.Attempt; readonly call: ToolCall }
+  FinishProvider: { readonly attempt: SessionStep.Attempt; readonly stream: Exit.Exit<void, AIError> }
+  RecoverOverflow: { readonly attempt: SessionStep.Attempt; readonly settlement: SessionStep.Settlement }
+  SettleAttempt: { readonly attempt: SessionStep.Attempt; readonly settlement: SessionStep.Settlement }
+  Retry: {
+    readonly context: Context
+    readonly outcome: BackoffOutcome
+  }
+}>
+export const Operation = Data.taggedEnum<Operation>()
+
+export type Capabilities<Failure, RetryFailure, Requirements> = {
+  readonly prepare: (context: Context) => Effect.Effect<Preparation, Failure, Requirements>
+  readonly retry: (context: Context, outcome: BackoffOutcome) => Effect.Effect<void, RetryFailure, Requirements>
+  readonly publishSynthetic: Effect.Effect<void, Failure, Requirements>
+}
+
+export const run = Effect.fn("SessionStepMachine.run")(function* <Failure, RetryFailure, Requirements>(
+  assistantMessageID: SessionMessage.ID,
+  capabilities: Capabilities<Failure, RetryFailure, Requirements>,
+) {
+  const execute = Operation.$match({
+    PrepareAttempt: (operation) =>
+      Effect.suspend(() => {
+        const context = operation.freshAssistant
+          ? { ...operation.context, assistantMessageID: SessionMessage.ID.create() }
+          : operation.context
+        return capabilities.prepare(context).pipe(Effect.map((preparation) => ({ context, preparation })))
+      }).pipe(
+        Effect.exit,
+        Effect.map((exit) => Event.Prepared({ exit })),
+      ),
+    ObserveProvider: (operation) =>
+      operation.attempt.observeUntilBoundary().pipe(
+        Effect.exit,
+        Effect.map((exit) => Event.ProviderObserved({ exit })),
+      ),
+    RunTool: (operation) =>
+      operation.attempt.runTool(operation.call).pipe(
+        Effect.exit,
+        Effect.map((exit) => Event.ToolFinished({ call: operation.call, exit })),
+      ),
+    FinishProvider: (operation) =>
+      operation.attempt.finishProvider(operation.stream).pipe(
+        Effect.exit,
+        Effect.map((exit) => Event.ProviderFinished({ exit })),
+      ),
+    RecoverOverflow: (operation) =>
+      operation.attempt.recoverOverflow(operation.settlement).pipe(
+        Effect.exit,
+        Effect.map((exit) => Event.OverflowRecovered({ exit })),
+      ),
+    SettleAttempt: (operation) =>
+      operation.attempt.settle(operation.settlement).pipe(
+        Effect.exit,
+        Effect.map((exit) => Event.AttemptSettled({ exit })),
+      ),
+    Retry: (operation) =>
+      capabilities.retry(operation.context, operation.outcome).pipe(
+        Effect.andThen(operation.outcome._tag === "Continue" ? capabilities.publishSynthetic : Effect.void),
+        Effect.exit,
+        Effect.map((exit) => Event.RetryFinished({ exit })),
+      ),
+  })
+  const result = yield* StateMachine.run(definition<Failure, RetryFailure>(assistantMessageID), execute)
+  return yield* result
+})
+
+export const definition = <Failure, RetryFailure>(assistantMessageID: SessionMessage.ID) => {
+  const context = {
+    assistantMessageID,
+    recoverOverflow: true,
+    recoverContinuation: true,
+  }
+  type MachineFailure = Failure | RetryFailure | AttemptFailure
+  type Decision = StateMachine.Decision<State, Operation, Exit.Exit<boolean, MachineFailure>>
+
+  const prepare = (context: Context, freshAssistant = false): StateMachine.Continue<State, Operation> =>
+    StateMachine.next(
+      State.PreparingAttempt({ context }),
+      StateMachine.invoke(PREPARATION, Operation.PrepareAttempt({ context, freshAssistant })),
+    )
+
+  const pull = (active: ActiveAttempt): Decision =>
+    StateMachine.next(
+      State.ObservingProvider({ active }),
+      StateMachine.invoke(PROVIDER, Operation.ObserveProvider({ attempt: active.attempt })),
+    )
+
+  const settlement = (active: ActiveAttempt, stream: Exit.Exit<void, AIError>): SessionStep.Settlement => ({
+    stream,
+    tools: Array.from(active.tools.values()).flatMap((tool) =>
+      tool.exit ? [{ call: tool.call, exit: tool.exit }] : [],
+    ),
+  })
+
+  const settle = (active: ActiveAttempt, stream: Exit.Exit<void, AIError>, stopping?: Cause.Cause<never>): Decision =>
+    StateMachine.next(
+      State.SettlingAttempt({ active, stopping }),
+      StateMachine.invoke(
+        SETTLEMENT,
+        Operation.SettleAttempt({
+          attempt: active.attempt,
+          settlement: settlement(active, stream),
+        }),
+      ),
+    )
+
+  const afterProvider = (active: ActiveAttempt, stream: Exit.Exit<void, AIError>): Decision => {
+    if (Array.from(active.tools.values()).some((tool) => tool.exit === undefined))
+      return StateMachine.next(State.AwaitingTools({ active, stream }))
+    if (!active.context.recoverOverflow) return settle(active, stream)
+    return StateMachine.next(
+      State.RecoveringOverflow({ active, stream }),
+      StateMachine.invoke(
+        COMPACTION,
+        Operation.RecoverOverflow({
+          attempt: active.attempt,
+          settlement: settlement(active, stream),
+        }),
+      ),
+    )
+  }
+
+  const finishProvider = (
+    active: ActiveAttempt,
+    stream: Exit.Exit<void, AIError>,
+    stopping?: Cause.Cause<never>,
+  ): Decision =>
+    StateMachine.next(
+      State.FinalizingProvider({ active, stream, stopping }),
+      StateMachine.invoke(
+        PROVIDER,
+        Operation.FinishProvider({
+          attempt: active.attempt,
+          stream,
+        }),
+      ),
+    )
+
+  const stop = (cause: Cause.Cause<never>, ids: ReadonlyArray<string>, from?: AttemptState): Decision => {
+    return StateMachine.next(
+      State.Stopping({ cause, from }),
+      StateMachine.stopAndJoin("step", ids, from?._tag === "FinalizingProvider" ? [PROVIDER] : []),
+    )
+  }
+
+  const interrupt = (state: State, cause: Cause.Cause<never>): Decision => {
+    const stopAttempt = (state: Exclude<AttemptState, { readonly _tag: "RecoveringOverflow" }>) =>
+      stop(
+        cause,
+        [
+          ...(state._tag === "ObservingProvider" ? [PROVIDER] : []),
+          ...Array.from(state.active.tools.values()).flatMap((tool) =>
+            tool.exit === undefined ? [toolID(tool.call)] : [],
+          ),
+        ],
+        state,
+      )
+    return State.$match(state, {
+      PreparingAttempt: () => stop(cause, [PREPARATION]),
+      ObservingProvider: stopAttempt,
+      FinalizingProvider: stopAttempt,
+      AwaitingTools: stopAttempt,
+      SettlingAttempt: (state) => StateMachine.next(State.SettlingAttempt({ active: state.active, stopping: cause })),
+      RecoveringOverflow: (state) => stop(cause, [COMPACTION], state),
+      BackingOff: () => stop(cause, [RETRY]),
+      Stopping: (state) => StateMachine.next(state),
+    })
+  }
+
+  return StateMachine.define<
+    State,
+    Event<Failure | RetryFailure>,
+    Operation,
+    never,
+    Exit.Exit<boolean, MachineFailure>
+  >({
+    initial: prepare(context),
+    interruption: Event.CancelRequested(),
+    transition: (state, runtimeEvent): Decision => {
+      if (runtimeEvent._tag === "Input") return interrupt(state, runtimeEvent.cause ?? Cause.interrupt(undefined))
+      if (runtimeEvent._tag === "InvocationsStopped") {
+        if (state._tag !== "Stopping") return unexpected(state, runtimeEvent)
+        if (!state.from) return StateMachine.done(Exit.failCause(state.cause))
+        const finished = runtimeEvent.exits.map(completed)
+        if (state.from._tag === "RecoveringOverflow") {
+          const recovered = finished.some(
+            (event) => event._tag === "OverflowRecovered" && Exit.isSuccess(event.exit) && event.exit.value,
+          )
+          return recovered
+            ? StateMachine.done(Exit.failCause(state.cause))
+            : settle(state.from.active, Exit.failCause(state.cause), state.cause)
+        }
+        const tools = new Map(state.from.active.tools)
+        finished.forEach((event) => {
+          if (event._tag === "ToolFinished") tools.set(event.call.id, { call: event.call, exit: event.exit })
+        })
+        const active = { ...state.from.active, tools }
+        if (state.from._tag === "ObservingProvider")
+          return finishProvider(active, Exit.failCause(state.cause), state.cause)
+        const provider = finished.find((event) => event._tag === "ProviderFinished")
+        const stream =
+          provider && Exit.isFailure(provider.exit) ? Exit.failCause(provider.exit.cause) : state.from.stream
+        return settle(active, stream, state.cause)
+      }
+
+      const event = completed(runtimeEvent)
+      if (event._tag === "ToolFinished") {
+        if (
+          state._tag === "ObservingProvider" ||
+          state._tag === "FinalizingProvider" ||
+          state._tag === "AwaitingTools"
+        ) {
+          const tools = new Map(state.active.tools)
+          tools.set(event.call.id, { call: event.call, exit: event.exit })
+          const active = { ...state.active, tools }
+          return state._tag === "AwaitingTools"
+            ? afterProvider(active, state.stream)
+            : StateMachine.next({ ...state, active })
+        }
+        return unexpected(state, event)
+      }
+
+      return State.$match(state, {
+        PreparingAttempt: (state) => {
+          if (event._tag !== "Prepared") return unexpected(state, event)
+          if (Exit.isFailure(event.exit)) return StateMachine.done(Exit.failCause(event.exit.cause))
+          if (event.exit.value.preparation._tag === "Rebuilt") return prepare(event.exit.value.context, true)
+          const active = {
+            context: event.exit.value.context,
+            attempt: event.exit.value.preparation.attempt,
+            tools: new Map<string, ToolRun>(),
+          }
+          return pull(active)
+        },
+        ObservingProvider: (state) => {
+          if (event._tag !== "ProviderObserved") return unexpected(state, event)
+          if (Exit.isFailure(event.exit)) return finishProvider(state.active, Exit.failCause(event.exit.cause))
+          const observed = event.exit.value
+          if (observed._tag === "ProviderEnd") return finishProvider(state.active, Exit.succeed(undefined))
+          const tools = new Map(state.active.tools)
+          tools.set(observed.call.id, { call: observed.call })
+          const next = { ...state.active, tools }
+          return StateMachine.next(
+            State.ObservingProvider({ active: next }),
+            StateMachine.invoke<Operation>(
+              toolID(observed.call),
+              Operation.RunTool({
+                attempt: next.attempt,
+                call: observed.call,
+              }),
+            ),
+            StateMachine.invoke<Operation>(PROVIDER, Operation.ObserveProvider({ attempt: next.attempt })),
+          )
+        },
+        FinalizingProvider: (state) => {
+          if (event._tag !== "ProviderFinished") return unexpected(state, event)
+          const stream = Exit.isFailure(event.exit) ? Exit.failCause(event.exit.cause) : state.stream
+          return state.stopping ? settle(state.active, stream, state.stopping) : afterProvider(state.active, stream)
+        },
+        RecoveringOverflow: (state) => {
+          if (event._tag !== "OverflowRecovered") return unexpected(state, event)
+          if (Exit.isFailure(event.exit)) return StateMachine.done(Exit.failCause(event.exit.cause))
+          if (!event.exit.value) return settle(state.active, state.stream)
+          const context = { ...state.active.context, recoverOverflow: false }
+          return prepare(context, true)
+        },
+        SettlingAttempt: (state) => {
+          if (event._tag !== "AttemptSettled") return unexpected(state, event)
+          if (state.stopping) return StateMachine.done(Exit.failCause(state.stopping))
+          if (Exit.isFailure(event.exit)) return StateMachine.done(Exit.failCause(event.exit.cause))
+          const backoff = (outcome: BackoffOutcome) =>
+            StateMachine.next(
+              State.BackingOff({ context: state.active.context, outcome }),
+              StateMachine.invoke(RETRY, Operation.Retry({ context: state.active.context, outcome })),
+            )
+          return SessionStep.Outcome.$match(event.exit.value, {
+            Completed: (outcome) => StateMachine.done(Exit.succeed(outcome.needsContinuation)),
+            Retry: backoff,
+            Continue: backoff,
+            RecoverFull: () => prepare({ ...state.active.context, recoverContinuation: false }),
+          })
+        },
+        BackingOff: (state) => {
+          if (event._tag !== "RetryFinished") return unexpected(state, event)
+          if (Exit.isFailure(event.exit)) return StateMachine.done(Exit.failCause(event.exit.cause))
+          return prepare(state.context, state.outcome._tag === "Continue")
+        },
+        AwaitingTools: (state) => unexpected(state, event),
+        Stopping: (state) => unexpected(state, event),
+      })
+    },
+  })
+}
+
+const toolID = (call: ToolCall) => `tool:${call.id}`
+
+// Pre-start interruption can bypass the interpreter's Effect.exit.
+// Normalize outer failures once without erasing operation-specific error types.
+function completed<Failure>(
+  invocation: StateMachine.InvocationExited<Event<Failure>, Operation, never>,
+): Event<Failure> {
+  if (Exit.isSuccess(invocation.exit)) return invocation.exit.value
+  const exit = Exit.failCause(invocation.exit.cause)
+  return Operation.$match(invocation.operation, {
+    PrepareAttempt: () => Event.Prepared({ exit }),
+    ObserveProvider: () => Event.ProviderObserved({ exit }),
+    RunTool: (operation) => Event.ToolFinished({ call: operation.call, exit }),
+    FinishProvider: () => Event.ProviderFinished({ exit }),
+    RecoverOverflow: () => Event.OverflowRecovered({ exit }),
+    SettleAttempt: () => Event.AttemptSettled({ exit }),
+    Retry: () => Event.RetryFinished({ exit }),
+  })
+}
+
+function unexpected(state: State, event: { readonly _tag: string }): never {
+  throw new Error(`Unexpected ${event._tag} event while Session Step machine is ${state._tag}`)
+}

packages/core/src/session/runner/step.ts

@@ -9,13 +9,12 @@ import {
   type ProviderErrorEvent,
   type ToolCall,
 } from "@opencode-ai/ai"
-import { Cause, Data, Effect, Exit, Fiber, Option, Stream } from "effect"
+import { Cause, Data, Effect, Exit, Option, Pull, Scope, Stream } from "effect"
 import { SessionError } from "@opencode-ai/schema/session-error"
 import { Agent } from "../../agent.js"
 import { Bus } from "../../bus.js"
 import { Permission } from "../../permission.js"
 import { Snapshot } from "../../snapshot.js"
-import { Tool } from "../../tool.js"
 import { ToolOutput } from "../../tool-output.js"
 import { QuestionTool } from "../../tool/plugin/question.js"
 import { StepFailedError } from "../error.js"
@@ -34,11 +33,10 @@ export type Outcome = Data.TaggedEnum<{
   Retry: { readonly cause: AIError; readonly error: SessionError.Error }
   Continue: { readonly cause: AIError; readonly error: SessionError.Error }
   RecoverFull: {}
-  Compacted: {}
 }>
 export const Outcome = Data.taggedEnum<Outcome>()
 
-interface Input {
+export interface Input {
   readonly sessionID: SessionSchema.ID
   readonly assistantMessageID: SessionMessage.ID
   readonly agent: Agent.ID
@@ -49,6 +47,27 @@ interface Input {
   readonly recoverOverflow: Effect.Effect<boolean>
 }
 
+export type ProviderObservation = Data.TaggedEnum<{
+  ToolCall: { readonly call: ToolCall }
+  ProviderEnd: {}
+}>
+export const ProviderObservation = Data.taggedEnum<ProviderObservation>()
+
+export type ToolExit = Exit.Exit<void, Permission.DeclinedError | QuestionTool.CancelledError>
+
+export interface Settlement {
+  readonly stream: Exit.Exit<void, AIError>
+  readonly tools: ReadonlyArray<{ readonly call: ToolCall; readonly exit: ToolExit }>
+}
+
+export interface Attempt {
+  readonly observeUntilBoundary: () => Effect.Effect<ProviderObservation, AIError>
+  readonly runTool: (call: ToolCall) => Effect.Effect<void, Permission.DeclinedError | QuestionTool.CancelledError>
+  readonly finishProvider: (stream: Exit.Exit<void, AIError>) => Effect.Effect<void>
+  readonly recoverOverflow: (settlement: Settlement) => Effect.Effect<boolean>
+  readonly settle: (settlement: Settlement) => Effect.Effect<Outcome, AIError | StepFailedError>
+}
+
 const TOOLS_INTERRUPTED = { type: "aborted", message: "Tool execution interrupted" } as const
 const STEP_INTERRUPTED = { type: "aborted", message: "Step interrupted" } as const
 const RESULT_MISSING = { type: "tool.result-missing", message: "Provider did not return a tool result" } as const
@@ -60,7 +79,7 @@ export const make = Effect.gen(function* () {
   const snapshots = yield* Snapshot.Service
   const toolOutput = yield* ToolOutput.Service
 
-  const attempt = Effect.fn("SessionStep.attempt")(function* (input: Input) {
+  const open = Effect.fn("SessionStep.open")(function* (input: Input) {
     const startSnapshot = yield* snapshots.capture()
     const publisher = createLLMEventPublisher(bus, {
       sessionID: input.sessionID,
@@ -70,185 +89,197 @@ export const make = Effect.gen(function* () {
       providerMetadataKey: input.model.model.route.providerMetadataKey ?? input.model.model.provider,
       snapshot: startSnapshot,
     })
-    const toolRuns: Array<{
-      readonly call: ToolCall
-      readonly fiber: Fiber.Fiber<void, Permission.DeclinedError | QuestionTool.CancelledError>
-    }> = []
-    const interruptTools = Effect.suspend(() => Fiber.interruptAll(toolRuns.map((run) => run.fiber)))
-    const executeTool = (call: ToolCall) => {
-      if (input.prepared.request.toolChoice?.type === "none")
-        return new Tool.Error({ message: "Tools are disabled after the maximum agent steps" })
-      return input.prepared.executeTool({
-        sessionID: input.sessionID,
-        agent: input.agent,
-        messageID: input.assistantMessageID,
-        call,
-        progress: (update) => publisher.progress(call.id, update),
-      })
-    }
-
-    // Provider and tool fibers retain per-source order without a shared writer queue.
-    // A local execution starts only after its Tool.Called publication completes.
+    const scope = yield* Scope.Scope
+    const providerScope = yield* Scope.fork(scope)
+    const pull = yield* llm
+      .stream(input.prepared.request, input.prepared.options)
+      .pipe(Stream.ensuring(publisher.flush()), Stream.toPull, Scope.provide(providerScope))
+    let buffered: ReadonlyArray<LLMEvent> = []
+    let offset = 0
     let overflowFailure: ProviderErrorEvent | undefined
-    // Read to the end, not just the finish event, so the next request can reuse this response.
-    const providerStream = llm.stream(input.prepared.request, input.prepared.options).pipe(
-      Stream.runForEach((event) =>
-        Effect.gen(function* () {
-          if (overflowFailure || publisher.hasProviderError()) return
+
+    const observeUntilBoundary = Effect.fnUntraced(function* (): Effect.fn.Return<ProviderObservation, AIError> {
+      while (true) {
+        const event = buffered[offset]
+        if (event) {
+          offset += 1
+          if (overflowFailure || publisher.hasProviderError()) continue
           if (
             LLMEvent.is.providerError(event) &&
             isContextOverflowFailure(event) &&
             !publisher.record().outputStarted
           ) {
             overflowFailure = event
-            return
+            continue
           }
-          yield* publisher.publish(event)
-          if (event.type !== "tool-call" || event.providerExecuted) return
-          toolRuns.push({
-            call: event,
-            fiber: yield* Effect.uninterruptibleMask((restore) =>
-              restore(executeTool(event)).pipe(
-                Effect.flatMap(toolOutput.truncate),
-                Effect.flatMap((outcome) => publisher.toolExecution(event.id, event.name, outcome)),
-                Effect.catchTag("Tool.Error", (error) =>
-                  publisher.failTool(event.id, toSessionError(error), error.metadata).pipe(Effect.asVoid),
-                ),
-              ),
-            ).pipe(Effect.forkScoped),
-          })
-        }),
-      ),
-      Effect.ensuring(publisher.flush()),
-    )
-
-    // Keep the final tool and Step events uninterruptible, even when the work itself is cancelled.
-    return yield* Effect.uninterruptibleMask((restore) =>
-      Effect.gen(function* () {
-        const stream = yield* restore(providerStream).pipe(Effect.exit)
-        const streamFailure = Option.getOrUndefined(Exit.findErrorOption(stream))
-        const streamInterrupted = Exit.hasInterrupts(stream)
-        if (!overflowFailure && publisher.hasStarted()) yield* publisher.streamed()
-        if (streamInterrupted) yield* interruptTools
-        const joined = yield* restore(Fiber.awaitAll(toolRuns.map((run) => run.fiber))).pipe(Effect.exit)
-        if (Exit.isFailure(joined)) yield* interruptTools
-        const tools = classifyToolExits(joined, toolRuns)
-
-        if (
-          !publisher.record().outputStarted &&
-          isContextOverflowFailure(overflowFailure ?? streamFailure) &&
-          (yield* restore(input.recoverOverflow))
-        )
-          return Outcome.Compacted()
-
-        if (overflowFailure) yield* publisher.publish(overflowFailure)
-        const recorded = publisher.record()
-        const unknownFinish =
-          Exit.isSuccess(stream) && recorded.finish?.finish === "unknown"
-            ? new AIError({
-                reason: new InvalidProviderOutputError({
-                  message: "The provider response ended with an unknown finish reason.",
-                  classification: "incomplete-stream",
-                }),
-              })
-            : undefined
-        const llmFailure = streamFailure instanceof AIError ? streamFailure : unknownFinish
-        const llmError = llmFailure && !recorded.providerFailed ? toSessionError(llmFailure) : undefined
-        if (
-          input.recoverContinuation &&
-          llmFailure?.reason._tag === "Transport" &&
-          (llmFailure.reason.recovery === "retry-full" || llmFailure.reason.recovery === "rotate-and-retry-full") &&
-          !recorded.outputStarted
-        )
-          return Outcome.RecoverFull()
-        if (llmFailure && llmError && SessionRunnerRetry.isRetryable(llmFailure) && !recorded.outputStarted) {
-          // Retry state projects onto the existing assistant, even before it has produced output.
-          yield* publisher.startAssistant()
-          return Outcome.Retry({ cause: llmFailure, error: llmError })
+          // Keep the publisher's in-memory mark and durable write indivisible under cancellation.
+          yield* publisher.publish(event).pipe(Effect.uninterruptible)
+          if (event.type === "tool-call" && !event.providerExecuted)
+            return ProviderObservation.ToolCall({ call: event })
+          continue
         }
-        if (llmError) yield* publisher.failAssistant(llmError)
+        const chunk = yield* pull.pipe(Pull.catchDone(() => Effect.succeed(undefined)))
+        if (!chunk) return ProviderObservation.ProviderEnd()
+        buffered = chunk
+        offset = 0
+      }
+    })
 
-        for (const decline of tools.declines)
-          yield* publisher.failTool(decline.call.id, {
-            type: "aborted",
-            message:
-              decline.reason._tag === "QuestionTool.CancelledError"
-                ? decline.reason.message
-                : "The user declined this tool call",
-          })
-        const interrupted = tools.declines.length > 0 || streamInterrupted || tools.interrupted
-        const toolFailure = interrupted
-          ? TOOLS_INTERRUPTED
-          : tools.failure !== undefined
-            ? toSessionError(Cause.squash(tools.failure))
-            : recorded.providerFailed
-              ? TOOLS_INTERRUPTED
-              : undefined
-        if (toolFailure) yield* publisher.failUnsettledTools(toolFailure)
-        if (interrupted) yield* publisher.failAssistant(STEP_INTERRUPTED)
+    const runTool = Effect.fnUntraced(function* (call: ToolCall) {
+      return yield* Effect.uninterruptibleMask((restore) => {
+        if (input.prepared.request.toolChoice?.type === "none")
+          return publisher
+            .failTool(call.id, { type: "tool.execution", message: "Tools are disabled after the maximum agent steps" })
+            .pipe(Effect.asVoid)
+        return restore(
+          input.prepared.executeTool({
+            sessionID: input.sessionID,
+            agent: input.agent,
+            messageID: input.assistantMessageID,
+            call,
+            progress: (update) => publisher.progress(call.id, update),
+          }),
+        ).pipe(
+          Effect.flatMap(toolOutput.truncate),
+          Effect.flatMap((outcome) => publisher.toolExecution(call.id, call.name, outcome)),
+          Effect.catchTag("Tool.Error", (error) =>
+            publisher.failTool(call.id, toSessionError(error), error.metadata).pipe(Effect.asVoid),
+          ),
+        )
+      })
+    })
 
-        // All local fibers have joined; only provider-hosted results can still be missing.
-        if (llmError || (Exit.isSuccess(stream) && !recorded.providerFailed)) {
-          const missing = yield* publisher.failUnsettledTools(RESULT_MISSING, "hosted")
-          if (missing && !llmError && !recorded.finish) yield* publisher.failAssistant(RESULT_MISSING)
-        }
+    const finishProvider = Effect.fnUntraced(function* (stream: Exit.Exit<void, AIError>) {
+      yield* Scope.close(providerScope, stream)
+      if (!overflowFailure && publisher.hasStarted()) yield* publisher.streamed()
+    }, Effect.uninterruptible)
 
-        const record = publisher.record()
-        if (record.finish || record.failure) {
-          const snapshot = yield* snapshots.capture()
-          const files =
-            startSnapshot && snapshot
-              ? startSnapshot === snapshot
-                ? []
-                : yield* snapshots
-                    .files({ from: startSnapshot, to: snapshot })
-                    .pipe(Effect.orElseSucceed(() => undefined))
-              : undefined
-          const usage = record.finish
-            ? { cost: SessionUsage.calculateCost(input.model.cost, record.finish.tokens), tokens: record.finish.tokens }
-            : undefined
-          if (record.failure) yield* publisher.publishStepFailure({ ...usage, snapshot, files })
-          if (record.finish && usage && !record.failure)
-            yield* bus.publish(SessionEvent.Step.Ended, {
-              sessionID: input.sessionID,
-              assistantMessageID: yield* publisher.startAssistant(),
-              finish: record.finish.finish,
-              rawFinish: record.finish.rawFinish,
-              providerState: record.finish.providerState,
-              ...usage,
-              snapshot,
-              files,
+    const recoverOverflow = (settlement: Settlement) => {
+      if (publisher.record().outputStarted) return Effect.succeed(false)
+      const failure = overflowFailure ?? Option.getOrUndefined(Exit.findErrorOption(settlement.stream))
+      return isContextOverflowFailure(failure) ? input.recoverOverflow : Effect.succeed(false)
+    }
+
+    const settle = Effect.fn("SessionStep.settle")(function* (settlement: Settlement) {
+      const streamFailure = Option.getOrUndefined(Exit.findErrorOption(settlement.stream))
+      const streamInterrupted = Exit.hasInterrupts(settlement.stream)
+      const tools = classifyToolExits(settlement.tools)
+
+      if (overflowFailure) yield* publisher.publish(overflowFailure)
+      const recorded = publisher.record()
+      const unknownFinish =
+        Exit.isSuccess(settlement.stream) && recorded.finish?.finish === "unknown"
+          ? new AIError({
+              reason: new InvalidProviderOutputError({
+                message: "The provider response ended with an unknown finish reason.",
+                classification: "incomplete-stream",
+              }),
             })
-        }
+          : undefined
+      const llmFailure = streamFailure instanceof AIError ? streamFailure : unknownFinish
+      const llmError = llmFailure && !recorded.providerFailed ? toSessionError(llmFailure) : undefined
+      if (
+        input.recoverContinuation &&
+        llmFailure?.reason._tag === "Transport" &&
+        (llmFailure.reason.recovery === "retry-full" || llmFailure.reason.recovery === "rotate-and-retry-full") &&
+        !recorded.outputStarted
+      )
+        return Outcome.RecoverFull()
+      if (llmFailure && llmError && SessionRunnerRetry.isRetryable(llmFailure) && !recorded.outputStarted) {
+        yield* publisher.startAssistant()
+        return Outcome.Retry({ cause: llmFailure, error: llmError })
+      }
+      if (llmError) yield* publisher.failAssistant(llmError)
 
-        // After durable output, recovery continues instead of replaying: the
-        // partial assistant message is already persisted history. Any failure
-        // the pre-output gate would retry is continued here, plus interrupted
-        // streams, whose read failures may carry delivery states the retry
-        // policy rejects for full resends.
-        if (
-          llmFailure &&
-          llmError &&
-          (isInterruptedStream(llmFailure) || SessionRunnerRetry.isRetryable(llmFailure)) &&
-          record.outputStarted &&
-          tools.declines.length === 0 &&
-          !tools.interrupted
-        )
-          return Outcome.Continue({ cause: llmFailure, error: llmError })
-
-        if (Exit.isFailure(stream)) return yield* Effect.failCause(stream.cause)
-        if (tools.declines.length > 0) return yield* Effect.interrupt
-        if (tools.interrupted && tools.failure) return yield* Effect.failCause(tools.failure)
-        if (tools.interrupted && Exit.isFailure(joined)) return yield* Effect.failCause(joined.cause)
-        if (record.failure) return yield* new StepFailedError({ error: record.failure })
-        return Outcome.Completed({
-          needsContinuation: input.prepared.request.toolChoice?.type !== "none" && record.needsContinuation,
+      for (const decline of tools.declines)
+        yield* publisher.failTool(decline.call.id, {
+          type: "aborted",
+          message:
+            decline.reason._tag === "QuestionTool.CancelledError"
+              ? decline.reason.message
+              : "The user declined this tool call",
         })
-      }),
-    )
-  }, Effect.scoped)
+      const interrupted = tools.declines.length > 0 || streamInterrupted || tools.interrupted
+      const toolFailure = interrupted
+        ? TOOLS_INTERRUPTED
+        : tools.failure !== undefined
+          ? toSessionError(Cause.squash(tools.failure))
+          : recorded.providerFailed
+            ? TOOLS_INTERRUPTED
+            : undefined
+      if (toolFailure) yield* publisher.failUnsettledTools(toolFailure)
+      if (interrupted) yield* publisher.failAssistant(STEP_INTERRUPTED)
 
-  return { attempt }
+      if (llmError || (Exit.isSuccess(settlement.stream) && !recorded.providerFailed)) {
+        const missing = yield* publisher.failUnsettledTools(RESULT_MISSING, "hosted")
+        if (missing && !llmError && !recorded.finish) yield* publisher.failAssistant(RESULT_MISSING)
+      }
+
+      const record = publisher.record()
+      if (record.finish || record.failure) {
+        const snapshot = yield* snapshots.capture()
+        const files =
+          startSnapshot && snapshot
+            ? startSnapshot === snapshot
+              ? []
+              : yield* snapshots
+                  .files({ from: startSnapshot, to: snapshot })
+                  .pipe(Effect.orElseSucceed(() => undefined))
+            : undefined
+        const usage = record.finish
+          ? {
+              cost: SessionUsage.calculateCost(input.model.cost, record.finish.tokens),
+              tokens: record.finish.tokens,
+            }
+          : undefined
+        if (record.failure) yield* publisher.publishStepFailure({ ...usage, snapshot, files })
+        if (record.finish && usage && !record.failure)
+          yield* bus.publish(SessionEvent.Step.Ended, {
+            sessionID: input.sessionID,
+            assistantMessageID: yield* publisher.startAssistant(),
+            finish: record.finish.finish,
+            rawFinish: record.finish.rawFinish,
+            providerState: record.finish.providerState,
+            ...usage,
+            snapshot,
+            files,
+          })
+      }
+
+      // After durable output, recovery continues instead of replaying: the
+      // partial assistant message is already persisted history. Any failure
+      // the pre-output gate would retry is continued here, plus interrupted
+      // streams, whose read failures may carry delivery states the retry
+      // policy rejects for full resends.
+      if (
+        llmFailure &&
+        llmError &&
+        (isInterruptedStream(llmFailure) || SessionRunnerRetry.isRetryable(llmFailure)) &&
+        record.outputStarted &&
+        tools.declines.length === 0 &&
+        !tools.interrupted
+      )
+        return Outcome.Continue({ cause: llmFailure, error: llmError })
+
+      if (Exit.isFailure(settlement.stream)) return yield* Effect.failCause(settlement.stream.cause)
+      if (tools.declines.length > 0) return yield* Effect.interrupt
+      if (tools.interrupted && tools.failure) return yield* Effect.failCause(tools.failure)
+      if (record.failure) return yield* new StepFailedError({ error: record.failure })
+      return Outcome.Completed({
+        needsContinuation: input.prepared.request.toolChoice?.type !== "none" && record.needsContinuation,
+      })
+    }, Effect.uninterruptible)
+
+    return {
+      observeUntilBoundary,
+      runTool,
+      finishProvider,
+      recoverOverflow,
+      settle,
+    } satisfies Attempt
+  })
+
+  return { open }
 })
 
 const isInterruptedStream = (failure: AIError) => {
@@ -259,20 +290,19 @@ const isInterruptedStream = (failure: AIError) => {
 
 /** Tool.Error settles in each fiber; only user declines remain in the typed error channel. */
 const classifyToolExits = (
-  settled: Exit.Exit<Array<Exit.Exit<void, Permission.DeclinedError | QuestionTool.CancelledError>>>,
-  runs: ReadonlyArray<{ readonly call: ToolCall }>,
+  runs: ReadonlyArray<{
+    readonly call: ToolCall
+    readonly exit: ToolExit
+  }>,
 ) => {
-  const exits = Exit.isSuccess(settled) ? settled.value : []
-  const declines = exits.flatMap((exit, index) =>
-    Exit.isFailure(exit)
-      ? exit.cause.reasons.flatMap((reason) =>
-          Cause.isFailReason(reason) ? [{ call: runs[index].call, reason: reason.error }] : [],
+  const declines = runs.flatMap((run) =>
+    Exit.isFailure(run.exit)
+      ? run.exit.cause.reasons.flatMap((reason) =>
+          Cause.isFailReason(reason) ? [{ call: run.call, reason: reason.error }] : [],
         )
       : [],
   )
-  const causes = Exit.isFailure(settled)
-    ? [settled.cause]
-    : exits.flatMap((exit) => (Exit.isFailure(exit) ? [exit.cause] : []))
+  const causes = runs.flatMap((run) => (Exit.isFailure(run.exit) ? [run.exit.cause] : []))
   const failure = causes
     .flatMap((cause) => {
       if (Cause.hasInterrupts(cause)) return []

packages/core/test/effect/state-machine.test.ts

@@ -0,0 +1,415 @@
+import { describe, expect } from "bun:test"
+import { Cause, Deferred, Effect, Exit, Fiber, Option, Ref, Scheduler } from "effect"
+import { StateMachine } from "@opencode-ai/core/effect/state-machine"
+import { it } from "../lib/effect"
+
+describe("StateMachine", () => {
+  it.effect("runs invoked operations through pure transitions", () => {
+    type Event = { readonly _tag: "Completed"; readonly value: number }
+    type Operation = { readonly _tag: "Work" }
+    const definition = StateMachine.define<"running", Event, Operation, never, number>({
+      initial: StateMachine.next("running", StateMachine.invoke("work", { _tag: "Work" })),
+      transition: (state, event) => {
+        expect(state).toBe("running")
+        expect(event._tag).toBe("InvocationExited")
+        if (event._tag !== "InvocationExited" || Exit.isFailure(event.exit)) return StateMachine.done(-1)
+        return StateMachine.done(event.exit.value.value)
+      },
+    })
+    return StateMachine.run(definition, () => Effect.succeed({ _tag: "Completed", value: 42 })).pipe(
+      Effect.map((output) => expect(output).toBe(42)),
+    )
+  })
+
+  it.effect("preserves the operation Cause", () => {
+    type Operation = { readonly _tag: "Work" }
+    const definition = StateMachine.define<"running", never, Operation, string, Cause.Cause<string>>({
+      initial: StateMachine.next("running", StateMachine.invoke("work", { _tag: "Work" })),
+      transition: (_, event) => {
+        if (event._tag === "InvocationExited" && Exit.isFailure(event.exit)) return StateMachine.done(event.exit.cause)
+        throw new Error("Expected the invocation to fail")
+      },
+    })
+    return StateMachine.run(definition, () => Effect.fail("boom")).pipe(
+      Effect.map((cause) => {
+        expect(Option.getOrUndefined(Cause.findErrorOption(cause))).toBe("boom")
+      }),
+    )
+  })
+
+  it.effect("settles owned work before propagating interruption", () =>
+    Effect.gen(function* () {
+      const finalized = yield* Deferred.make<void>()
+      type State = "running" | "stopping"
+      type Event = { readonly _tag: "Cancel" }
+      type Operation = { readonly _tag: "Work" }
+      const definition = StateMachine.define<State, Event, Operation, never, "cancelled">({
+        initial: StateMachine.next("running", StateMachine.invoke("work", { _tag: "Work" })),
+        interruption: { _tag: "Cancel" } as const,
+        transition: (state, event) => {
+          if (event._tag === "Input") {
+            expect(state).toBe("running")
+            return StateMachine.next("stopping" as const, StateMachine.stop("work"))
+          }
+          expect(state).toBe("stopping")
+          if (event._tag !== "InvocationExited") throw new Error("Expected the invocation to stop")
+          expect(Exit.hasInterrupts(event.exit)).toBe(true)
+          return StateMachine.done("cancelled" as const)
+        },
+      })
+      const machine = yield* StateMachine.run(definition, () =>
+        Effect.never.pipe(Effect.ensuring(Deferred.succeed(finalized, undefined))),
+      ).pipe(Effect.forkChild({ startImmediately: true }))
+
+      yield* Effect.yieldNow
+      yield* Fiber.interrupt(machine)
+      const exit = yield* Fiber.await(machine)
+      expect(Exit.hasInterrupts(exit)).toBe(true)
+      expect(yield* Deferred.isDone(finalized)).toBe(true)
+    }),
+  )
+
+  it.effect("runs cleanup invocations after interruption", () =>
+    Effect.gen(function* () {
+      const workStarted = yield* Deferred.make<void>()
+      const cleanupRan = yield* Deferred.make<void>()
+      type State = "running" | "stopping" | "cleaning"
+      type Event = { readonly _tag: "Cancel" } | { readonly _tag: "WorkDone" } | { readonly _tag: "CleanupDone" }
+      type Operation = { readonly _tag: "Work" } | { readonly _tag: "Cleanup" }
+      const definition = StateMachine.define<State, Event, Operation, never, void>({
+        initial: StateMachine.next("running", StateMachine.invoke("phase", { _tag: "Work" })),
+        interruption: { _tag: "Cancel" },
+        transition: (state, event) => {
+          if (event._tag === "Input")
+            return StateMachine.next("stopping", StateMachine.stopAndJoin("interruption", ["phase"]))
+          if (state === "stopping") {
+            if (event._tag !== "InvocationsStopped") throw new Error("Expected the aggregate stop result")
+            expect(event.id).toBe("interruption")
+            expect(event.exits).toMatchObject([{ id: "phase", operation: { _tag: "Work" } }])
+            expect(Exit.hasInterrupts(event.exits[0].exit)).toBe(true)
+            return StateMachine.next("cleaning", StateMachine.invoke("cleanup", { _tag: "Cleanup" }))
+          }
+          if (state === "cleaning") return StateMachine.done(undefined)
+          throw new Error("Unexpected state machine transition")
+        },
+      })
+      const machine = yield* StateMachine.run(definition, (operation) => {
+        if (operation._tag === "Cleanup")
+          return Deferred.succeed(cleanupRan, undefined).pipe(Effect.as({ _tag: "CleanupDone" } as const))
+        return Deferred.succeed(workStarted, undefined).pipe(
+          Effect.andThen(Effect.never),
+          Effect.as({ _tag: "WorkDone" } as const),
+        )
+      }).pipe(Effect.forkChild({ startImmediately: true }))
+
+      yield* Deferred.await(workStarted)
+      yield* Fiber.interrupt(machine)
+      expect(Exit.hasInterrupts(yield* Fiber.await(machine))).toBe(true)
+      expect(yield* Deferred.isDone(cleanupRan)).toBe(true)
+    }),
+  )
+
+  it.effect("stops invocations together and joins cross-dependent finalizers", () =>
+    Effect.gen(function* () {
+      const started = { left: yield* Deferred.make<void>(), right: yield* Deferred.make<void>() }
+      const finalizing = { left: yield* Deferred.make<void>(), right: yield* Deferred.make<void>() }
+      const finalized = yield* Ref.make<ReadonlyArray<string>>([])
+      type State = "running" | "stopping" | "verifying"
+      type Event = "ready" | "verified"
+      type Operation = "left" | "right" | "trigger" | "verify"
+      const definition = StateMachine.define<State, Event, Operation, never, boolean>({
+        initial: StateMachine.next(
+          "running",
+          StateMachine.invoke<Operation>("left", "left"),
+          StateMachine.invoke<Operation>("right", "right"),
+          StateMachine.invoke<Operation>("trigger", "trigger"),
+        ),
+        transition: (state, event) => {
+          if (event._tag === "InvocationExited" && event.operation === "trigger")
+            return StateMachine.next("stopping", StateMachine.stopAndJoin("workers", ["left", "right"]))
+          if (event._tag === "InvocationsStopped") {
+            expect(state).toBe("stopping")
+            expect(event.id).toBe("workers")
+            expect(event.exits).toMatchObject([
+              { _tag: "InvocationExited", id: "left", generation: 1, operation: "left" },
+              { _tag: "InvocationExited", id: "right", generation: 2, operation: "right" },
+            ])
+            expect(event.exits.every((invocation) => Exit.hasInterrupts(invocation.exit))).toBe(true)
+            return StateMachine.next("verifying", StateMachine.invoke("verify", "verify"))
+          }
+          if (event._tag === "InvocationExited" && event.operation === "verify") {
+            expect(state).toBe("verifying")
+            expect(event.exit).toEqual(Exit.succeed("verified"))
+            return StateMachine.done(true)
+          }
+          throw new Error("Unexpected state machine transition")
+        },
+      })
+      const output = yield* StateMachine.run(definition, (operation) => {
+        if (operation === "trigger")
+          return Deferred.await(started.left).pipe(Effect.andThen(Deferred.await(started.right)), Effect.as("ready"))
+        if (operation === "verify")
+          return Ref.get(finalized).pipe(
+            Effect.map((value) => {
+              expect(value.toSorted()).toEqual(["left", "right"])
+              return "verified" as const
+            }),
+          )
+        return Deferred.succeed(started[operation], undefined).pipe(
+          Effect.andThen(Effect.never),
+          Effect.ensuring(
+            Deferred.succeed(finalizing[operation], undefined).pipe(
+              Effect.andThen(Deferred.await(finalizing[operation === "left" ? "right" : "left"])),
+              Effect.andThen(Ref.update(finalized, (value) => [...value, operation])),
+            ),
+          ),
+        )
+      })
+      expect(output).toBe(true)
+    }),
+  )
+
+  it.effect("aggregates queued and never-started exits once without affecting reused keys", () =>
+    Effect.gen(function* () {
+      const completed = yield* Deferred.make<Fiber.Fiber<unknown, unknown>>()
+      const releaseCompleted = yield* Deferred.make<void>()
+      const gateStarted = yield* Deferred.make<void>()
+      const childStarted = yield* Deferred.make<void>()
+      type Event = "completed" | "triggered" | "replaced"
+      type Operation = "complete" | "gate" | "trigger" | "never-started" | "replacement"
+      type Seen = ReadonlyArray<StateMachine.RuntimeEvent<Event, Operation, never>>
+      const definition = StateMachine.define<Seen, Event, Operation, never, Seen>({
+        initial: StateMachine.next(
+          [],
+          StateMachine.invoke<Operation>("completed", "complete"),
+          StateMachine.invoke<Operation>("gate", "gate"),
+          StateMachine.invoke<Operation>("trigger", "trigger"),
+        ),
+        transition: (state, event) => {
+          const seen = [...state, event]
+          if (event._tag === "InvocationExited" && event.operation === "trigger")
+            return StateMachine.next(
+              seen,
+              StateMachine.stop("gate"),
+              StateMachine.invoke<Operation>("child", "never-started"),
+              StateMachine.stopAndJoin("batch", ["completed", "gate", "child"]),
+              StateMachine.invoke<Operation>("completed", "replacement"),
+              StateMachine.invoke<Operation>("child", "replacement"),
+            )
+          return seen.length === 4 ? StateMachine.done(seen) : StateMachine.next(seen)
+        },
+      })
+      const seen = yield* StateMachine.run(definition, (operation) => {
+        if (operation === "complete")
+          return Effect.withFiber((fiber) => Deferred.succeed(completed, fiber)).pipe(
+            Effect.andThen(Deferred.await(releaseCompleted)),
+            Effect.as("completed"),
+          )
+        if (operation === "gate")
+          return Deferred.succeed(gateStarted, undefined).pipe(
+            Effect.andThen(Effect.never),
+            // Hold the command loop until the completed child's exit is queued.
+            Effect.ensuring(
+              Deferred.succeed(releaseCompleted, undefined).pipe(
+                Effect.andThen(Deferred.await(completed)),
+                Effect.flatMap(Fiber.await),
+              ),
+            ),
+          )
+        if (operation === "trigger")
+          return Deferred.await(completed).pipe(Effect.andThen(Deferred.await(gateStarted)), Effect.as("triggered"))
+        if (operation === "never-started")
+          return Deferred.succeed(childStarted, undefined).pipe(Effect.andThen(Effect.never))
+        return Effect.succeed("replaced")
+      }).pipe(
+        // Keep the adjacent invoke/stop commands in one scheduler slice.
+        Effect.provideService(Scheduler.PreventSchedulerYield, true),
+      )
+
+      expect(seen.map((event) => (event._tag === "InvocationExited" ? event.operation : event._tag))).toEqual([
+        "trigger",
+        "InvocationsStopped",
+        "replacement",
+        "replacement",
+      ])
+      const stopped = seen[1]
+      if (stopped._tag !== "InvocationsStopped") throw new Error("Expected the aggregate stop result")
+      expect(stopped.id).toBe("batch")
+      expect(stopped.exits).toMatchObject([
+        {
+          _tag: "InvocationExited",
+          id: "completed",
+          generation: 1,
+          operation: "complete",
+          exit: Exit.succeed("completed"),
+        },
+        { _tag: "InvocationExited", id: "gate", generation: 2, operation: "gate" },
+        { _tag: "InvocationExited", id: "child", generation: 4, operation: "never-started" },
+      ])
+      expect(stopped.exits.slice(1).every((invocation) => Exit.hasInterrupts(invocation.exit))).toBe(true)
+      expect(seen.slice(2)).toMatchObject([
+        { _tag: "InvocationExited", id: "completed", generation: 5, exit: Exit.succeed("replaced") },
+        { _tag: "InvocationExited", id: "child", generation: 6, exit: Exit.succeed("replaced") },
+      ])
+      expect(yield* Deferred.isDone(childStarted)).toBe(false)
+    }),
+  )
+
+  it.effect("emits an empty aggregate for an empty stop batch", () => {
+    const definition = StateMachine.define<"stopping", never, never, never, boolean>({
+      initial: StateMachine.next("stopping", StateMachine.stopAndJoin("empty", [])),
+      transition: (_, event) => {
+        expect(event).toEqual({ _tag: "InvocationsStopped", id: "empty", exits: [] })
+        return StateMachine.done(true)
+      },
+    })
+    return StateMachine.run(definition, () => Effect.die("Unexpected operation")).pipe(
+      Effect.map((output) => expect(output).toBe(true)),
+    )
+  })
+
+  it.effect("awaits a never-started finalizer without interrupting it", () =>
+    Effect.gen(function* () {
+      const finalized = yield* Ref.make(0)
+      type Operation = "work" | "finalize"
+      const definition = StateMachine.define<"stopping", "finalized", Operation, never, boolean>({
+        initial: StateMachine.next(
+          "stopping",
+          StateMachine.invoke<Operation>("work", "work"),
+          StateMachine.invoke<Operation>("finalizer", "finalize"),
+          StateMachine.stopAndJoin("batch", ["work"], ["finalizer"]),
+        ),
+        transition: (_, event) => {
+          if (event._tag !== "InvocationsStopped") throw new Error("Expected only the joined batch")
+          expect(event.exits).toHaveLength(2)
+          expect(event.exits[0].id).toBe("work")
+          expect(Exit.hasInterrupts(event.exits[0].exit)).toBe(true)
+          expect(event.exits[1]).toMatchObject({ id: "finalizer", exit: Exit.succeed("finalized") })
+          return StateMachine.done(true)
+        },
+      })
+      expect(
+        yield* StateMachine.run(definition, (operation) =>
+          operation === "work"
+            ? Effect.never
+            : Ref.update(finalized, (count) => count + 1).pipe(Effect.as("finalized" as const)),
+        ).pipe(Effect.provideService(Scheduler.PreventSchedulerYield, true)),
+      ).toBe(true)
+      expect(yield* Ref.get(finalized)).toBe(1)
+    }),
+  )
+
+  it.effect("defects when a stop batch contains an unknown invocation", () =>
+    Effect.gen(function* () {
+      const definition = StateMachine.define<"stopping", never, "work", never, never>({
+        initial: StateMachine.next(
+          "stopping",
+          StateMachine.invoke("known", "work"),
+          StateMachine.stopAndJoin("batch", ["known", "unknown"]),
+        ),
+        transition: () => {
+          throw new Error("Unexpected state machine transition")
+        },
+      })
+      const exit = yield* StateMachine.run(definition, () => Effect.never).pipe(Effect.exit)
+      if (Exit.isSuccess(exit)) throw new Error("Expected an unknown invocation defect")
+      expect(Cause.hasDies(exit.cause)).toBe(true)
+      expect(Cause.prettyErrors(exit.cause).map((error) => error.message)).toEqual([
+        "Unknown state machine invocation in StopAndJoin",
+      ])
+    }),
+  )
+
+  it.effect("observes an individual exit when a deferred child is stopped before starting", () =>
+    Effect.gen(function* () {
+      const started = yield* Deferred.make<void>()
+      const definition = StateMachine.define<"stopping", never, "work", never, boolean>({
+        initial: StateMachine.next("stopping", StateMachine.invoke("work", "work"), StateMachine.stop("work")),
+        transition: (_, event) => {
+          if (event._tag !== "InvocationExited") throw new Error("Expected the invocation to stop")
+          expect(event.id).toBe("work")
+          expect(Exit.hasInterrupts(event.exit)).toBe(true)
+          return StateMachine.done(true)
+        },
+      })
+      const output = yield* StateMachine.run(definition, () =>
+        Deferred.succeed(started, undefined).pipe(Effect.andThen(Effect.never)),
+      ).pipe(Effect.provideService(Scheduler.PreventSchedulerYield, true))
+      expect(output).toBe(true)
+      expect(yield* Deferred.isDone(started)).toBe(false)
+    }),
+  )
+
+  it.effect("waits for replaced invocation cleanup and ignores its stale exit", () =>
+    Effect.gen(function* () {
+      const firstStarted = yield* Deferred.make<void>()
+      const releaseTrigger = yield* Deferred.make<void>()
+      const events = yield* Ref.make<ReadonlyArray<string>>([])
+      type State = "first" | "second"
+      type Event = { readonly _tag: "Triggered" } | { readonly _tag: "SecondDone" }
+      type Operation = { readonly _tag: "First" } | { readonly _tag: "Trigger" } | { readonly _tag: "Second" }
+      const definition = StateMachine.define<State, Event, Operation, never, string>({
+        initial: StateMachine.next(
+          "first",
+          StateMachine.invoke<Operation>("work", { _tag: "First" }),
+          StateMachine.invoke<Operation>("trigger", { _tag: "Trigger" }),
+        ),
+        transition: (state, event) => {
+          if (event._tag !== "InvocationExited" || Exit.isFailure(event.exit)) return StateMachine.done("unexpected")
+          if (event.operation._tag === "Trigger") {
+            return StateMachine.next("second" as const, StateMachine.invoke("work", { _tag: "Second" } as const))
+          }
+          if (state === "second") return StateMachine.done(event.exit.value._tag)
+          return StateMachine.next(state)
+        },
+      })
+      const output = yield* StateMachine.run(definition, (operation) => {
+        if (operation._tag === "Trigger")
+          return Deferred.await(releaseTrigger).pipe(Effect.as({ _tag: "Triggered" } as const))
+        if (operation._tag === "Second") {
+          return Ref.update(events, (value) => [...value, "second started"]).pipe(
+            Effect.as({ _tag: "SecondDone" } as const),
+          )
+        }
+        return Deferred.succeed(firstStarted, undefined).pipe(
+          Effect.andThen(Effect.never),
+          Effect.ensuring(Ref.update(events, (value) => [...value, "first finalized"])),
+        )
+      }).pipe(Effect.forkChild({ startImmediately: true }))
+
+      yield* Deferred.await(firstStarted)
+      yield* Deferred.succeed(releaseTrigger, undefined)
+      expect(yield* Fiber.join(output)).toBe("SecondDone")
+      expect(yield* Ref.get(events)).toEqual(["first finalized", "second started"])
+    }),
+  )
+
+  it.effect("does not start the next invocation when interruption is pending at the transition boundary", () =>
+    Effect.gen(function* () {
+      const releaseFirst = yield* Deferred.make<void>()
+      const secondStarted = yield* Deferred.make<void>()
+      type State = "first" | "second"
+      type Event = { readonly _tag: "FirstDone" } | { readonly _tag: "SecondDone" }
+      type Operation = { readonly _tag: "First" } | { readonly _tag: "Second" }
+      let machine: Fiber.Fiber<string> | undefined
+      const definition = StateMachine.define<State, Event, Operation, never, string>({
+        initial: StateMachine.next("first", StateMachine.invoke("work", { _tag: "First" })),
+        transition: (state, event) => {
+          if (event._tag !== "InvocationExited" || Exit.isFailure(event.exit)) return StateMachine.done("unexpected")
+          if (state === "second") return StateMachine.done("completed")
+          machine?.interruptUnsafe(123)
+          return StateMachine.next("second", StateMachine.invoke("work", { _tag: "Second" }))
+        },
+      })
+      machine = yield* StateMachine.run(definition, (operation) =>
+        operation._tag === "First"
+          ? Deferred.await(releaseFirst).pipe(Effect.as({ _tag: "FirstDone" } as const))
+          : Deferred.succeed(secondStarted, undefined).pipe(Effect.as({ _tag: "SecondDone" } as const)),
+      ).pipe(Effect.forkChild({ startImmediately: true }))
+
+      yield* Deferred.succeed(releaseFirst, undefined)
+      expect(Exit.hasInterrupts(yield* Fiber.await(machine))).toBe(true)
+      expect(yield* Deferred.isDone(secondStarted)).toBe(false)
+    }),
+  )
+})

packages/core/test/session-runner.test.ts

@@ -582,6 +582,13 @@ const scenario = (
     }),
   )
 
+const nextRetryScheduled = (s: Scenario) =>
+  s.bus.subscribe(SessionEvent.RetryScheduled).pipe(
+    Stream.filter((event) => event.data.sessionID === sessionID),
+    Stream.runHead,
+    Effect.forkScoped({ startImmediately: true }),
+  )
+
 const providerUnavailable = () =>
   new AIError({
     reason: new TransportError({
@@ -4357,8 +4364,9 @@ describe("SessionRunnerLLM", () => {
     yield* s.llm.push(Stream.fail(providerUnavailable()))
     yield* s.llm.push(TestLLM.text("Recovered", "retry-success"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("1599 millis")
     expect(s.requests).toHaveLength(1)
     yield* TestClock.adjust("801 millis")
@@ -4379,11 +4387,7 @@ describe("SessionRunnerLLM", () => {
   scenario("does not start another physical attempt after interruption during retry backoff", function* (s) {
     yield* s.admit("Interrupt retry backoff")
     yield* s.llm.push(Stream.fail(providerUnavailable()), TestLLM.text("Must not run", "unused-retry"))
-    const scheduled = yield* s.bus.subscribe(SessionEvent.RetryScheduled).pipe(
-      Stream.filter((event) => event.data.sessionID === sessionID),
-      Stream.runHead,
-      Effect.forkScoped({ startImmediately: true }),
-    )
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
     yield* Fiber.join(scheduled)
     yield* s.session.interrupt(sessionID)
@@ -4425,8 +4429,9 @@ describe("SessionRunnerLLM", () => {
     yield* s.llm.push(Stream.fail(incompleteStream()))
     yield* s.llm.push(TestLLM.text("Recovered", "incomplete-stream-success"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("2400 millis")
     yield* Fiber.join(run)
 
@@ -4446,8 +4451,9 @@ describe("SessionRunnerLLM", () => {
     ])
     yield* s.llm.push(TestLLM.text("Recovered", "unknown-finish-success"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("2400 millis")
     yield* Fiber.join(run)
 
@@ -4464,8 +4470,9 @@ describe("SessionRunnerLLM", () => {
     yield* s.llm.push(Stream.fail(rateLimited(5_000)))
     yield* s.llm.push(TestLLM.text("Recovered", "retry-after-success"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("4999 millis")
     expect(s.requests).toHaveLength(1)
     yield* TestClock.adjust("1 millis")
@@ -4478,8 +4485,9 @@ describe("SessionRunnerLLM", () => {
     yield* s.llm.push(Stream.fail(rateLimited(3_600_000)))
     yield* s.llm.push(TestLLM.text("Recovered", "retry-cap-success"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("899999 millis")
     expect(s.requests).toHaveLength(1)
     yield* TestClock.adjust("1 millis")
@@ -4500,8 +4508,9 @@ describe("SessionRunnerLLM", () => {
     )
     yield* s.llm.push(TestLLM.text(" continuation", "continued-text"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("2400 millis")
     yield* Fiber.join(run)
 
@@ -4549,8 +4558,9 @@ describe("SessionRunnerLLM", () => {
     ])
     yield* s.llm.push(TestLLM.text(" continuation", "unknown-continuation"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("2400 millis")
     yield* Fiber.join(run)
 
@@ -4579,8 +4589,9 @@ describe("SessionRunnerLLM", () => {
     )
     yield* s.llm.push(TestLLM.text(" continuation", "rate-limit-continuation"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("4999 millis")
     expect(s.requests).toHaveLength(1)
     yield* TestClock.adjust("1 millis")
@@ -4617,8 +4628,9 @@ describe("SessionRunnerLLM", () => {
     )
     yield* s.llm.push(TestLLM.text(" continuation", "unknown-failure-continuation"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("2400 millis")
     yield* Fiber.join(run)
 
@@ -4647,8 +4659,9 @@ describe("SessionRunnerLLM", () => {
     )
     yield* s.llm.push(TestLLM.text("Recovered", "reasoning-recovery"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("2400 millis")
     yield* Fiber.join(run)
 
@@ -4689,8 +4702,9 @@ describe("SessionRunnerLLM", () => {
     )
     yield* s.llm.push(TestLLM.text("Recovered", "reasoning-transport-recovery"))
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("2400 millis")
     yield* Fiber.join(run)
 
@@ -4831,11 +4845,16 @@ describe("SessionRunnerLLM", () => {
       ),
     )
 
+    const scheduled = yield* Queue.unbounded<void>()
+    yield* s.bus.subscribe(SessionEvent.RetryScheduled).pipe(
+      Stream.filter((event) => event.data.sessionID === sessionID),
+      Stream.runForEach(() => Queue.offer(scheduled, undefined)),
+      Effect.forkScoped({ startImmediately: true }),
+    )
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
-    for (const [index, delay] of [2_400, 4_800, 9_600, 19_200].entries()) {
+    for (const delay of [2_400, 4_800, 9_600, 19_200]) {
+      yield* Queue.take(scheduled)
       yield* TestClock.adjust(delay)
-      yield* s.llm.wait(index + 2)
     }
     expect(yield* Fiber.join(run).pipe(Effect.flip)).toBe(failure)
     expect(s.requests).toHaveLength(5)
@@ -4849,11 +4868,16 @@ describe("SessionRunnerLLM", () => {
     const failure = providerUnavailable()
     yield* s.llm.always(Stream.fail(failure))
 
+    const scheduled = yield* Queue.unbounded<void>()
+    yield* s.bus.subscribe(SessionEvent.RetryScheduled).pipe(
+      Stream.filter((event) => event.data.sessionID === sessionID),
+      Stream.runForEach(() => Queue.offer(scheduled, undefined)),
+      Effect.forkScoped({ startImmediately: true }),
+    )
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
-    for (const [index, delay] of [2_400, 4_800, 9_600, 19_200].entries()) {
+    for (const delay of [2_400, 4_800, 9_600, 19_200]) {
+      yield* Queue.take(scheduled)
       yield* TestClock.adjust(delay)
-      yield* s.llm.wait(index + 2)
     }
     expect(yield* Fiber.join(run).pipe(Effect.flip)).toBe(failure)
     expect(s.requests).toHaveLength(5)
@@ -4898,8 +4922,9 @@ describe("SessionRunnerLLM", () => {
     yield* s.llm.push(Stream.fail(failure))
     yield* s.llm.push(TestLLM.tool("call-after-retry", "echo", { text: "recovered" }), TestLLM.stop())
 
+    const scheduled = yield* nextRetryScheduled(s)
     const run = yield* s.resume.pipe(Effect.forkChild)
-    yield* s.llm.wait(1)
+    yield* Fiber.join(scheduled)
     yield* TestClock.adjust("2400 millis")
     yield* Fiber.join(run)
 

packages/core/test/session-step-machine.test.ts

@@ -0,0 +1,527 @@
+import { describe, expect, test } from "bun:test"
+import { AIError, TransportError, type LLMEvent } from "@opencode-ai/ai"
+import { SessionMessage } from "@opencode-ai/core/session/message"
+import { SessionStep } from "@opencode-ai/core/session/runner/step"
+import { SessionStepMachine } from "@opencode-ai/core/session/runner/step-machine"
+import { Cause, Deferred, Effect, Exit, Fiber, Ref, Scheduler } from "effect"
+import { it } from "./lib/effect"
+
+const firstID = SessionMessage.ID.make("msg_first")
+const failure = new AIError({
+  reason: new TransportError({ message: "Provider unavailable", transport: "http", operation: "request" }),
+})
+const error = { type: "provider.transport", message: "Provider unavailable" } as const
+describe("SessionStepMachine", () => {
+  it.effect("completes a logical Step", () =>
+    Effect.gen(function* () {
+      const attempts = yield* Ref.make<ReadonlyArray<SessionStepMachine.Context>>([])
+      const result = yield* SessionStepMachine.run(firstID, {
+        prepare: (context) =>
+          Ref.update(attempts, (values) => [...values, context]).pipe(
+            Effect.as(
+              SessionStepMachine.Preparation.Ready({
+                attempt: makeAttempt(SessionStep.Outcome.Completed({ needsContinuation: true })),
+              }),
+            ),
+          ),
+        retry: () => Effect.void,
+        publishSynthetic: Effect.void,
+      })
+      expect(result).toBe(true)
+      expect(yield* Ref.get(attempts)).toEqual([
+        { assistantMessageID: firstID, recoverOverflow: true, recoverContinuation: true },
+      ])
+    }),
+  )
+
+  it.effect("pulls, publishes, and runs a local tool before settlement", () =>
+    Effect.gen(function* () {
+      const operations = yield* Ref.make<ReadonlyArray<string>>([])
+      const call = { type: "tool-call", id: "call_1", name: "lookup", input: {} } satisfies Extract<
+        LLMEvent,
+        { type: "tool-call" }
+      >
+      const attempt = makeAttempt(SessionStep.Outcome.Completed({ needsContinuation: false }), {
+        events: [call],
+        operations,
+      })
+      yield* SessionStepMachine.run(firstID, {
+        prepare: () => Effect.succeed(SessionStepMachine.Preparation.Ready({ attempt })),
+        retry: () => Effect.void,
+        publishSynthetic: Effect.void,
+      })
+      const observed = yield* Ref.get(operations)
+      expect(observed.indexOf("publish:tool-call")).toBeLessThan(observed.indexOf("tool:call_1"))
+      expect(observed.at(-1)).toBe("settle")
+    }),
+  )
+
+  it.effect("retries transparently with the same assistant", () =>
+    Effect.gen(function* () {
+      const outcomes: Array<SessionStep.Outcome> = [
+        SessionStep.Outcome.Retry({ cause: failure, error }),
+        SessionStep.Outcome.Completed({ needsContinuation: false }),
+      ]
+      const operations = yield* Ref.make<ReadonlyArray<string>>([])
+      const result = yield* SessionStepMachine.run(firstID, {
+        prepare: (context) =>
+          Ref.update(operations, (values) => [...values, `attempt:${context.assistantMessageID}`]).pipe(
+            Effect.map(() =>
+              SessionStepMachine.Preparation.Ready({
+                attempt: makeAttempt(outcomes.shift() ?? SessionStep.Outcome.Completed({ needsContinuation: false })),
+              }),
+            ),
+          ),
+        retry: (context) => Ref.update(operations, (values) => [...values, `retry:${context.assistantMessageID}`]),
+        publishSynthetic: Effect.void,
+      })
+      expect(result).toBe(false)
+      expect(yield* Ref.get(operations)).toEqual([`attempt:${firstID}`, `retry:${firstID}`, `attempt:${firstID}`])
+    }),
+  )
+
+  it.effect("continues partial output only after retry and synthetic publication", () =>
+    Effect.gen(function* () {
+      const outcomes: Array<SessionStep.Outcome> = [
+        SessionStep.Outcome.Continue({ cause: failure, error }),
+        SessionStep.Outcome.Completed({ needsContinuation: false }),
+      ]
+      const operations = yield* Ref.make<ReadonlyArray<string>>([])
+      yield* SessionStepMachine.run(firstID, {
+        prepare: (context) =>
+          Ref.update(operations, (values) => [...values, `attempt:${context.assistantMessageID}`]).pipe(
+            Effect.map(() =>
+              SessionStepMachine.Preparation.Ready({
+                attempt: makeAttempt(outcomes.shift() ?? SessionStep.Outcome.Completed({ needsContinuation: false })),
+              }),
+            ),
+          ),
+        retry: () => Ref.update(operations, (values) => [...values, "retry"]),
+        publishSynthetic: Ref.update(operations, (values) => [...values, "synthetic"]),
+      })
+      const observed = yield* Ref.get(operations)
+      expect(observed.slice(0, 3)).toEqual([`attempt:${firstID}`, "retry", "synthetic"])
+      expect(observed.at(3)).toStartWith("attempt:msg_")
+      expect(observed.at(3)).not.toBe(`attempt:${firstID}`)
+    }),
+  )
+
+  it.effect("tracks independent recovery allowances", () =>
+    Effect.gen(function* () {
+      const outcomes: Array<SessionStep.Outcome> = [
+        SessionStep.Outcome.RecoverFull(),
+        SessionStep.Outcome.Completed({ needsContinuation: false }),
+        SessionStep.Outcome.Completed({ needsContinuation: false }),
+      ]
+      const recoveries = [false, true, false]
+      const attempts = yield* Ref.make<ReadonlyArray<SessionStepMachine.Context>>([])
+      yield* SessionStepMachine.run(firstID, {
+        prepare: (context) =>
+          Ref.update(attempts, (values) => [...values, context]).pipe(
+            Effect.map(() =>
+              SessionStepMachine.Preparation.Ready({
+                attempt: makeAttempt(outcomes.shift() ?? SessionStep.Outcome.Completed({ needsContinuation: false }), {
+                  recoverOverflow: recoveries.shift(),
+                }),
+              }),
+            ),
+          ),
+        retry: () => Effect.void,
+        publishSynthetic: Effect.void,
+      })
+      const observed = yield* Ref.get(attempts)
+      expect(observed.slice(0, 2)).toEqual([
+        { assistantMessageID: firstID, recoverOverflow: true, recoverContinuation: true },
+        { assistantMessageID: firstID, recoverOverflow: true, recoverContinuation: false },
+      ])
+      expect(observed.at(2)).toMatchObject({ recoverOverflow: false, recoverContinuation: false })
+      expect(observed.at(2)?.assistantMessageID).not.toBe(firstID)
+    }),
+  )
+
+  it.effect("does not begin another attempt when retry is interrupted", () =>
+    Effect.gen(function* () {
+      const retryStarted = yield* Deferred.make<void>()
+      const retryFinalized = yield* Deferred.make<void>()
+      const attempts = yield* Ref.make(0)
+      const machine = yield* SessionStepMachine.run(firstID, {
+        prepare: () =>
+          Ref.update(attempts, (value) => value + 1).pipe(
+            Effect.as(
+              SessionStepMachine.Preparation.Ready({
+                attempt: makeAttempt(SessionStep.Outcome.Retry({ cause: failure, error })),
+              }),
+            ),
+          ),
+        retry: () =>
+          Deferred.succeed(retryStarted, undefined).pipe(
+            Effect.andThen(Effect.never),
+            Effect.ensuring(Deferred.succeed(retryFinalized, undefined)),
+          ),
+        publishSynthetic: Effect.void,
+      }).pipe(Effect.forkChild({ startImmediately: true }))
+
+      yield* Deferred.await(retryStarted)
+      yield* Fiber.interrupt(machine)
+      expect(Exit.hasInterrupts(yield* Fiber.await(machine))).toBe(true)
+      expect(yield* Deferred.isDone(retryFinalized)).toBe(true)
+      expect(yield* Ref.get(attempts)).toBe(1)
+    }),
+  )
+
+  for (const outcome of [
+    SessionStep.Outcome.Completed({ needsContinuation: true }),
+    SessionStep.Outcome.Retry({ cause: failure, error }),
+    SessionStep.Outcome.Continue({ cause: failure, error }),
+    SessionStep.Outcome.RecoverFull(),
+  ]) {
+    it.effect(`cancellation during settlement prevents ${outcome._tag} from starting more work`, () =>
+      Effect.gen(function* () {
+        const started = yield* Deferred.make<void>()
+        const release = yield* Deferred.make<void>()
+        const operations = yield* Ref.make<ReadonlyArray<string>>([])
+        const attempt = {
+          ...makeAttempt(outcome),
+          settle: () =>
+            Deferred.succeed(started, undefined).pipe(
+              Effect.andThen(Deferred.await(release)),
+              Effect.andThen(Ref.update(operations, (values) => [...values, "settled"])),
+              Effect.as(outcome),
+              Effect.uninterruptible,
+            ),
+        }
+        const machine = yield* SessionStepMachine.run(firstID, {
+          prepare: () =>
+            Ref.update(operations, (values) => [...values, "prepare"]).pipe(
+              Effect.as(SessionStepMachine.Preparation.Ready({ attempt })),
+            ),
+          retry: () => Ref.update(operations, (values) => [...values, "retry"]),
+          publishSynthetic: Ref.update(operations, (values) => [...values, "synthetic"]),
+        }).pipe(Effect.forkChild({ startImmediately: true }))
+
+        yield* Deferred.await(started)
+        const interrupted = yield* Fiber.interrupt(machine).pipe(Effect.forkChild({ startImmediately: true }))
+        yield* Deferred.succeed(release, undefined)
+        yield* Fiber.join(interrupted)
+        expect(Exit.hasInterrupts(yield* Fiber.await(machine))).toBe(true)
+        expect(yield* Ref.get(operations)).toEqual(["prepare", "settled"])
+      }),
+    )
+  }
+
+  it.effect("cancels provider and tools together, then closes and settles once", () =>
+    Effect.gen(function* () {
+      const providerStarted = yield* Deferred.make<void>()
+      const providerStopped = yield* Deferred.make<void>()
+      const toolStarted = yield* Deferred.make<void>()
+      const toolStopped = yield* Deferred.make<void>()
+      const operations = yield* Ref.make<ReadonlyArray<string>>([])
+      const calls = [{ type: "tool-call", id: "call_parallel", name: "lookup", input: {} }] as const
+      const pending = [...calls]
+      const attempt: SessionStep.Attempt = {
+        ...makeAttempt(SessionStep.Outcome.Completed({ needsContinuation: false }), { operations }),
+        observeUntilBoundary: () =>
+          Effect.suspend(() => {
+            const call = pending.shift()
+            if (call) return Effect.succeed(SessionStep.ProviderObservation.ToolCall({ call }))
+            return Deferred.succeed(providerStarted, undefined).pipe(
+              Effect.andThen(Effect.never),
+              Effect.ensuring(
+                Deferred.succeed(providerStopped, undefined).pipe(Effect.andThen(Deferred.await(toolStopped))),
+              ),
+            )
+          }),
+        runTool: () =>
+          Deferred.succeed(toolStarted, undefined).pipe(
+            Effect.andThen(Effect.never),
+            Effect.ensuring(
+              Deferred.succeed(toolStopped, undefined).pipe(Effect.andThen(Deferred.await(providerStopped))),
+            ),
+          ),
+        settle: (settlement) =>
+          Effect.sync(() => {
+            expect(Exit.hasInterrupts(settlement.stream)).toBe(true)
+            expect(settlement.tools).toHaveLength(1)
+            expect(settlement.tools[0]?.call).toEqual(calls[0])
+            expect(settlement.tools.every((tool) => Exit.hasInterrupts(tool.exit))).toBe(true)
+          }).pipe(
+            Effect.andThen(Ref.update(operations, (values) => [...values, "settle"])),
+            Effect.as(SessionStep.Outcome.Completed({ needsContinuation: false })),
+          ),
+      }
+      const machine = yield* SessionStepMachine.run(firstID, {
+        prepare: () => Effect.succeed(SessionStepMachine.Preparation.Ready({ attempt })),
+        retry: () => Effect.die("Unexpected retry"),
+        publishSynthetic: Effect.die("Unexpected continuation"),
+      }).pipe(Effect.forkChild({ startImmediately: true }))
+
+      yield* Deferred.await(providerStarted)
+      yield* Deferred.await(toolStarted)
+      yield* Fiber.interrupt(machine)
+      expect(Exit.hasInterrupts(yield* Fiber.await(machine))).toBe(true)
+      expect(yield* Ref.get(operations)).toEqual(["finish-provider", "settle"])
+    }),
+  )
+
+  it.effect("does not finalize the provider twice when cancellation races with finalization", () =>
+    Effect.gen(function* () {
+      const started = yield* Deferred.make<void>()
+      const release = yield* Deferred.make<void>()
+      const operations = yield* Ref.make<ReadonlyArray<string>>([])
+      const attempt = {
+        ...makeAttempt(SessionStep.Outcome.Completed({ needsContinuation: false }), { operations }),
+        finishProvider: () =>
+          Deferred.succeed(started, undefined).pipe(
+            Effect.andThen(Deferred.await(release)),
+            Effect.andThen(Ref.update(operations, (values) => [...values, "finish-provider"])),
+            Effect.uninterruptible,
+          ),
+      }
+      const machine = yield* SessionStepMachine.run(firstID, {
+        prepare: () => Effect.succeed(SessionStepMachine.Preparation.Ready({ attempt })),
+        retry: () => Effect.die("Unexpected retry"),
+        publishSynthetic: Effect.die("Unexpected continuation"),
+      }).pipe(Effect.forkChild({ startImmediately: true }))
+
+      yield* Deferred.await(started)
+      const interrupted = yield* Fiber.interrupt(machine).pipe(Effect.forkChild({ startImmediately: true }))
+      yield* Deferred.succeed(release, undefined)
+      yield* Fiber.join(interrupted)
+      expect(Exit.hasInterrupts(yield* Fiber.await(machine))).toBe(true)
+      expect(yield* Ref.get(operations)).toEqual(["read:end", "finish-provider", "settle"])
+    }),
+  )
+
+  test("cancellation awaits provider finalization and stops pending tools before settling", () => {
+    const definition = SessionStepMachine.definition<never, never>(firstID)
+    const cause = Cause.interrupt(123)
+    const call = { type: "tool-call", id: "call_pending", name: "lookup", input: {} } as const
+    const completed = { ...call, id: "call_completed" }
+    const state = SessionStepMachine.State.FinalizingProvider({
+      active: {
+        context: { assistantMessageID: firstID, recoverOverflow: true, recoverContinuation: true },
+        attempt: makeAttempt(SessionStep.Outcome.Completed({ needsContinuation: false })),
+        tools: new Map([
+          [completed.id, { call: completed, exit: Exit.succeed(undefined) }],
+          [call.id, { call }],
+        ]),
+      },
+      stream: Exit.succeed(undefined),
+    })
+    const stopping = definition.transition(state, {
+      _tag: "Input",
+      input: SessionStepMachine.Event.CancelRequested(),
+      cause,
+    })
+    if (stopping._tag !== "Continue") throw new Error("Expected cancellation to await owned invocations")
+    expect(stopping.state).toEqual({ _tag: "Stopping", from: state, cause })
+    expect(stopping.commands).toEqual([
+      { _tag: "StopAndJoin", id: "step", ids: ["tool:call_pending"], waitFor: ["provider"] },
+    ])
+    expect(
+      definition.transition(stopping.state, {
+        _tag: "Input",
+        input: SessionStepMachine.Event.CancelRequested(),
+        cause,
+      }),
+    ).toEqual({ _tag: "Continue", state: stopping.state, commands: [] })
+
+    const settling = definition.transition(stopping.state, {
+      _tag: "InvocationsStopped",
+      id: "step",
+      exits: [
+        {
+          _tag: "InvocationExited",
+          id: "tool:call_pending",
+          generation: 1,
+          operation: SessionStepMachine.Operation.RunTool({ attempt: state.active.attempt, call }),
+          exit: Exit.interrupt(456),
+        },
+        {
+          _tag: "InvocationExited",
+          id: "provider",
+          generation: 2,
+          operation: SessionStepMachine.Operation.FinishProvider({
+            attempt: state.active.attempt,
+            stream: state.stream,
+          }),
+          exit: Exit.succeed(SessionStepMachine.Event.ProviderFinished({ exit: Exit.succeed(undefined) })),
+        },
+      ],
+    })
+    if (settling._tag !== "Continue") throw new Error("Expected settlement after the joined batch")
+    expect(settling.state).toMatchObject({ _tag: "SettlingAttempt", stopping: cause })
+    expect(settling.commands).toEqual([
+      {
+        _tag: "Invoke",
+        id: "settlement",
+        operation: {
+          _tag: "SettleAttempt",
+          attempt: state.active.attempt,
+          settlement: {
+            stream: state.stream,
+            tools: [
+              { call: completed, exit: Exit.succeed(undefined) },
+              { call, exit: Exit.interrupt(456) },
+            ],
+          },
+        },
+      },
+    ])
+  })
+
+  for (const fixture of [
+    { name: "never-started", exit: Exit.interrupt(456), replaced: false },
+    {
+      name: "queued false",
+      exit: Exit.succeed(SessionStepMachine.Event.OverflowRecovered({ exit: Exit.succeed(false) })),
+      replaced: false,
+    },
+    {
+      name: "queued failure",
+      exit: Exit.succeed(SessionStepMachine.Event.OverflowRecovered({ exit: Exit.die("Recovery failed") })),
+      replaced: false,
+    },
+    {
+      name: "queued true",
+      exit: Exit.succeed(SessionStepMachine.Event.OverflowRecovered({ exit: Exit.succeed(true) })),
+      replaced: true,
+    },
+  ] as const) {
+    test(`cancellation reconciles ${fixture.name} overflow recovery before deciding settlement`, () => {
+      const definition = SessionStepMachine.definition<never, never>(firstID)
+      const cause = Cause.interrupt(123)
+      const state = SessionStepMachine.State.RecoveringOverflow({
+        active: {
+          context: { assistantMessageID: firstID, recoverOverflow: true, recoverContinuation: true },
+          attempt: makeAttempt(SessionStep.Outcome.Completed({ needsContinuation: true })),
+          tools: new Map(),
+        },
+        stream: Exit.succeed(undefined),
+      })
+      const stopping = definition.transition(state, {
+        _tag: "Input",
+        input: SessionStepMachine.Event.CancelRequested(),
+        cause,
+      })
+      if (stopping._tag !== "Continue") throw new Error("Expected cancellation to await recovery")
+      expect(stopping.state).toEqual({ _tag: "Stopping", from: state, cause })
+      expect(stopping.commands).toEqual([{ _tag: "StopAndJoin", id: "step", ids: ["compaction"], waitFor: [] }])
+
+      const settled = definition.transition(stopping.state, {
+        _tag: "InvocationsStopped",
+        id: "step",
+        exits: [
+          {
+            _tag: "InvocationExited",
+            id: "compaction",
+            generation: 1,
+            operation: SessionStepMachine.Operation.RecoverOverflow({
+              attempt: state.active.attempt,
+              settlement: { stream: state.stream, tools: [] },
+            }),
+            exit: fixture.exit,
+          },
+        ],
+      })
+      if (fixture.replaced) {
+        expect(settled).toEqual({ _tag: "Done", output: Exit.failCause(cause) })
+        return
+      }
+      if (settled._tag !== "Continue") throw new Error("Expected the unreplaced attempt to settle")
+      expect(settled.state).toEqual({ _tag: "SettlingAttempt", active: state.active, stopping: cause })
+      expect(settled.commands).toEqual([
+        {
+          _tag: "Invoke",
+          id: "settlement",
+          operation: {
+            _tag: "SettleAttempt",
+            attempt: state.active.attempt,
+            settlement: { stream: Exit.failCause(cause), tools: [] },
+          },
+        },
+      ])
+      const command = settled.commands[0]
+      if (command?._tag !== "Invoke") throw new Error("Expected a settlement invocation")
+      expect(
+        definition.transition(settled.state, {
+          _tag: "InvocationExited",
+          id: command.id,
+          generation: 2,
+          operation: command.operation,
+          exit: Exit.succeed(
+            SessionStepMachine.Event.AttemptSettled({
+              exit: Exit.succeed(SessionStep.Outcome.Completed({ needsContinuation: true })),
+            }),
+          ),
+        }),
+      ).toEqual({ _tag: "Done", output: Exit.failCause(cause) })
+    })
+  }
+
+  for (const target of ["finishProvider", "recoverOverflow"] as const) {
+    it.effect(`settles once when cancellation precedes ${target} execution`, () =>
+      Effect.gen(function* () {
+        const operations = yield* Ref.make<ReadonlyArray<string>>([])
+        const attempt = makeAttempt(SessionStep.Outcome.Completed({ needsContinuation: true }), { operations })
+        const machine = yield* Effect.withFiber((fiber) =>
+          SessionStepMachine.run(firstID, {
+            prepare: () =>
+              Effect.succeed(
+                SessionStepMachine.Preparation.Ready({
+                  attempt: {
+                    ...attempt,
+                    // Interrupt during construction, before the deferred invocation starts.
+                    finishProvider: (stream) => {
+                      if (target === "finishProvider") fiber.interruptUnsafe(123)
+                      return attempt.finishProvider(stream)
+                    },
+                    recoverOverflow: (settlement) => {
+                      if (target === "recoverOverflow") fiber.interruptUnsafe(123)
+                      return attempt.recoverOverflow(settlement)
+                    },
+                  },
+                }),
+              ),
+            retry: () => Effect.die("Unexpected retry"),
+            publishSynthetic: Effect.die("Unexpected continuation"),
+          }),
+        ).pipe(
+          Effect.provideService(Scheduler.PreventSchedulerYield, true),
+          Effect.forkChild({ startImmediately: true }),
+        )
+
+        expect(Exit.hasInterrupts(yield* Fiber.await(machine))).toBe(true)
+        expect(yield* Ref.get(operations)).toEqual(["read:end", "finish-provider", "settle"])
+      }),
+    )
+  }
+})
+
+function makeAttempt(
+  outcome: SessionStep.Outcome,
+  options?: {
+    readonly events?: ReadonlyArray<LLMEvent>
+    readonly operations?: Ref.Ref<ReadonlyArray<string>>
+    readonly recoverOverflow?: boolean
+  },
+): SessionStep.Attempt {
+  const events = [...(options?.events ?? [])]
+  const log = (value: string) =>
+    options?.operations ? Ref.update(options.operations, (values) => [...values, value]) : Effect.void
+  return {
+    observeUntilBoundary: () =>
+      Effect.gen(function* () {
+        const event = events.shift()
+        yield* log(event ? `read:${event.type}` : "read:end")
+        if (!event) return SessionStep.ProviderObservation.ProviderEnd()
+        yield* log(`publish:${event.type}`)
+        if (event.type !== "tool-call") return SessionStep.ProviderObservation.ProviderEnd()
+        return SessionStep.ProviderObservation.ToolCall({ call: event })
+      }),
+    runTool: (call) => log(`tool:${call.id}`),
+    finishProvider: () => log("finish-provider"),
+    recoverOverflow: () => Effect.succeed(options?.recoverOverflow ?? false),
+    settle: () => log("settle").pipe(Effect.as(outcome)),
+  }
+}

packages/core/test/session-step.test.ts

@@ -1,5 +1,5 @@
 import { expect } from "bun:test"
-import { LanguageModel, LLM, LLMEvent } from "@opencode-ai/ai"
+import { AIError, LanguageModel, LLM, LLMEvent, TransportError } from "@opencode-ai/ai"
 import { OpenAIChat } from "@opencode-ai/ai/protocols/openai-chat"
 import { TestLLM } from "@opencode-ai/ai/testing"
 import { Agent } from "@opencode-ai/core/agent"
@@ -11,17 +11,19 @@ import { Project } from "@opencode-ai/core/project"
 import { ProjectTable } from "@opencode-ai/core/project/sql"
 import { AbsolutePath, RelativePath } from "@opencode-ai/core/schema"
 import { Session } from "@opencode-ai/core/session"
+import { SessionEvent } from "@opencode-ai/core/session/event"
 import { SessionMessage } from "@opencode-ai/core/session/message"
 import { SessionProjector } from "@opencode-ai/core/session/projector"
 import { SessionRunnerModel } from "@opencode-ai/core/session/runner/model"
 import { SessionStep } from "@opencode-ai/core/session/runner/step"
+import { SessionStepMachine } from "@opencode-ai/core/session/runner/step-machine"
 import { SessionMessageTable, SessionTable } from "@opencode-ai/core/session/sql"
 import { Snapshot } from "@opencode-ai/core/snapshot"
 import { ToolOutput } from "@opencode-ai/core/tool-output"
 import { Money } from "@opencode-ai/schema/money"
 import { LayerNode } from "@opencode-ai/util/effect/layer-node"
 import { asc, eq } from "drizzle-orm"
-import { Effect, Exit, Layer } from "effect"
+import { Deferred, Effect, Exit, Fiber, Layer, Stream } from "effect"
 import { testEffect } from "./lib/effect"
 
 const it = testEffect(
@@ -40,49 +42,21 @@ for (const fixture of [
 ] as const) {
   it.effect(`settles ${fixture.finish} with tool choice ${fixture.toolChoice ?? "default"}`, () =>
     Effect.gen(function* () {
-      const db = (yield* Database.Service).db
-      const llm = yield* TestLLM.Test
-      const sessionID = Session.ID.create()
-      const assistantMessageID = SessionMessage.ID.create()
       const start = Snapshot.ID.make("before")
       const end = Snapshot.ID.make("after")
       const files = [RelativePath.make("changed.ts")]
       let captures = 0
       let executions = 0
-      const steps = yield* SessionStep.make.pipe(
-        Effect.provide(
-          Layer.mock(Snapshot.Service)({
-            capture: () => Effect.sync(() => (captures++ === 0 ? start : end)),
-            files: (input) => {
-              expect(input).toEqual({ from: start, to: end })
-              return Effect.succeed(files)
-            },
-          }),
-        ),
-      )
-      yield* db
-        .insert(ProjectTable)
-        .values({ id: Project.ID.global, worktree: AbsolutePath.make("/project"), sandboxes: [] })
-        .run()
-      yield* db
-        .insert(SessionTable)
-        .values({ id: sessionID, project_id: Project.ID.global, slug: "step", directory: "/project", version: "test" })
-        .run()
-      const model = SessionRunnerModel.resolved(
-        LanguageModel.make({ id: "test-model", provider: "test", route: OpenAIChat.route }),
-        {
-          capabilities: { tools: true, input: ["text"], output: ["text"] },
-          limit: { context: 100_000, output: 1_000 },
-          cost: [
-            {
-              input: Money.USDPerMillionTokens.make(1),
-              output: Money.USDPerMillionTokens.make(2),
-              cache: { read: Money.USDPerMillionTokens.make(0.1), write: Money.USDPerMillionTokens.make(0.5) },
-            },
-          ],
+      const s = yield* setup({
+        snapshot: {
+          capture: () => Effect.sync(() => (captures++ === 0 ? start : end)),
+          files: (input) => {
+            expect(input).toEqual({ from: start, to: end })
+            return Effect.succeed(files)
+          },
         },
-      )
-      yield* llm.push(
+      })
+      yield* s.llm.push(
         TestLLM.complete(
           {
             reason: { normalized: fixture.finish },
@@ -98,52 +72,33 @@ for (const fixture of [
           LLMEvent.toolCall({ id: "call-test", name: "test", input: {} }),
         ),
       )
-      const result = yield* steps
-        .attempt({
-          sessionID,
-          assistantMessageID,
-          agent: Agent.defaultID,
-          model,
-          prepared: {
-            request: LLM.request({ model: model.model, prompt: "Run one tool", toolChoice: fixture.toolChoice }),
-            options: {},
+      const result = yield* SessionStepMachine.run(s.assistantMessageID, {
+        prepare: (context) =>
+          s.prepare(context, {
+            toolChoice: fixture.toolChoice,
             executeTool: () =>
               Effect.sync(() => {
                 executions++
                 return { content: "Completed tool" }
               }),
-          },
-          recoverContinuation: true,
-          recoverOverflow: Effect.succeed(false),
-        })
-        .pipe(Effect.exit)
+          }),
+        retry: () => Effect.die("Unexpected retry"),
+        publishSynthetic: Effect.die("Unexpected continuation"),
+      }).pipe(Effect.exit)
       expect(Exit.isSuccess(result)).toBe(fixture.finish === "stop")
       expect(executions).toBe(fixture.toolChoice === "none" ? 0 : 1)
-      if (Exit.isSuccess(result))
-        expect(result.value).toEqual(
-          SessionStep.Outcome.Completed({ needsContinuation: fixture.toolChoice !== "none" }),
-        )
-      expect(yield* llm.requests()).toHaveLength(1)
+      if (Exit.isSuccess(result)) expect(result.value).toBe(fixture.toolChoice !== "none")
+      expect(yield* s.llm.requests()).toHaveLength(1)
       expect(captures).toBe(2)
-      const message = yield* db
-        .select()
-        .from(SessionMessageTable)
-        .where(eq(SessionMessageTable.id, assistantMessageID))
-        .get()
-      expect(message?.data).toMatchObject({
+      const message = yield* s.message
+      expect(message).toMatchObject({
         finish: fixture.finish,
         tokens: { input: 10, output: 4, reasoning: 2, cache: { read: 3, write: 2 } },
         snapshot: { start, end, files },
         content: [{ type: "tool", state: { status: fixture.toolChoice === "none" ? "error" : "completed" } }],
       })
-      expect(message?.data).toHaveProperty("cost", expect.closeTo(0.0000233, 10))
-      const events = yield* db
-        .select({ type: EventTable.type })
-        .from(EventTable)
-        .where(eq(EventTable.aggregate_id, sessionID))
-        .orderBy(asc(EventTable.seq))
-        .all()
-      const types = events.map((event) => event.type)
+      expect(message).toHaveProperty("cost", expect.closeTo(0.0000233, 10))
+      const types = yield* s.events
       const terminal = fixture.finish === "stop" ? "session.step.ended.1" : "session.step.failed.1"
       expect(types.filter((type) => type === terminal)).toHaveLength(1)
       expect(
@@ -152,3 +107,285 @@ for (const fixture of [
     }),
   )
 }
+
+it.effect("closes provider stream resources before the next physical retry", () =>
+  Effect.gen(function* () {
+    const s = yield* setup()
+    const cleanupStarted = yield* Deferred.make<void>()
+    const cleanupRelease = yield* Deferred.make<void>()
+    const operations: string[] = []
+    yield* s.llm.push(
+      Stream.unwrap(
+        Effect.acquireRelease(
+          Effect.sync(() => operations.push("acquire")),
+          () =>
+            Deferred.succeed(cleanupStarted, undefined).pipe(
+              Effect.andThen(Deferred.await(cleanupRelease)),
+              Effect.andThen(Effect.sync(() => operations.push("release"))),
+            ),
+        ).pipe(
+          Effect.as(
+            Stream.fail(
+              new AIError({
+                reason: new TransportError({ message: "Request failed", transport: "http", operation: "request" }),
+              }),
+            ),
+          ),
+        ),
+      ),
+      TestLLM.stop(),
+    )
+    const run = yield* SessionStepMachine.run(s.assistantMessageID, {
+      prepare: (context) => Effect.sync(() => operations.push("prepare")).pipe(Effect.andThen(s.prepare(context))),
+      retry: () => Effect.sync(() => operations.push("retry")).pipe(Effect.asVoid),
+      publishSynthetic: Effect.die("Unexpected continuation"),
+    }).pipe(Effect.forkScoped({ startImmediately: true }))
+    yield* Effect.addFinalizer(() => Deferred.succeed(cleanupRelease, undefined))
+    yield* Deferred.await(cleanupStarted)
+
+    expect(operations).toEqual(["prepare", "acquire"])
+    expect(yield* s.llm.requests()).toHaveLength(1)
+    expect(run.pollUnsafe()).toBeUndefined()
+    yield* Deferred.succeed(cleanupRelease, undefined)
+    expect(yield* Fiber.join(run)).toBe(false)
+    expect(operations).toEqual(["prepare", "acquire", "release", "retry", "prepare"])
+    expect(yield* s.llm.requests()).toHaveLength(2)
+    expect(yield* s.message).toMatchObject({ finish: "stop" })
+  }),
+)
+
+for (const providerExecuted of [false, true]) {
+  it.effect(
+    `commits ${providerExecuted ? "provider-hosted" : "local"} tool success during cancellation under the bus lock`,
+    () =>
+      Effect.gen(function* () {
+        const ready = yield* Deferred.make<void>()
+        const resultRelease = yield* Deferred.make<void>()
+        const publishing = yield* Deferred.make<void>()
+        const held = yield* Deferred.make<void>()
+        const lockRelease = yield* Deferred.make<void>()
+        const s = yield* setup({
+          observePublish: (type) =>
+            type === SessionEvent.Tool.Success.type ? Deferred.succeed(publishing, undefined) : Effect.void,
+        })
+        const call = LLMEvent.toolCall({ id: "call-race", name: "lookup", input: {}, providerExecuted })
+        let executions = 0
+        yield* s.llm.push(
+          providerExecuted
+            ? Stream.fromIterable([LLMEvent.stepStart({ index: 0 }), call]).pipe(
+                Stream.concat(
+                  Stream.unwrap(
+                    Deferred.succeed(ready, undefined).pipe(
+                      Effect.andThen(Deferred.await(resultRelease)),
+                      Effect.as(
+                        Stream.make(
+                          LLMEvent.toolResult({
+                            id: call.id,
+                            name: call.name,
+                            providerExecuted: true,
+                            result: { type: "text", value: "Durable result" },
+                          }),
+                        ),
+                      ),
+                    ),
+                  ),
+                ),
+                Stream.concat(Stream.never),
+              )
+            : TestLLM.hangAfter(LLMEvent.stepStart({ index: 0 }), call),
+        )
+        const run = yield* SessionStepMachine.run(s.assistantMessageID, {
+          prepare: (context) =>
+            s.prepare(context, {
+              executeTool: () =>
+                Effect.gen(function* () {
+                  executions++
+                  yield* Deferred.succeed(ready, undefined)
+                  yield* Deferred.await(resultRelease)
+                  return { content: "Durable result" }
+                }),
+            }),
+          retry: () => Effect.die("Unexpected retry"),
+          publishSynthetic: Effect.die("Unexpected continuation"),
+        }).pipe(Effect.forkScoped({ startImmediately: true }))
+        yield* Deferred.await(ready)
+        yield* Effect.acquireRelease(
+          s.bus.listen((event) =>
+            event.type === SessionEvent.Renamed.type
+              ? Deferred.succeed(held, undefined).pipe(Effect.andThen(Deferred.await(lockRelease)))
+              : Effect.void,
+          ),
+          (unsubscribe) => unsubscribe,
+        )
+        // Notifications hold the real aggregate lock after the unrelated event commits.
+        const holder = yield* s.bus
+          .publish(SessionEvent.Renamed, { sessionID: s.sessionID, title: "Hold publication" })
+          .pipe(Effect.forkScoped({ startImmediately: true }))
+        yield* Effect.addFinalizer(() => Deferred.succeed(lockRelease, undefined))
+        yield* Deferred.await(held)
+        yield* Deferred.succeed(resultRelease, undefined)
+        yield* Deferred.await(publishing)
+        const cancellation = yield* Fiber.interrupt(run).pipe(Effect.forkChild({ startImmediately: true }))
+        yield* Effect.yieldNow
+
+        expect(cancellation.pollUnsafe()).toBeUndefined()
+        expect(yield* s.events).not.toContain("session.tool.success.2")
+        yield* Deferred.succeed(lockRelease, undefined)
+        yield* Fiber.join(holder)
+        yield* Fiber.join(cancellation)
+        expect(Exit.hasInterrupts(yield* Fiber.await(run))).toBe(true)
+        expect(executions).toBe(providerExecuted ? 0 : 1)
+        expect(yield* s.llm.requests()).toHaveLength(1)
+        const events = yield* s.events
+        expect(events.filter((type) => type === "session.tool.success.2")).toHaveLength(1)
+        expect(events).not.toContain("session.tool.failed.2")
+        expect(events.filter((type) => type === "session.step.failed.1")).toHaveLength(1)
+        expect(events.indexOf("session.tool.success.2")).toBeLessThan(events.indexOf("session.step.failed.1"))
+        expect(yield* s.message).toMatchObject({
+          finish: "error",
+          error: { type: "aborted" },
+          content: [
+            {
+              type: "tool",
+              id: call.id,
+              executed: providerExecuted,
+              state: { status: "completed", content: [{ type: "text", text: "Durable result" }] },
+            },
+          ],
+        })
+      }),
+  )
+}
+
+it.effect("recovers overflow instead of generically retrying a subsequent transport failure", () =>
+  Effect.gen(function* () {
+    const s = yield* setup()
+    const contexts: SessionStepMachine.Context[] = []
+    const operations: string[] = []
+    yield* s.llm.push(
+      TestLLM.failAfter(
+        new AIError({
+          reason: new TransportError({ message: "Read failed", transport: "http", operation: "read" }),
+        }),
+        LLMEvent.stepStart({ index: 0 }),
+        LLMEvent.providerError({ message: "Prompt too long", classification: "context-overflow" }),
+      ),
+      TestLLM.stop(),
+    )
+    const result = yield* SessionStepMachine.run(s.assistantMessageID, {
+      prepare: (context) =>
+        Effect.sync(() => contexts.push(context)).pipe(
+          Effect.andThen(
+            s.prepare(context, {
+              recoverOverflow: Effect.sync(() => {
+                operations.push("compact")
+                return true
+              }),
+            }),
+          ),
+        ),
+      retry: () => Effect.sync(() => operations.push("retry")).pipe(Effect.asVoid),
+      publishSynthetic: Effect.die("Unexpected continuation"),
+    })
+
+    expect(result).toBe(false)
+    expect(operations).toEqual(["compact"])
+    expect(yield* s.llm.requests()).toHaveLength(2)
+    expect(contexts).toHaveLength(2)
+    expect(contexts[0]).toMatchObject({ assistantMessageID: s.assistantMessageID, recoverOverflow: true })
+    expect(contexts[1]).toMatchObject({ recoverOverflow: false })
+    expect(contexts[1]?.assistantMessageID).not.toBe(s.assistantMessageID)
+    expect(yield* s.events).not.toContain("session.step.failed.1")
+  }),
+)
+
+const setup = Effect.fnUntraced(function* (
+  options: {
+    readonly snapshot?: Pick<Snapshot.Interface, "capture" | "files">
+    readonly observePublish?: (type: string) => Effect.Effect<unknown>
+  } = {},
+) {
+  const db = (yield* Database.Service).db
+  const bus = yield* Bus.Service
+  const llm = yield* TestLLM.Test
+  const sessionID = Session.ID.create()
+  const assistantMessageID = SessionMessage.ID.create()
+  yield* db
+    .insert(ProjectTable)
+    .values({ id: Project.ID.global, worktree: AbsolutePath.make("/project"), sandboxes: [] })
+    .run()
+  yield* db
+    .insert(SessionTable)
+    .values({ id: sessionID, project_id: Project.ID.global, slug: "step", directory: "/project", version: "test" })
+    .run()
+  const model = SessionRunnerModel.resolved(
+    LanguageModel.make({ id: "test-model", provider: "test", route: OpenAIChat.route }),
+    {
+      capabilities: { tools: true, input: ["text"], output: ["text"] },
+      limit: { context: 100_000, output: 1_000 },
+      cost: [
+        {
+          input: Money.USDPerMillionTokens.make(1),
+          output: Money.USDPerMillionTokens.make(2),
+          cache: { read: Money.USDPerMillionTokens.make(0.1), write: Money.USDPerMillionTokens.make(0.5) },
+        },
+      ],
+    },
+  )
+  const steps = yield* SessionStep.make.pipe(
+    Effect.provide(
+      Layer.mock(Snapshot.Service)(
+        options.snapshot ?? { capture: () => Effect.undefined, files: () => Effect.succeed([]) },
+      ),
+    ),
+    Effect.provideService(Bus.Service, {
+      ...bus,
+      publish: (definition, data, publishOptions) =>
+        (options.observePublish?.(definition.type) ?? Effect.void).pipe(
+          Effect.andThen(bus.publish(definition, data, publishOptions)),
+        ),
+    }),
+  )
+  return {
+    bus,
+    llm,
+    sessionID,
+    assistantMessageID,
+    prepare: (
+      context: SessionStepMachine.Context,
+      input?: {
+        readonly toolChoice?: "none"
+        readonly executeTool?: SessionStep.Input["prepared"]["executeTool"]
+        readonly recoverOverflow?: Effect.Effect<boolean>
+      },
+    ) =>
+      steps
+        .open({
+          sessionID,
+          assistantMessageID: context.assistantMessageID,
+          agent: Agent.defaultID,
+          model,
+          prepared: {
+            request: LLM.request({ model: model.model, prompt: "Run one step", toolChoice: input?.toolChoice }),
+            options: {},
+            executeTool: input?.executeTool ?? (() => Effect.die("Unexpected tool execution")),
+          },
+          recoverContinuation: context.recoverContinuation,
+          recoverOverflow: input?.recoverOverflow ?? Effect.succeed(false),
+        })
+        .pipe(Effect.map((attempt) => SessionStepMachine.Preparation.Ready({ attempt }))),
+    message: db
+      .select()
+      .from(SessionMessageTable)
+      .where(eq(SessionMessageTable.id, assistantMessageID))
+      .get()
+      .pipe(Effect.map((row) => row?.data)),
+    events: db
+      .select({ type: EventTable.type })
+      .from(EventTable)
+      .where(eq(EventTable.aggregate_id, sessionID))
+      .orderBy(asc(EventTable.seq))
+      .all()
+      .pipe(Effect.map((rows) => rows.map((row) => row.type))),
+  }
+})