From 4112698e726870f51c966bb1d1defe6bb1289501 Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Thu, 27 Aug 2026 14:42:12 -0400 Subject: [PATCH] refactor(core): simplify session runner control flow (#45614) --- packages/core/src/session/execution.ts | 36 ++++++++-------- packages/core/src/session/runner/llm.ts | 55 ++++++++++++------------ packages/core/src/session/runner/step.ts | 25 ++++++----- 3 files changed, 57 insertions(+), 59 deletions(-) diff --git a/packages/core/src/session/execution.ts b/packages/core/src/session/execution.ts index c49526973e0..dcd619b9bf5 100644 --- a/packages/core/src/session/execution.ts +++ b/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 { - 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 { + 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({ started: (sessionID) => reportLifecycle( diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index d797f5e941b..15550fc6eef 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/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 } }) diff --git a/packages/core/src/session/runner/step.ts b/packages/core/src/session/runner/step.ts index 60aaec0b70a..01495404ce9 100644 --- a/packages/core/src/session/runner/step.ts +++ b/packages/core/src/session/runner/step.ts @@ -36,7 +36,7 @@ export type Outcome = Data.TaggedEnum<{ RecoverFull: {} Compacted: {} }> -const Outcome = Data.taggedEnum() +export const Outcome = Data.taggedEnum() 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>>, calls: ReadonlyArray, ) => { - 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 []