refactor(core): clarify provider turn phases

packages/core/src/session/runner/run-turn.ts

@@ -1,5 +1,15 @@
 export * as RunTurn from "./run-turn"
 
+/**
+ * Drives one logical provider turn to settlement.
+ *
+ * A logical turn may rebuild its immutable preparation when concurrent Session,
+ * agent, model, or Context Epoch changes make a prepared request stale. Each
+ * prepared attempt invokes `llm.stream` at most once. A pre-output context
+ * overflow may compact and rebuild once; later rebuilds do not restore that
+ * recovery budget.
+ */
+
 import {
   LLM,
   LLMClient,
@@ -40,15 +50,10 @@ export type Run = (
   promotion: SessionInput.Delivery | undefined,
 ) => Effect.Effect<boolean, RunError>
 
-type TurnTransition =
-  | { readonly _tag: "RebuildPreparedTurn"; readonly promotion?: SessionInput.Delivery }
-  | { readonly _tag: "ContinueAfterOverflowCompaction" }
-
-class TurnTransitionError extends Error {
-  constructor(readonly transition: TurnTransition) {
-    super()
-  }
-}
+const TurnTransition = Schema.TaggedUnion({
+  RebuildPreparedTurn: { promotion: SessionInput.Delivery.pipe(Schema.optional) },
+  ContinueAfterOverflowCompaction: {},
+})
 
 export const make = Effect.gen(function* () {
   const events = yield* EventV2.Service
@@ -74,14 +79,12 @@ export const make = Effect.gen(function* () {
   const isQuestionRejected = (cause: Cause.Cause<unknown>) =>
     cause.reasons.some((reason) => Cause.isDieReason(reason) && reason.defect instanceof QuestionV2.RejectedError)
   const rebuildPreparedTurn = (promotion?: SessionInput.Delivery) =>
-    new TurnTransitionError({ _tag: "RebuildPreparedTurn", promotion })
-  const continueAfterOverflowCompaction = new TurnTransitionError({
-    _tag: "ContinueAfterOverflowCompaction",
-  })
+    TurnTransition.cases.RebuildPreparedTurn.make({ promotion })
+  const continueAfterOverflowCompaction = TurnTransition.cases.ContinueAfterOverflowCompaction.make({})
   const retryAgentMismatch = (promotion: SessionInput.Delivery | undefined) =>
     Effect.catchDefect((defect) =>
       defect instanceof SessionContextEpoch.AgentMismatch
-        ? Effect.die(rebuildPreparedTurn(promotion))
+        ? Effect.fail(rebuildPreparedTurn(promotion))
         : Effect.die(defect),
     )
   const sameModel = Schema.toEquivalence(Schema.UndefinedOr(ModelV2.Ref))
@@ -90,10 +93,15 @@ export const make = Effect.gen(function* () {
       Effect.map(SystemContext.combine),
     )
 
-  const runAttempt = Effect.fn("SessionRunner.runTurn")(function* (
+  /**
+   * Promotes admitted input and builds one coherent immutable request snapshot.
+   *
+   * Rebuild transitions before promotion preserve the requested delivery;
+   * transitions after promotion clear it so queued input cannot be promoted twice.
+   */
+  const prepareTurn = Effect.fn("SessionRunner.prepareTurn")(function* (
     sessionID: SessionSchema.ID,
     promotion: SessionInput.Delivery | undefined,
-    recoverOverflow?: typeof compaction.compactAfterOverflow,
   ) {
     const session = yield* getSession(sessionID)
     if (session.location.directory !== location.directory || session.location.workspaceID !== location.workspaceID)
@@ -106,8 +114,6 @@ export const make = Effect.gen(function* () {
       session.location,
       agent.id,
     ).pipe(retryAgentMismatch(promotion))
-    const toolFibers = yield* FiberSet.make<void, ToolOutputStore.Error>()
-    let needsContinuation = false
     if (promotion) {
       const cutoff = yield* SessionInput.latestSeq(db, session.id)
       if (promotion === "steer") yield* SessionInput.promoteSteers(db, events, session.id, cutoff)
@@ -128,7 +134,7 @@ export const make = Effect.gen(function* () {
       ).pipe(retryAgentMismatch(undefined)))
     const current = yield* getSession(sessionID)
     if ((yield* agents.select(current.agent)).id !== agent.id || !sameModel(current.model, session.model))
-      return yield* Effect.die(rebuildPreparedTurn())
+      return yield* Effect.fail(rebuildPreparedTurn())
     const model = yield* models.resolve(session)
     const entries = yield* SessionHistory.entriesForRunner(db, session.id, system.baselineSeq)
     const context = entries.map((entry) => entry.message)
@@ -144,47 +150,76 @@ export const make = Effect.gen(function* () {
       tools: toolMaterialization.definitions,
     })
     if (yield* compaction.compactIfNeeded({ sessionID: session.id, entries, model, request }))
-      return yield* Effect.die(rebuildPreparedTurn())
+      return yield* Effect.fail(rebuildPreparedTurn())
+    return { session, agent, model, entries, request, system, toolMaterialization }
+  })
+
+  type PreparedTurn = Effect.Success<ReturnType<typeof prepareTurn>>
+
+  /**
+   * Allocates the mutable state shared by provider consumption and settlement.
+   * Publication is serialized because provider events and local tool results may
+   * arrive concurrently but mutate one durable publisher state machine.
+   */
+  const makeRuntime = Effect.fnUntraced(function* (prepared: PreparedTurn) {
     const publisher = createLLMEventPublisher(events, {
-      sessionID: session.id,
-      agent: agent.id,
+      sessionID: prepared.session.id,
+      agent: prepared.agent.id,
       model: {
-        id: ModelV2.ID.make(model.id),
-        providerID: ProviderV2.ID.make(model.provider),
-        ...(session.model?.variant === undefined ? {} : { variant: session.model.variant }),
+        id: ModelV2.ID.make(prepared.model.id),
+        providerID: ProviderV2.ID.make(prepared.model.provider),
+        ...(prepared.session.model?.variant === undefined ? {} : { variant: prepared.session.model.variant }),
       },
     })
     const withPublication = Semaphore.makeUnsafe(1).withPermit
-    const publish = (event: LLMEvent, outputPaths: ReadonlyArray<string> = []) =>
-      withPublication(publisher.publish(event, outputPaths))
-    let overflowFailure: ProviderErrorEvent | undefined
-    if (!(yield* SessionContextEpoch.current(db, session.id, agent.id, system.revision)))
-      return yield* Effect.die(rebuildPreparedTurn())
-    const providerStream = llm.stream(request).pipe(
+    return {
+      publisher,
+      withPublication,
+      publish: (event: LLMEvent, outputPaths: ReadonlyArray<string> = []) =>
+        withPublication(publisher.publish(event, outputPaths)),
+      toolFibers: yield* FiberSet.make<void, ToolOutputStore.Error>(),
+      needsContinuation: false,
+      overflowFailure: undefined as ProviderErrorEvent | undefined,
+    }
+  })
+
+  type TurnRuntime = Effect.Success<ReturnType<typeof makeRuntime>>
+
+  /**
+   * Consumes exactly one provider stream.
+   *
+   * Every event is durably published before a local tool starts. Tool settlement
+   * is registered with the turn FiberSet before interruption can resume. A
+   * recoverable pre-output overflow is withheld until the settlement phase.
+   */
+  const consumeProvider = (prepared: PreparedTurn, runtime: TurnRuntime) =>
+    llm.stream(prepared.request).pipe(
       Stream.runForEach((event) =>
         Effect.gen(function* () {
-          if (overflowFailure || publisher.hasProviderError()) return
-          if (LLMEvent.is.providerError(event)) {
-            if (isContextOverflowFailure(event) && !publisher.hasAssistantStarted()) {
-              overflowFailure = event
-              return
-            }
+          if (runtime.overflowFailure || runtime.publisher.hasProviderError()) return
+          if (
+            LLMEvent.is.providerError(event) &&
+            isContextOverflowFailure(event) &&
+            !runtime.publisher.hasAssistantStarted()
+          ) {
+            runtime.overflowFailure = event
+            return
           }
-          yield* publish(event)
+          yield* runtime.publish(event)
           if (event.type !== "tool-call" || event.providerExecuted) return
-          needsContinuation = true
-          const assistantMessageID = yield* publisher.assistantMessageID(event.id)
+          runtime.needsContinuation = true
+          const assistantMessageID = yield* runtime.publisher.assistantMessageID(event.id)
           yield* Effect.uninterruptibleMask((restore) =>
             restore(
-              toolMaterialization.settle({
-                sessionID: session.id,
-                agent: agent.id,
+              prepared.toolMaterialization.settle({
+                sessionID: prepared.session.id,
+                agent: prepared.agent.id,
                 assistantMessageID,
                 call: event,
               }),
             ).pipe(
               Effect.flatMap((settlement) =>
-                publish(
+                runtime.publish(
                   LLMEvent.toolResult({
                     id: event.id,
                     name: event.name,
@@ -195,94 +230,119 @@ export const make = Effect.gen(function* () {
                 ),
               ),
             ),
-          ).pipe(FiberSet.run(toolFibers))
+          ).pipe(FiberSet.run(runtime.toolFibers))
         }),
       ),
-      Effect.ensuring(withPublication(publisher.flush())),
+      Effect.ensuring(runtime.withPublication(runtime.publisher.flush())),
     )
 
+  /**
+   * Runs one prepared provider attempt and settles every local tool it starts.
+   *
+   * The interruption mask keeps the handoff from stream completion to tool
+   * settlement atomic, while provider consumption and tool work remain interruptible.
+   */
+  const runAttempt = Effect.fn("SessionRunner.runTurn")(function* (
+    sessionID: SessionSchema.ID,
+    promotion: SessionInput.Delivery | undefined,
+    recoverOverflow?: typeof compaction.compactAfterOverflow,
+  ) {
+    const prepared = yield* prepareTurn(sessionID, promotion)
+    const runtime = yield* makeRuntime(prepared)
+    if (!(yield* SessionContextEpoch.current(db, prepared.session.id, prepared.agent.id, prepared.system.revision)))
+      return yield* Effect.fail(rebuildPreparedTurn())
+
     return yield* Effect.uninterruptibleMask((restore) =>
       Effect.gen(function* () {
-        const stream = yield* restore(providerStream).pipe(Effect.exit)
+        const stream = yield* restore(consumeProvider(prepared, runtime)).pipe(Effect.exit)
         const failure =
           stream._tag === "Failure" ? Option.getOrUndefined(Cause.findErrorOption(stream.cause)) : undefined
         if (
           recoverOverflow &&
-          !publisher.hasAssistantStarted() &&
-          isContextOverflowFailure(overflowFailure ?? failure) &&
-          (yield* restore(recoverOverflow({ sessionID: session.id, entries, model, request })))
+          !runtime.publisher.hasAssistantStarted() &&
+          isContextOverflowFailure(runtime.overflowFailure ?? failure) &&
+          (yield* restore(
+            recoverOverflow({
+              sessionID: prepared.session.id,
+              entries: prepared.entries,
+              model: prepared.model,
+              request: prepared.request,
+            }),
+          ))
         )
-          return yield* Effect.die(continueAfterOverflowCompaction)
-        if (overflowFailure) yield* publish(overflowFailure)
+          return yield* Effect.fail(continueAfterOverflowCompaction)
+        if (runtime.overflowFailure) yield* runtime.publish(runtime.overflowFailure)
         const llmFailure = failure instanceof LLMError ? failure : undefined
-        if (llmFailure && !publisher.hasProviderError()) {
-          yield* withPublication(publisher.failUnsettledTools("Provider did not return a tool result", true))
-          yield* withPublication(
+        if (llmFailure && !runtime.publisher.hasProviderError()) {
+          yield* runtime.withPublication(
+            runtime.publisher.failUnsettledTools("Provider did not return a tool result", true),
+          )
+          yield* runtime.withPublication(
             events.publish(SessionEvent.Step.Failed, {
-              sessionID: session.id,
+              sessionID: prepared.session.id,
               timestamp: yield* DateTime.now,
-              assistantMessageID: yield* publisher.startAssistant(),
+              assistantMessageID: yield* runtime.publisher.startAssistant(),
               error: { type: "unknown", message: llmFailure.reason.message },
             }),
           )
         }
-        if (stream._tag === "Failure" && Cause.hasInterrupts(stream.cause)) yield* FiberSet.clear(toolFibers)
-        const settled = yield* restore(awaitToolFibers(toolFibers)).pipe(Effect.exit)
+        if (stream._tag === "Failure" && Cause.hasInterrupts(stream.cause)) yield* FiberSet.clear(runtime.toolFibers)
+        const settled = yield* restore(awaitToolFibers(runtime.toolFibers)).pipe(Effect.exit)
         if (settled._tag === "Failure" && isQuestionRejected(settled.cause)) {
-          yield* FiberSet.clear(toolFibers)
-          yield* withPublication(publisher.failUnsettledTools("Tool execution interrupted"))
+          yield* FiberSet.clear(runtime.toolFibers)
+          yield* runtime.withPublication(runtime.publisher.failUnsettledTools("Tool execution interrupted"))
           return yield* Effect.interrupt
         }
         if (
           (stream._tag === "Failure" && Cause.hasInterrupts(stream.cause)) ||
           (settled._tag === "Failure" && Cause.hasInterrupts(settled.cause))
         ) {
-          yield* FiberSet.clear(toolFibers)
-          yield* withPublication(publisher.failUnsettledTools("Tool execution interrupted"))
+          yield* FiberSet.clear(runtime.toolFibers)
+          yield* runtime.withPublication(runtime.publisher.failUnsettledTools("Tool execution interrupted"))
         }
         if (settled._tag === "Failure" && !Cause.hasInterrupts(settled.cause)) {
           const failure = Cause.squash(settled.cause)
           const message = failure instanceof Error ? failure.message : String(failure)
-          yield* withPublication(publisher.failUnsettledTools(`Tool execution failed: ${message}`))
+          yield* runtime.withPublication(runtime.publisher.failUnsettledTools(`Tool execution failed: ${message}`))
         }
-        if (publisher.hasProviderError())
-          yield* withPublication(publisher.failUnsettledTools("Tool execution interrupted"))
-        if (stream._tag === "Success" && !publisher.hasProviderError())
-          yield* withPublication(publisher.failUnsettledTools("Provider did not return a tool result", true))
+        if (runtime.publisher.hasProviderError())
+          yield* runtime.withPublication(runtime.publisher.failUnsettledTools("Tool execution interrupted"))
+        if (stream._tag === "Success" && !runtime.publisher.hasProviderError())
+          yield* runtime.withPublication(
+            runtime.publisher.failUnsettledTools("Provider did not return a tool result", true),
+          )
         if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause)
         if (settled._tag === "Failure") return yield* Effect.failCause(settled.cause)
-        return !publisher.hasProviderError() && needsContinuation
+        return !runtime.publisher.hasProviderError() && runtime.needsContinuation
       }),
     )
   }, Effect.scoped)
 
-  const runAfterOverflowCompaction: Run = Effect.fnUntraced(function* (sessionID, promotion) {
-    return yield* runAttempt(sessionID, promotion).pipe(
-      Effect.catchDefect(
-        Effect.fnUntraced(function* (defect) {
-          if (!(defect instanceof TurnTransitionError)) return yield* Effect.die(defect)
-          if (defect.transition._tag === "ContinueAfterOverflowCompaction")
-            return yield* Effect.die("Post-compaction provider attempt cannot recover another overflow")
+  /** Rebuilds stale attempts while preserving the single overflow-recovery budget. */
+  const runState = Effect.fnUntraced(function* (
+    sessionID: SessionSchema.ID,
+    promotion: SessionInput.Delivery | undefined,
+    canRecoverOverflow: boolean,
+  ): Effect.fn.Return<boolean, RunError> {
+    return yield* runAttempt(
+      sessionID,
+      promotion,
+      canRecoverOverflow ? compaction.compactAfterOverflow : undefined,
+    ).pipe(
+      Effect.catchTags({
+        ContinueAfterOverflowCompaction: Effect.fnUntraced(function* () {
           yield* Effect.yieldNow
-          return yield* runAfterOverflowCompaction(sessionID, defect.transition.promotion)
+          return yield* runState(sessionID, undefined, false)
         }),
-      ),
+        RebuildPreparedTurn: Effect.fnUntraced(function* (transition) {
+          yield* Effect.yieldNow
+          return yield* runState(sessionID, transition.promotion, canRecoverOverflow)
+        }),
+      }),
     )
   })
 
-  const run: Run = Effect.fnUntraced(function* (sessionID, promotion) {
-    return yield* runAttempt(sessionID, promotion, compaction.compactAfterOverflow).pipe(
-      Effect.catchDefect(
-        Effect.fnUntraced(function* (defect) {
-          if (!(defect instanceof TurnTransitionError)) return yield* Effect.die(defect)
-          yield* Effect.yieldNow
-          if (defect.transition._tag === "ContinueAfterOverflowCompaction")
-            return yield* runAfterOverflowCompaction(sessionID, undefined)
-          return yield* run(sessionID, defect.transition.promotion)
-        }),
-      ),
-    )
-  })
+  const run: Run = (sessionID, promotion) => runState(sessionID, promotion, true)
 
   return run
 })