From 497be8941e718855c2af87b3660599e4f37207ca Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Sun, 21 Jun 2026 19:53:11 -0400 Subject: [PATCH] refactor(core): type runner retry control flow --- packages/core/src/session/context-epoch.ts | 20 +++--- packages/core/src/session/runner/llm.ts | 77 ++++++++++++---------- packages/core/test/session-runner.test.ts | 2 +- 3 files changed, 53 insertions(+), 46 deletions(-) diff --git a/packages/core/src/session/context-epoch.ts b/packages/core/src/session/context-epoch.ts index 1fb8df92e6..2c9a1ebeaf 100644 --- a/packages/core/src/session/context-epoch.ts +++ b/packages/core/src/session/context-epoch.ts @@ -18,7 +18,7 @@ type DatabaseService = Database.Interface["db"] class RevisionMismatch extends Error {} class LocationMismatch extends Error {} -export class AgentMismatch extends Error {} +export class AgentMismatch extends Schema.TaggedErrorClass()("SessionContextEpoch.AgentMismatch", {}) {} export class AgentReplacementBlocked extends Schema.TaggedErrorClass()( "SessionContextEpoch.AgentReplacementBlocked", { sessionID: SessionSchema.ID, previous: AgentV2.ID, current: AgentV2.ID }, @@ -45,7 +45,7 @@ export function initialize( sessionID: SessionSchema.ID, location: Location.Ref, agent: AgentV2.ID, -): Effect.Effect { +): Effect.Effect { return retryRevisionMismatch(() => initializeOnce(db, context, sessionID, location, agent)).pipe( Effect.withSpan("SessionContextEpoch.initialize"), ) @@ -58,7 +58,10 @@ export function prepare( sessionID: SessionSchema.ID, location: Location.Ref, agent: AgentV2.ID, -): Effect.Effect { +): Effect.Effect< + Prepared, + SystemContext.InitializationBlocked | ContextSnapshotDecodeError | AgentMismatch | AgentReplacementBlocked +> { return retryRevisionMismatch(() => prepareOnce(db, events, context, sessionID, location, agent)).pipe( Effect.withSpan("SessionContextEpoch.prepare"), ) @@ -153,7 +156,7 @@ const requireAgentSelection = Effect.fnUntraced(function* ( .where(eq(SessionTable.id, sessionID)) .get() .pipe(Effect.orDie) - if (!selected || (selected.agent !== null && selected.agent !== agent)) return yield* Effect.die(new AgentMismatch()) + if (!selected || (selected.agent !== null && selected.agent !== agent)) return yield* new AgentMismatch({}) }) export const requestReplacement = Effect.fn("SessionContextEpoch.requestReplacement")(function* ( @@ -212,7 +215,7 @@ const insert = Effect.fnUntraced(function* ( .get() .pipe(Effect.orDie) if (!placed) return yield* Effect.die(new LocationMismatch()) - if (placed.agent !== null && placed.agent !== agent) return yield* Effect.die(new AgentMismatch()) + if (placed.agent !== null && placed.agent !== agent) return yield* new AgentMismatch({}) const baselineSeq = yield* SessionInput.latestSeq(db, sessionID) yield* db .insert(SessionContextEpochTable) @@ -235,7 +238,7 @@ const insert = Effect.fnUntraced(function* ( }), { behavior: "immediate" }, ) - .pipe(Effect.orDie) + .pipe(Effect.catch((error) => (error instanceof AgentMismatch ? Effect.fail(error) : Effect.die(error)))) }) const replace = Effect.fnUntraced(function* ( @@ -274,7 +277,7 @@ const replace = Effect.fnUntraced(function* ( }), { behavior: "immediate" }, ) - .pipe(Effect.orDie) + .pipe(Effect.catch((error) => (error instanceof AgentMismatch ? Effect.fail(error) : Effect.die(error)))) }) const fence = Effect.fnUntraced(function* ( @@ -290,8 +293,7 @@ const fence = Effect.fnUntraced(function* ( .where(eq(SessionContextEpochTable.session_id, sessionID)) .get() .pipe(Effect.orDie) - if (!current || (current.selected !== null && current.selected !== agent)) - return yield* Effect.die(new AgentMismatch()) + if (!current || (current.selected !== null && current.selected !== agent)) return yield* new AgentMismatch({}) if (current.revision !== expectedRevision) return yield* Effect.die(new RevisionMismatch()) }) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index ec5c5246b2..8eee294594 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -36,6 +36,18 @@ import { createLLMEventPublisher } from "./publish-llm-event" import { toLLMMessages } from "./to-llm-message" import { MAX_STEPS_PROMPT } from "./max-steps" +type TurnState = Data.TaggedEnum<{ + OverflowAvailable: { readonly promotion: SessionInput.Delivery | undefined } + OverflowExhausted: {} +}> +const TurnState = Data.taggedEnum() + +type TurnResult = Data.TaggedEnum<{ + Retry: { readonly state: TurnState } + Complete: { readonly needsContinuation: boolean } +}> +const TurnResult = Data.taggedEnum() + /** * Runs one durable coding-agent Session until it settles. * @@ -140,21 +152,12 @@ export const layer = Layer.effect( const isQuestionRejected = (cause: Cause.Cause) => cause.reasons.some((reason) => Cause.isDieReason(reason) && reason.defect instanceof QuestionV2.RejectedError) - type TurnResult = Data.TaggedEnum<{ - // Request preparation observed a concurrent Session change and must restart from durable state. - Rebuild: { readonly nextPromotion: SessionInput.Delivery | undefined } - // Overflow compaction completed; rebuild once through the path without overflow recovery. - OverflowCompacted: {} - Complete: { readonly needsContinuation: boolean } - }> - const TurnResult = Data.taggedEnum() - - const optionOnAgentMismatch = (effect: Effect.Effect): Effect.Effect, E> => + const optionOnAgentMismatch = ( + effect: Effect.Effect, + ): Effect.Effect, E, R> => effect.pipe( - Effect.map(Option.some), - Effect.catchDefect((defect) => - defect instanceof SessionContextEpoch.AgentMismatch ? Effect.succeed(Option.none()) : Effect.die(defect), - ), + Effect.asSome, + Effect.catchTag("SessionContextEpoch.AgentMismatch", () => Effect.succeedNone), ) const sameModel = Schema.toEquivalence(Schema.UndefinedOr(ModelV2.Ref)) @@ -165,17 +168,16 @@ export const layer = Layer.effect( type RunTurnAttempt = ( sessionID: SessionSchema.ID, - promotion: SessionInput.Delivery | undefined, + state: TurnState, step: number, - recoverOverflow?: typeof compaction.compactAfterOverflow, ) => Effect.Effect const runTurnAttempt: RunTurnAttempt = Effect.fn("SessionRunner.runTurn")(function* ( sessionID: SessionSchema.ID, - promotion: SessionInput.Delivery | undefined, + state: TurnState, step: number, - recoverOverflow?: typeof compaction.compactAfterOverflow, ) { + const promotion = state._tag === "OverflowAvailable" ? state.promotion : undefined const session = yield* getSession(sessionID) if (session.location.directory !== location.directory || session.location.workspaceID !== location.workspaceID) return yield* Effect.interrupt @@ -183,7 +185,7 @@ export const layer = Layer.effect( const initialization = yield* optionOnAgentMismatch( SessionContextEpoch.initialize(db, loadSystemContext(agent), session.id, session.location, agent.id), ) - if (Option.isNone(initialization)) return TurnResult.Rebuild({ nextPromotion: promotion }) + if (Option.isNone(initialization)) return TurnResult.Retry({ state }) const initialized = initialization.value const toolFibers = yield* FiberSet.make() let needsContinuation = false @@ -201,11 +203,13 @@ export const layer = Layer.effect( SessionContextEpoch.prepare(db, events, loadSystemContext(agent), session.id, session.location, agent.id), ) : Option.some(initialized) - if (Option.isNone(preparation)) return TurnResult.Rebuild({ nextPromotion: undefined }) + const nextState = + state._tag === "OverflowAvailable" ? TurnState.OverflowAvailable({ promotion: undefined }) : state + if (Option.isNone(preparation)) return TurnResult.Retry({ state: nextState }) const system = preparation.value const current = yield* getSession(sessionID) if ((yield* agents.select(current.agent)).id !== agent.id || !sameModel(current.model, session.model)) - return TurnResult.Rebuild({ nextPromotion: undefined }) + return TurnResult.Retry({ state: nextState }) const model = yield* models.resolve(session) const entries = yield* SessionHistory.entriesForRunner(db, session.id, system.baselineSeq) const context = entries.map((entry) => entry.message) @@ -223,7 +227,7 @@ export const layer = Layer.effect( toolChoice: isLastStep ? "none" : undefined, }) if (yield* compaction.compactIfNeeded({ sessionID: session.id, entries, model, request })) - return TurnResult.Rebuild({ nextPromotion: undefined }) + return TurnResult.Retry({ state: nextState }) const publisher = createLLMEventPublisher(events, { sessionID: session.id, agent: agent.id, @@ -238,7 +242,7 @@ export const layer = Layer.effect( withPublication(publisher.publish(event, outputPaths)) let overflowFailure: ProviderErrorEvent | undefined if (!(yield* SessionContextEpoch.current(db, session.id, agent.id, system.revision))) - return TurnResult.Rebuild({ nextPromotion: undefined }) + return TurnResult.Retry({ state: nextState }) const providerStream = llm.stream(request).pipe( Stream.runForEach((event) => Effect.gen(function* () { @@ -290,12 +294,12 @@ export const layer = Layer.effect( const failure = stream._tag === "Failure" ? Option.getOrUndefined(Cause.findErrorOption(stream.cause)) : undefined if ( - recoverOverflow && + state._tag === "OverflowAvailable" && !publisher.hasAssistantStarted() && isContextOverflowFailure(overflowFailure ?? failure) && - (yield* restore(recoverOverflow({ sessionID: session.id, entries, model, request }))) + (yield* restore(compaction.compactAfterOverflow({ sessionID: session.id, entries, model, request }))) ) - return TurnResult.OverflowCompacted() + return TurnResult.Retry({ state: TurnState.OverflowExhausted() }) if (overflowFailure) yield* publish(overflowFailure) const llmFailure = failure instanceof LLMError ? failure : undefined if (llmFailure && !publisher.hasProviderError()) { @@ -340,18 +344,19 @@ export const layer = Layer.effect( ) => Effect.Effect const runTurn: RunTurn = Effect.fnUntraced(function* (sessionID, promotion, step) { - let nextPromotion = promotion - let recoverOverflow: typeof compaction.compactAfterOverflow | undefined = compaction.compactAfterOverflow + let state: TurnState = TurnState.OverflowAvailable({ promotion }) while (true) { - const result = yield* runTurnAttempt(sessionID, nextPromotion, step, recoverOverflow) - if (result._tag === "Complete") return result.needsContinuation - yield* Effect.yieldNow - if (result._tag === "OverflowCompacted") { - nextPromotion = undefined - recoverOverflow = undefined - continue + const result: TurnResult = yield* runTurnAttempt(sessionID, state, step) + switch (result._tag) { + case "Complete": + return result.needsContinuation + case "Retry": + state = result.state + yield* Effect.yieldNow + break + default: + return result satisfies never } - nextPromotion = result.nextPromotion } }) diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index eb5ccb277d..e93fc65122 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -1142,7 +1142,7 @@ describe("SessionRunnerLLM", () => { sessionID, location, AgentV2.defaultID, - ).pipe(Effect.catchDefect(Effect.succeed)), + ).pipe(Effect.flip), ).toBeInstanceOf(SessionContextEpoch.AgentMismatch) expect(