From ad048465937530344058b4759c560fdb2fa782ec Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Sat, 6 Jun 2026 22:06:12 -0400 Subject: [PATCH] refactor(core): clarify provider turn phases --- packages/core/src/session/runner/run-turn.ts | 242 ++++++++++++------- 1 file changed, 151 insertions(+), 91 deletions(-) diff --git a/packages/core/src/session/runner/run-turn.ts b/packages/core/src/session/runner/run-turn.ts index 29b4956b938..6a0c6c35b8f 100644 --- a/packages/core/src/session/runner/run-turn.ts +++ b/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 -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) => 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() - 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> + + /** + * 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 = []) => - 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 = []) => + withPublication(publisher.publish(event, outputPaths)), + toolFibers: yield* FiberSet.make(), + needsContinuation: false, + overflowFailure: undefined as ProviderErrorEvent | undefined, + } + }) + + type TurnRuntime = Effect.Success> + + /** + * 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 { + 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 })