refactor(core): clean up callModel readability (#38706)

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

@@ -35,6 +35,33 @@ type CallOutcome = Data.TaggedEnum<{
 }>
 const CallOutcome = Data.taggedEnum<CallOutcome>()
 
+// Declining an interactive prompt halts the drain instead of becoming model-facing tool output.
+const isUserDeclined = (cause: Cause.Cause<unknown>) =>
+  cause.reasons.some(
+    (reason) =>
+      Cause.isDieReason(reason) &&
+      (reason.defect instanceof PermissionV2.DeclinedError || reason.defect instanceof QuestionTool.CancelledError),
+  )
+
+/**
+ * Classifies how the owned tool fibers ended. Interrupts and interactive declines abort
+ * the step; a defect from a tool implementation becomes a failed tool call the model can
+ * read; a typed infrastructure failure must fail the assistant and then the drain.
+ */
+const classifyToolExits = (settled: Exit.Exit<Array<Exit.Exit<void, ToolOutputStore.Error>>, never>) => {
+  const causes =
+    settled._tag === "Failure"
+      ? [settled.cause]
+      : settled.value.flatMap((exit) => (exit._tag === "Failure" ? [exit.cause] : []))
+  const failure = causes.find((cause) => !Cause.hasInterrupts(cause) && !isUserDeclined(cause))
+  return {
+    interrupted: causes.some(Cause.hasInterrupts),
+    declined: causes.some(isUserDeclined),
+    failure,
+    infraError: failure === undefined ? undefined : Option.getOrUndefined(Cause.findErrorOption(failure)),
+  }
+}
+
 const layer = Layer.effect(
   Service,
   Effect.gen(function* () {
@@ -52,44 +79,88 @@ const layer = Layer.effect(
     // re-fire a redundant LLM call; `SessionTitle` itself is idempotent based on durable history.
     const titleStarted = new Set<SessionSchema.ID>()
     const forkTitle = yield* FiberSet.makeRuntime<never, void, never>()
-    const getSession = Effect.fn("SessionRunner.getSession")(function* (sessionID: SessionSchema.ID) {
-      const session = yield* store.get(sessionID)
-      if (!session) return yield* Effect.die(new Error(`Session not found: ${sessionID}`))
-      return session
+    /**
+     * Drains eligible manual compaction and user input until the Session becomes idle.
+     * Execution lifecycle is published per busy period by SessionExecution, not here.
+     */
+    const drain = Effect.fn("SessionRunner.drain")(function* (input: {
+      readonly sessionID: SessionSchema.ID
+      readonly force: boolean
+    }) {
+      if (!input.force && !(yield* SessionPending.has(db, input.sessionID, "any"))) return
+      yield* settleStaleToolCalls(input.sessionID)
+      yield* runPendingCompaction(input.sessionID)
+      if (!input.force && !(yield* SessionPending.has(db, input.sessionID, "input"))) return
+      do {
+        yield* runSteps(input.sessionID)
+      } while (yield* SessionPending.has(db, input.sessionID, "input"))
     })
-    /** Fires title generation once per process after the first step makes a user message visible. */
-    const startTitleOnce = Effect.fnUntraced(function* (sessionID: SessionSchema.ID) {
-      if (titleStarted.has(sessionID)) return
-      titleStarted.add(sessionID)
-      forkTitle(title.generateForFirstPrompt(yield* getSession(sessionID)).pipe(Effect.ignore))
+
+    /**
+     * Runs logical steps until no tool result or newly admitted steer requires another
+     * model call. Queued inputs remain pending until the current model work reaches idle.
+     */
+    const runSteps = Effect.fn("SessionRunner.runSteps")(function* (sessionID: SessionSchema.ID) {
+      // Fresh work may promote queued input; later steps absorb steers only.
+      let promotable: SessionPending.Promotable = "input"
+      let step = 1
+      while (true) {
+        const result = yield* runStep(sessionID, promotable, step)
+        yield* startTitleOnce(sessionID)
+        yield* runPendingCompaction(sessionID)
+        if (!result.needsContinuation && !(yield* SessionPending.has(db, sessionID, "steer"))) return
+        promotable = "steer"
+        step = result.step + 1
+      }
     })
-    /** Closes stale tool calls left active by an earlier interrupted drain. */
-    const settleStaleToolCalls = Effect.fn("SessionRunner.settleStaleToolCalls")(function* (
+
+    /** Completes one logical model step, transparently retrying or rebuilding after compaction. */
+    const runStep = Effect.fnUntraced(function* (
       sessionID: SessionSchema.ID,
+      promotable: SessionPending.Promotable,
+      step: number,
     ) {
-      for (const message of yield* store.context(sessionID)) {
-        if (message.type !== "assistant") continue
-        for (const tool of message.content) {
-          if (tool.type !== "tool" || (tool.state.status !== "streaming" && tool.state.status !== "running")) continue
-          yield* events.publish(SessionEvent.Tool.Failed, {
-            sessionID,
-            assistantMessageID: message.id,
-            callID: tool.id,
-            error: { type: "aborted", message: `Tool execution interrupted: ${tool.name}` },
-            executed: tool.executed === true,
-          })
-        }
+      const retry = yield* Schedule.toStepWithSleep(SessionRunnerRetry.schedule(events, sessionID))
+      /**
+       * Consumes one retry allowance: sleeps the scheduled backoff and reports what the next
+       * attempt should reuse, or publishes Step.Failed and fails once attempts are exhausted.
+       * The step loop performs the retry itself on the next iteration.
+       */
+      const waitForRetry = (failure: SessionRunnerRetry.RetryableFailure) =>
+        retry(failure).pipe(
+          Effect.as(CallOutcome.Retry({ step: failure.step, assistantMessageID: failure.assistantMessageID })),
+          Pull.catchDone(() =>
+            events
+              .publish(SessionEvent.Step.Failed, {
+                sessionID,
+                assistantMessageID: failure.assistantMessageID,
+                error: failure.error,
+              })
+              .pipe(Effect.andThen(Effect.fail(failure.cause))),
+          ),
+        )
+      let currentPromotable: SessionPending.Promotable | undefined = promotable
+      let currentStep = step
+      let assistantMessageID: SessionMessage.ID | undefined
+      // Overflow recovery is one-shot: a call after recovery must not recover another overflow.
+      let recoverOverflow = true
+      while (true) {
+        const outcome = yield* callModel(
+          sessionID,
+          currentPromotable,
+          currentStep,
+          recoverOverflow,
+          assistantMessageID,
+        ).pipe(Effect.catchTag("SessionRunner.RetryableFailure", waitForRetry))
+        if (outcome._tag === "Completed") return { needsContinuation: outcome.needsContinuation, step: outcome.step }
+        if (outcome._tag === "Retry") assistantMessageID = outcome.assistantMessageID
+        if (outcome._tag === "Restart" && outcome.recoveredOverflow) recoverOverflow = false
+        // Neither a retry nor a compaction restart re-promotes input.
+        currentPromotable = undefined
+        currentStep = outcome.step
       }
     })
 
-    // Declining an interactive prompt halts the drain instead of becoming model-facing tool output.
-    const isUserDeclined = (cause: Cause.Cause<unknown>) =>
-      cause.reasons.some(
-        (reason) =>
-          Cause.isDieReason(reason) &&
-          (reason.defect instanceof PermissionV2.DeclinedError || reason.defect instanceof QuestionTool.CancelledError),
-      )
-
     /**
      * Prepares and runs at most one model call, executes its local tools, and durably
      * settles the step. Compaction may instead request that the logical step restart.
@@ -105,11 +176,9 @@ const layer = Layer.effect(
       // Establish what the model knows before admitting what the user said, so
       // a blocked first step leaves pending inputs untouched.
       yield* InstructionState.prepare(db, events, selected.instructions, selected.session.id)
-      let currentStep = step
-      if (promotable) {
-        const promoted = yield* SessionPending.promote(db, events, selected.session.id, promotable)
-        if (promoted > 0) currentStep = 1
-      }
+      const promoted = promotable ? yield* SessionPending.promote(db, events, selected.session.id, promotable) : 0
+      // Promoted input opens a fresh step allowance.
+      const currentStep = promoted > 0 ? 1 : step
       const loaded = yield* context.load(selected)
       const { session, agent } = loaded
       const resolved = loaded.model
@@ -119,7 +188,8 @@ const layer = Layer.effect(
       const compactionInput = { session, messages: loaded.messages, model, cost: resolved.cost }
       if (compaction.required(compactionInput) && !(yield* SessionPending.compaction(db, session.id))) {
         const compacted = yield* compaction.compact(compactionInput)
-        if (compacted.status === "completed") return CallOutcome.Restart({ step: currentStep, recoveredOverflow: false })
+        if (compacted.status === "completed")
+          return CallOutcome.Restart({ step: currentStep, recoveredOverflow: false })
         return yield* new StepFailedError({ error: compacted.error })
       }
       const prepared = yield* modelRequests.prepare({
@@ -236,8 +306,7 @@ const layer = Layer.effect(
             recoverOverflow &&
             !publisher.hasRetryEvidence() &&
             isContextOverflowFailure(overflowFailure ?? streamFailure) &&
-            (yield* restore(compaction.compact({ session, messages: loaded.messages, model, cost: resolved.cost })))
-              .status === "completed"
+            (yield* restore(compaction.compact(compactionInput))).status === "completed"
           )
             return CallOutcome.Restart({ step: currentStep, recoveredOverflow: true })
 
@@ -267,30 +336,17 @@ const layer = Layer.effect(
           const settled = yield* restore(
             Effect.forEach(ownedToolFibers, Fiber.await, { concurrency: "unbounded" }),
           ).pipe(Effect.exit)
-          const settledCauses =
-            settled._tag === "Failure"
-              ? [settled.cause]
-              : settled.value.flatMap((exit) => (exit._tag === "Failure" ? [exit.cause] : []))
-          const toolsInterrupted = settledCauses.some(Cause.hasInterrupts)
-          const userDeclined = settledCauses.some(isUserDeclined)
-
           if (settled._tag === "Failure") yield* FiberSet.clear(toolFibers)
-          if (userDeclined || streamInterrupted || toolsInterrupted) {
+          const tools = classifyToolExits(settled)
+
+          if (tools.declined || streamInterrupted || tools.interrupted) {
             yield* serialized(publisher.failUnsettledTools({ type: "aborted", message: "Tool execution interrupted" }))
             yield* serialized(publisher.failAssistant({ type: "aborted", message: "Step interrupted" }))
           }
-          // A settled tool fiber failure is one of two things. A defect from a tool
-          // implementation becomes a failed tool call the model can read, and the step still
-          // settles so the model may recover. A typed infrastructure failure (tool output
-          // could not be persisted) also fails the assistant and then fails the drain.
-          const settledFailure = settledCauses.find((cause) => !Cause.hasInterrupts(cause) && !isUserDeclined(cause))
-          const infraError =
-            settledFailure === undefined ? undefined : Option.getOrUndefined(Cause.findErrorOption(settledFailure))
-          if (settledFailure !== undefined) {
-            const failure = infraError ?? Cause.squash(settledFailure)
-            const error = toSessionError(failure)
+          if (tools.failure !== undefined) {
+            const error = toSessionError(tools.infraError ?? Cause.squash(tools.failure))
             yield* serialized(publisher.failUnsettledTools(error))
-            if (infraError !== undefined) yield* serialized(publisher.failAssistant(error))
+            if (tools.infraError !== undefined) yield* serialized(publisher.failAssistant(error))
           }
 
           // Fail unresolved calls before the terminal step event. Local calls have joined, so
@@ -338,63 +394,16 @@ const layer = Layer.effect(
           }
 
           if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause)
-          if (userDeclined) return yield* Effect.interrupt
-          if ((toolsInterrupted || infraError !== undefined) && settledFailure)
-            return yield* Effect.failCause(settledFailure)
-          if (toolsInterrupted && settled._tag === "Failure") return yield* Effect.failCause(settled.cause)
+          if (tools.declined) return yield* Effect.interrupt
+          if ((tools.interrupted || tools.infraError !== undefined) && tools.failure)
+            return yield* Effect.failCause(tools.failure)
+          if (tools.interrupted && settled._tag === "Failure") return yield* Effect.failCause(settled.cause)
           if (stepFailure) return yield* new StepFailedError({ error: stepFailure })
           return CallOutcome.Completed({ needsContinuation, step: currentStep })
         }),
       )
     }, Effect.scoped)
 
-    /** Completes one logical model step, transparently retrying or rebuilding after compaction. */
-    const runStep = Effect.fnUntraced(function* (
-      sessionID: SessionSchema.ID,
-      promotable: SessionPending.Promotable,
-      step: number,
-    ) {
-      const retry = yield* Schedule.toStepWithSleep(SessionRunnerRetry.schedule(events, sessionID))
-      /**
-       * Consumes one retry allowance: sleeps the scheduled backoff and reports what the next
-       * attempt should reuse, or publishes Step.Failed and fails once attempts are exhausted.
-       * The step loop performs the retry itself on the next iteration.
-       */
-      const waitForRetry = (failure: SessionRunnerRetry.RetryableFailure) =>
-        retry(failure).pipe(
-          Effect.as(CallOutcome.Retry({ step: failure.step, assistantMessageID: failure.assistantMessageID })),
-          Pull.catchDone(() =>
-            events
-              .publish(SessionEvent.Step.Failed, {
-                sessionID,
-                assistantMessageID: failure.assistantMessageID,
-                error: failure.error,
-              })
-              .pipe(Effect.andThen(Effect.fail(failure.cause))),
-          ),
-        )
-      let currentPromotable: SessionPending.Promotable | undefined = promotable
-      let currentStep = step
-      let assistantMessageID: SessionMessage.ID | undefined
-      // Overflow recovery is one-shot: a call after recovery must not recover another overflow.
-      let recoverOverflow = true
-      while (true) {
-        const outcome = yield* callModel(
-          sessionID,
-          currentPromotable,
-          currentStep,
-          recoverOverflow,
-          assistantMessageID,
-        ).pipe(Effect.catchTag("SessionRunner.RetryableFailure", waitForRetry))
-        if (outcome._tag === "Completed") return { needsContinuation: outcome.needsContinuation, step: outcome.step }
-        if (outcome._tag === "Retry") assistantMessageID = outcome.assistantMessageID
-        if (outcome._tag === "Restart" && outcome.recoveredOverflow) recoverOverflow = false
-        // Neither a retry nor a compaction restart re-promotes input.
-        currentPromotable = undefined
-        currentStep = outcome.step
-      }
-    })
-
     /** Executes a previously admitted manual compaction request, if one is pending. */
     const runPendingCompaction = Effect.fn("SessionRunner.runPendingCompaction")(function* (
       sessionID: SessionSchema.ID,
@@ -429,39 +438,36 @@ const layer = Layer.effect(
       )
     })
 
-    /**
-     * Runs logical steps until no tool result or newly admitted steer requires another
-     * model call. Queued inputs remain pending until the current model work reaches idle.
-     */
-    const runSteps = Effect.fn("SessionRunner.runSteps")(function* (sessionID: SessionSchema.ID) {
-      // Fresh work may promote queued input; later steps absorb steers only.
-      let promotable: SessionPending.Promotable = "input"
-      let step = 1
-      while (true) {
-        const result = yield* runStep(sessionID, promotable, step)
-        yield* startTitleOnce(sessionID)
-        yield* runPendingCompaction(sessionID)
-        if (!result.needsContinuation && !(yield* SessionPending.has(db, sessionID, "steer"))) return
-        promotable = "steer"
-        step = result.step + 1
+    /** Closes stale tool calls left active by an earlier interrupted drain. */
+    const settleStaleToolCalls = Effect.fn("SessionRunner.settleStaleToolCalls")(function* (
+      sessionID: SessionSchema.ID,
+    ) {
+      for (const message of yield* store.context(sessionID)) {
+        if (message.type !== "assistant") continue
+        for (const tool of message.content) {
+          if (tool.type !== "tool" || (tool.state.status !== "streaming" && tool.state.status !== "running")) continue
+          yield* events.publish(SessionEvent.Tool.Failed, {
+            sessionID,
+            assistantMessageID: message.id,
+            callID: tool.id,
+            error: { type: "aborted", message: `Tool execution interrupted: ${tool.name}` },
+            executed: tool.executed === true,
+          })
+        }
       }
     })
 
-    /**
-     * Drains eligible manual compaction and user input until the Session becomes idle.
-     * Execution lifecycle is published per busy period by SessionExecution, not here.
-     */
-    const drain = Effect.fn("SessionRunner.drain")(function* (input: {
-      readonly sessionID: SessionSchema.ID
-      readonly force: boolean
-    }) {
-      if (!input.force && !(yield* SessionPending.has(db, input.sessionID, "any"))) return
-      yield* settleStaleToolCalls(input.sessionID)
-      yield* runPendingCompaction(input.sessionID)
-      if (!input.force && !(yield* SessionPending.has(db, input.sessionID, "input"))) return
-      do {
-        yield* runSteps(input.sessionID)
-      } while (yield* SessionPending.has(db, input.sessionID, "input"))
+    /** Fires title generation once per process after the first step makes a user message visible. */
+    const startTitleOnce = Effect.fnUntraced(function* (sessionID: SessionSchema.ID) {
+      if (titleStarted.has(sessionID)) return
+      titleStarted.add(sessionID)
+      forkTitle(title.generateForFirstPrompt(yield* getSession(sessionID)).pipe(Effect.ignore))
+    })
+
+    const getSession = Effect.fn("SessionRunner.getSession")(function* (sessionID: SessionSchema.ID) {
+      const session = yield* store.get(sessionID)
+      if (!session) return yield* Effect.die(new Error(`Session not found: ${sessionID}`))
+      return session
     })
 
     return Service.of({ drain })