refactor(core): simplify session runner control flow (#45614)

packages/core/src/session/execution.ts

@@ -77,29 +77,29 @@ export const layer = Layer.effect(
     const releaseOnCommit = (sessionID: SessionSchema.ID) => ({
       commit: () => store.release(sessionID),
     })
-    function drain(
+    const drain = Effect.fnUntraced(function* (
       sessionID: SessionSchema.ID,
       force: boolean,
       continuation?: SessionRunner.Continuation,
       promotable: SessionInbox.Promotable = "input",
-    ): Effect.Effect<void, SessionRunner.RunError> {
-      return Effect.gen(function* () {
-        const session = yield* store.get(sessionID)
-        if (!session) return yield* Effect.die(new Error(`Session not found: ${sessionID}`))
-        const result = yield* SessionRunner.Service.use((runner) =>
-          runner.drain({ sessionID, force, continuation, promotable }),
-        ).pipe(
-          Effect.provide(locations.get(session.location)),
-          Effect.tapCause((cause) =>
-            Cause.hasInterruptsOnly(cause)
-              ? Effect.void
-              : Effect.logError("Failed to drain Session", cause).pipe(Effect.annotateLogs({ sessionID })),
-          ),
-        )
-        if (result._tag === "Complete") return
-        return yield* drain(sessionID, false, result.continuation, promotable)
+    ): Effect.fn.Return<void, SessionRunner.RunError> {
+      const session = yield* store.get(sessionID)
+      if (!session) return yield* Effect.die(new Error(`Session not found: ${sessionID}`))
+      const result = yield* SessionRunner.Service.use((runner) =>
+        runner.drain({ sessionID, force, continuation, promotable }),
+      ).pipe(
+        Effect.provide(locations.get(session.location)),
+        Effect.tapCause((cause) =>
+          Cause.hasInterruptsOnly(cause)
+            ? Effect.void
+            : Effect.logError("Failed to drain Session", cause).pipe(Effect.annotateLogs({ sessionID })),
+        ),
+      )
+      return yield* SessionRunner.DrainResult.$match(result, {
+        Complete: () => Effect.void,
+        Moved: (result) => drain(sessionID, false, result.continuation, promotable),
       })
-    }
+    })
     const coordinator = yield* SessionRunCoordinator.make<SessionSchema.ID, SessionRunner.RunError, InterruptReason>({
       started: (sessionID) =>
         reportLifecycle(

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

@@ -85,14 +85,9 @@ const layer = Layer.effect(
       const promotable = input.promotable ?? "input"
       if (!force && !continuing) {
         const pending = yield* SessionInbox.nextPromotable(db, sessionID, "input")
-        if (
-          !pending ||
-          (pending.delivery === "queue" &&
-            promotable === "steer" &&
-            pending.type !== "compaction" &&
-            pending.type !== "move")
-        )
-          return DrainResult.Complete()
+        if (!pending) return DrainResult.Complete()
+        const control = pending.type === "compaction" || pending.type === "move"
+        if (promotable === "steer" && pending.delivery === "queue" && !control) return DrainResult.Complete()
       }
       yield* plugins.flush
       yield* settleStaleToolCalls(sessionID)
@@ -263,29 +258,33 @@ const layer = Layer.effect(
               : Effect.succeed(false),
           ),
         })
-        if (outcome._tag === "Completed") return outcome.needsContinuation
-        if (outcome._tag === "Retry" || outcome._tag === "Continue") {
-          yield* retry({ cause: outcome.cause, error: outcome.error, assistantMessageID }).pipe(
-            Pull.catchDone(() =>
-              Effect.gen(function* () {
-                if (outcome._tag === "Retry")
-                  yield* bus.publish(SessionEvent.Step.Failed, { sessionID, assistantMessageID, error: outcome.error })
-                return yield* outcome.cause
-              }),
+        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)),
+              ),
+              Effect.asVoid,
             ),
-          )
-          if (outcome._tag === "Continue") {
+          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()
-          }
-          continue
-        }
-        if (outcome._tag === "Compacted") {
-          recoverOverflow = false
-          assistantMessageID = SessionMessage.ID.create()
-          continue
-        }
-        recoverContinuation = false
+          }),
+          Compacted: Effect.fnUntraced(function* () {
+            recoverOverflow = false
+            assistantMessageID = SessionMessage.ID.create()
+          }),
+          RecoverFull: Effect.fnUntraced(function* () {
+            recoverContinuation = false
+          }),
+        })
+        if (completed !== undefined) return completed
       }
     })
 

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

