refactor(core): deepen provider turn runner

This commit is contained in:
Kit Langton 2026-06-06 22:24:53 -04:00
commit 652ec048e7
2 changed files with 76 additions and 92 deletions

View file

@ -108,7 +108,7 @@ export const layer = Layer.effect(
while (openActivity) { while (openActivity) {
let needsContinuation = true let needsContinuation = true
for (let step = 0; step < MAX_STEPS; step++) { for (let step = 0; step < MAX_STEPS; step++) {
needsContinuation = yield* runTurn(input.sessionID, promotion) needsContinuation = yield* runTurn({ sessionID: input.sessionID, delivery: promotion })
promotion = "steer" promotion = "steer"
if (!needsContinuation) needsContinuation = yield* SessionInput.hasPending(db, input.sessionID, "steer") if (!needsContinuation) needsContinuation = yield* SessionInput.hasPending(db, input.sessionID, "steer")
if (!needsContinuation) break if (!needsContinuation) break

View file

@ -44,10 +44,12 @@ import { SessionRunnerModel } from "./model"
import { createLLMEventPublisher } from "./publish-llm-event" import { createLLMEventPublisher } from "./publish-llm-event"
import { toLLMMessages } from "./to-llm-message" import { toLLMMessages } from "./to-llm-message"
export type Run = ( export interface Input {
sessionID: SessionSchema.ID, readonly sessionID: SessionSchema.ID
promotion: SessionInput.Delivery | undefined, readonly delivery?: SessionInput.Delivery
) => Effect.Effect<boolean, RunError> }
export type Run = (input: Input) => Effect.Effect<boolean, RunError>
const AttemptResult = Schema.TaggedUnion({ const AttemptResult = Schema.TaggedUnion({
Complete: { needsContinuation: Schema.Boolean }, Complete: { needsContinuation: Schema.Boolean },
@ -96,11 +98,11 @@ export const make = Effect.gen(function* () {
* Once input is promoted, retries load it from history instead of promoting * Once input is promoted, retries load it from history instead of promoting
* again. This matters for queued input because promotion opens the next item. * again. This matters for queued input because promotion opens the next item.
*/ */
const prepareTurn = Effect.fn("SessionRunner.prepareTurn")(function* ( const buildRequest = Effect.fn("SessionRunner.buildRequest")(function* (
sessionID: SessionSchema.ID, sessionID: SessionSchema.ID,
promotion: SessionInput.Delivery | undefined, delivery: SessionInput.Delivery | undefined,
) { ) {
let pendingPromotion = promotion let pendingDelivery = delivery
while (true) { while (true) {
const session = yield* getSession(sessionID) const session = yield* getSession(sessionID)
if (session.location.directory !== location.directory || session.location.workspaceID !== location.workspaceID) if (session.location.directory !== location.directory || session.location.workspaceID !== location.workspaceID)
@ -110,14 +112,14 @@ export const make = Effect.gen(function* () {
SessionContextEpoch.initialize(db, loadSystemContext(agent), session.id, session.location, agent.id), SessionContextEpoch.initialize(db, loadSystemContext(agent), session.id, session.location, agent.id),
) )
if (initialized === stale) continue if (initialized === stale) continue
if (pendingPromotion) { if (pendingDelivery) {
const cutoff = yield* SessionInput.latestSeq(db, session.id) const cutoff = yield* SessionInput.latestSeq(db, session.id)
if (pendingPromotion === "steer") yield* SessionInput.promoteSteers(db, events, session.id, cutoff) if (pendingDelivery === "steer") yield* SessionInput.promoteSteers(db, events, session.id, cutoff)
if (pendingPromotion === "queue") { if (pendingDelivery === "queue") {
yield* SessionInput.promoteNextQueued(db, events, session.id) yield* SessionInput.promoteNextQueued(db, events, session.id)
yield* SessionInput.promoteSteers(db, events, session.id, cutoff) yield* SessionInput.promoteSteers(db, events, session.id, cutoff)
} }
pendingPromotion = undefined pendingDelivery = undefined
} }
const prepared = const prepared =
initialized ?? initialized ??
@ -153,13 +155,20 @@ export const make = Effect.gen(function* () {
} }
}) })
type RequestSnapshot = Effect.Success<ReturnType<typeof prepareTurn>> type RequestSnapshot = Effect.Success<ReturnType<typeof buildRequest>>
/** /**
* Provider events and tool results can arrive concurrently. They share one * Reads one model response and finishes every tool it starts.
* permit so their durable Session events are written in order. *
* Provider events and tool results share one permit so durable events stay in
* order. Tool calls are recorded before their side effects begin. A pre-output
* overflow is held back so successful compaction does not leave a terminal
* error in Session history.
*/ */
const startTurn = Effect.fnUntraced(function* (prepared: RequestSnapshot) { const streamAndSettle = Effect.fn("SessionRunner.streamAndSettle")(function* (
prepared: RequestSnapshot,
recoverOverflow?: typeof compaction.compactAfterOverflow,
) {
const publisher = createLLMEventPublisher(events, { const publisher = createLLMEventPublisher(events, {
sessionID: prepared.session.id, sessionID: prepared.session.id,
agent: prepared.agent.id, agent: prepared.agent.id,
@ -170,41 +179,23 @@ export const make = Effect.gen(function* () {
}, },
}) })
const withPublication = Semaphore.makeUnsafe(1).withPermit const withPublication = Semaphore.makeUnsafe(1).withPermit
return { const publish = (event: LLMEvent, outputPaths: ReadonlyArray<string> = []) =>
publisher, withPublication(publisher.publish(event, outputPaths))
withPublication, const toolFibers = yield* FiberSet.make<void, ToolOutputStore.Error>()
publish: (event: LLMEvent, outputPaths: ReadonlyArray<string> = []) => let needsContinuation = false
withPublication(publisher.publish(event, outputPaths)), let overflowFailure: ProviderErrorEvent | undefined
toolFibers: yield* FiberSet.make<void, ToolOutputStore.Error>(), const providerStream = llm.stream(prepared.request).pipe(
needsContinuation: false,
overflowFailure: undefined as ProviderErrorEvent | undefined,
}
})
type ActiveTurn = Effect.Success<ReturnType<typeof startTurn>>
/**
* Reads one model response. A tool call is recorded before its side effect
* starts. An overflow error is held back briefly so successful compaction does
* not leave a terminal error in Session history.
*/
const consumeProvider = (prepared: RequestSnapshot, runtime: ActiveTurn) =>
llm.stream(prepared.request).pipe(
Stream.runForEach((event) => Stream.runForEach((event) =>
Effect.gen(function* () { Effect.gen(function* () {
if (runtime.overflowFailure || runtime.publisher.hasProviderError()) return if (overflowFailure || publisher.hasProviderError()) return
if ( if (LLMEvent.is.providerError(event) && isContextOverflowFailure(event) && !publisher.hasAssistantStarted()) {
LLMEvent.is.providerError(event) && overflowFailure = event
isContextOverflowFailure(event) &&
!runtime.publisher.hasAssistantStarted()
) {
runtime.overflowFailure = event
return return
} }
yield* runtime.publish(event) yield* publish(event)
if (event.type !== "tool-call" || event.providerExecuted) return if (event.type !== "tool-call" || event.providerExecuted) return
runtime.needsContinuation = true needsContinuation = true
const assistantMessageID = yield* runtime.publisher.assistantMessageID(event.id) const assistantMessageID = yield* publisher.assistantMessageID(event.id)
yield* Effect.uninterruptibleMask((restore) => yield* Effect.uninterruptibleMask((restore) =>
restore( restore(
prepared.toolMaterialization.settle({ prepared.toolMaterialization.settle({
@ -215,7 +206,7 @@ export const make = Effect.gen(function* () {
}), }),
).pipe( ).pipe(
Effect.flatMap((settlement) => Effect.flatMap((settlement) =>
runtime.publish( publish(
LLMEvent.toolResult({ LLMEvent.toolResult({
id: event.id, id: event.id,
name: event.name, name: event.name,
@ -226,32 +217,23 @@ export const make = Effect.gen(function* () {
), ),
), ),
), ),
).pipe(FiberSet.run(runtime.toolFibers)) ).pipe(FiberSet.run(toolFibers))
}), }),
), ),
Effect.ensuring(runtime.withPublication(runtime.publisher.flush())), Effect.ensuring(withPublication(publisher.flush())),
) )
/** // Keep cleanup protected after the response ends so no started tool is
* The model response and tools remain interruptible. The short handoff after // forgotten, while the response stream and tool work remain interruptible.
* the response ends is protected so no started tool is forgotten before cleanup.
*/
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* startTurn(prepared)
return yield* Effect.uninterruptibleMask((restore) => return yield* Effect.uninterruptibleMask((restore) =>
Effect.gen(function* () { Effect.gen(function* () {
const stream = yield* restore(consumeProvider(prepared, runtime)).pipe(Effect.exit) const stream = yield* restore(providerStream).pipe(Effect.exit)
const failure = const failure =
stream._tag === "Failure" ? Option.getOrUndefined(Cause.findErrorOption(stream.cause)) : undefined stream._tag === "Failure" ? Option.getOrUndefined(Cause.findErrorOption(stream.cause)) : undefined
if ( if (
recoverOverflow && recoverOverflow &&
!runtime.publisher.hasAssistantStarted() && !publisher.hasAssistantStarted() &&
isContextOverflowFailure(runtime.overflowFailure ?? failure) && isContextOverflowFailure(overflowFailure ?? failure) &&
(yield* restore( (yield* restore(
recoverOverflow({ recoverOverflow({
sessionID: prepared.session.id, sessionID: prepared.session.id,
@ -262,63 +244,65 @@ export const make = Effect.gen(function* () {
)) ))
) )
return AttemptResult.cases.CompactedOverflow.make({}) return AttemptResult.cases.CompactedOverflow.make({})
if (runtime.overflowFailure) yield* runtime.publish(runtime.overflowFailure) if (overflowFailure) yield* publish(overflowFailure)
const llmFailure = failure instanceof LLMError ? failure : undefined const llmFailure = failure instanceof LLMError ? failure : undefined
if (llmFailure && !runtime.publisher.hasProviderError()) { if (llmFailure && !publisher.hasProviderError()) {
yield* runtime.withPublication( yield* withPublication(publisher.failUnsettledTools("Provider did not return a tool result", true))
runtime.publisher.failUnsettledTools("Provider did not return a tool result", true), yield* withPublication(
)
yield* runtime.withPublication(
events.publish(SessionEvent.Step.Failed, { events.publish(SessionEvent.Step.Failed, {
sessionID: prepared.session.id, sessionID: prepared.session.id,
timestamp: yield* DateTime.now, timestamp: yield* DateTime.now,
assistantMessageID: yield* runtime.publisher.startAssistant(), assistantMessageID: yield* publisher.startAssistant(),
error: { type: "unknown", message: llmFailure.reason.message }, error: { type: "unknown", message: llmFailure.reason.message },
}), }),
) )
} }
if (stream._tag === "Failure" && Cause.hasInterrupts(stream.cause)) yield* FiberSet.clear(runtime.toolFibers) if (stream._tag === "Failure" && Cause.hasInterrupts(stream.cause)) yield* FiberSet.clear(toolFibers)
const settled = yield* restore(awaitToolFibers(runtime.toolFibers)).pipe(Effect.exit) const settled = yield* restore(awaitToolFibers(toolFibers)).pipe(Effect.exit)
if (settled._tag === "Failure" && isQuestionRejected(settled.cause)) { if (settled._tag === "Failure" && isQuestionRejected(settled.cause)) {
yield* FiberSet.clear(runtime.toolFibers) yield* FiberSet.clear(toolFibers)
yield* runtime.withPublication(runtime.publisher.failUnsettledTools("Tool execution interrupted")) yield* withPublication(publisher.failUnsettledTools("Tool execution interrupted"))
return yield* Effect.interrupt return yield* Effect.interrupt
} }
if ( if (
(stream._tag === "Failure" && Cause.hasInterrupts(stream.cause)) || (stream._tag === "Failure" && Cause.hasInterrupts(stream.cause)) ||
(settled._tag === "Failure" && Cause.hasInterrupts(settled.cause)) (settled._tag === "Failure" && Cause.hasInterrupts(settled.cause))
) { ) {
yield* FiberSet.clear(runtime.toolFibers) yield* FiberSet.clear(toolFibers)
yield* runtime.withPublication(runtime.publisher.failUnsettledTools("Tool execution interrupted")) yield* withPublication(publisher.failUnsettledTools("Tool execution interrupted"))
} }
if (settled._tag === "Failure" && !Cause.hasInterrupts(settled.cause)) { if (settled._tag === "Failure" && !Cause.hasInterrupts(settled.cause)) {
const failure = Cause.squash(settled.cause) const failure = Cause.squash(settled.cause)
const message = failure instanceof Error ? failure.message : String(failure) const message = failure instanceof Error ? failure.message : String(failure)
yield* runtime.withPublication(runtime.publisher.failUnsettledTools(`Tool execution failed: ${message}`)) yield* withPublication(publisher.failUnsettledTools(`Tool execution failed: ${message}`))
} }
if (runtime.publisher.hasProviderError()) if (publisher.hasProviderError())
yield* runtime.withPublication(runtime.publisher.failUnsettledTools("Tool execution interrupted")) yield* withPublication(publisher.failUnsettledTools("Tool execution interrupted"))
if (stream._tag === "Success" && !runtime.publisher.hasProviderError()) if (stream._tag === "Success" && !publisher.hasProviderError())
yield* runtime.withPublication( yield* withPublication(publisher.failUnsettledTools("Provider did not return a tool result", true))
runtime.publisher.failUnsettledTools("Provider did not return a tool result", true),
)
if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause) if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause)
if (settled._tag === "Failure") return yield* Effect.failCause(settled.cause) if (settled._tag === "Failure") return yield* Effect.failCause(settled.cause)
return AttemptResult.cases.Complete.make({ return AttemptResult.cases.Complete.make({
needsContinuation: !runtime.publisher.hasProviderError() && runtime.needsContinuation, needsContinuation: !publisher.hasProviderError() && needsContinuation,
}) })
}), }),
) )
}, Effect.scoped) }, Effect.scoped)
const run: Run = Effect.fnUntraced(function* (sessionID, promotion) { const run: Run = Effect.fn("SessionRunner.runTurn")(function* (input) {
const first = yield* runAttempt(sessionID, promotion, compaction.compactAfterOverflow) let pendingDelivery = input.delivery
if (first._tag === "Complete") return first.needsContinuation let canRecoverOverflow = true
const second = yield* runAttempt(sessionID, undefined) while (true) {
return AttemptResult.match(second, { const request = yield* buildRequest(input.sessionID, pendingDelivery)
Complete: (result) => result.needsContinuation, pendingDelivery = undefined
CompactedOverflow: () => false, const result = yield* streamAndSettle(request, canRecoverOverflow ? compaction.compactAfterOverflow : undefined)
}) const next = AttemptResult.match(result, {
Complete: (completed) => completed.needsContinuation,
CompactedOverflow: () => undefined,
})
if (next !== undefined) return next
canRecoverOverflow = false
}
}) })
return run return run