@@ -36,7 +36,7 @@ export type Outcome = Data.TaggedEnum<{
   RecoverFull: {}
   Compacted: {}
 }>
-const Outcome = Data.taggedEnum<Outcome>()
+export const Outcome = Data.taggedEnum<Outcome>()
 
 interface Input {
   readonly sessionID: SessionSchema.ID
@@ -127,11 +127,11 @@ export const make = Effect.gen(function* () {
       Effect.gen(function* () {
         const stream = yield* restore(providerStream).pipe(Effect.exit)
         const streamFailure = Option.getOrUndefined(Exit.findErrorOption(stream))
-        const streamInterrupted = stream._tag === "Failure" && Cause.hasInterrupts(stream.cause)
+        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 (joined._tag === "Failure") yield* interruptTools
+        if (Exit.isFailure(joined)) yield* interruptTools
         const tools = classifyToolExits(
           joined,
           toolRuns.map((run) => run.call),
@@ -147,7 +147,7 @@ export const make = Effect.gen(function* () {
         if (overflowFailure) yield* publisher.publish(overflowFailure)
         const recorded = publisher.record()
         const unknownFinish =
-          stream._tag === "Success" && recorded.finish?.finish === "unknown"
+          Exit.isSuccess(stream) && recorded.finish?.finish === "unknown"
             ? new AIError({
                 reason: new InvalidProviderOutputError({
                   message: "The provider response ended with an unknown finish reason.",
@@ -191,7 +191,7 @@ export const make = Effect.gen(function* () {
         if (interrupted) yield* publisher.failAssistant(STEP_INTERRUPTED)
 
         // All local fibers have joined; only provider-hosted results can still be missing.
-        if (llmError || (stream._tag === "Success" && !recorded.providerFailed)) {
+        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)
         }
@@ -234,10 +234,10 @@ export const make = Effect.gen(function* () {
         )
           return Outcome.Continue({ cause: llmFailure, error: llmError })
 
-        if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause)
+        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 && joined._tag === "Failure") return yield* Effect.failCause(joined.cause)
+        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.toolsDisabled && record.needsContinuation,
@@ -265,18 +265,17 @@ const classifyToolExits = (
   settled: Exit.Exit<Array<Exit.Exit<void, SessionModelRequest.ExecuteError>>>,
   calls: ReadonlyArray<ToolCall>,
 ) => {
-  const exits = settled._tag === "Success" ? settled.value : []
+  const exits = Exit.isSuccess(settled) ? settled.value : []
   const declines = exits.flatMap((exit, index) =>
-    exit._tag === "Failure"
+    Exit.isFailure(exit)
       ? exit.cause.reasons.flatMap((reason) =>
           Cause.isFailReason(reason) && isDecline(reason.error) ? [{ call: calls[index], reason: reason.error }] : [],
         )
       : [],
   )
-  const causes =
-    settled._tag === "Failure"
-      ? [settled.cause]
-      : exits.flatMap((exit) => (exit._tag === "Failure" ? [exit.cause] : []))
+  const causes = Exit.isFailure(settled)
+    ? [settled.cause]
+    : exits.flatMap((exit) => (Exit.isFailure(exit) ? [exit.cause] : []))
   const failure = causes
     .flatMap((cause) => {
       if (Cause.hasInterrupts(cause)) return []