From e65477ab1d53ad360b53a43c4bafdf03e11365ab Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Thu, 2 Jul 2026 22:01:53 -0400 Subject: [PATCH] refactor(core): polish runner drain and coordinator readability (#35051) --- packages/core/src/session/run-coordinator.ts | 18 +++- packages/core/src/session/runner/llm.ts | 97 ++++++++++---------- packages/core/test/session-runner.test.ts | 16 ++-- 3 files changed, 70 insertions(+), 61 deletions(-) diff --git a/packages/core/src/session/run-coordinator.ts b/packages/core/src/session/run-coordinator.ts index e7834b766b..8524b1c1f8 100644 --- a/packages/core/src/session/run-coordinator.ts +++ b/packages/core/src/session/run-coordinator.ts @@ -6,11 +6,11 @@ import { Deferred, Effect, Exit, Fiber, FiberSet, Scope } from "effect" export interface Coordinator { /** Snapshots keys with an execution owned by this coordinator. */ readonly active: Effect.Effect> - /** Starts an execution while idle or joins the active execution. */ + /** Starts an execution while idle, or joins the active execution and returns its exit. */ readonly run: (key: Key) => Effect.Effect - /** Registers one coalesced follow-up after newly recorded work. */ + /** Rings the doorbell: an idle key starts an execution; an active one drains again before settling. */ readonly wake: (key: Key) => Effect.Effect - /** Stops the active execution and waits for its cleanup. */ + /** Stops the active execution, clears its doorbell, and waits for cleanup. No-op when idle. */ readonly interrupt: (key: Key) => Effect.Effect /** Resolves once no execution is active for the key. Returns immediately when already idle and never starts work. */ readonly awaitIdle: (key: Key) => Effect.Effect @@ -30,6 +30,17 @@ type Execution = { stopping: boolean } +/** + * ```text + * wake | run + * idle ──────────────▶ execution (one fiber) + * drain ⟲ doorbell rung mid-drain + * │ exit (settled hook runs) + * doorbell quiet ◀───────┴───────▶ doorbell rung + * idle, waiters get exit successor execution, + * waiters get this exit + * ``` + */ export const make = (options: { readonly drain: (key: Key, force: boolean) => Effect.Effect /** @@ -84,6 +95,7 @@ export const make = (options: { Effect.uninterruptibleMask((restore) => { const execution = executions.get(key) if (execution !== undefined) { + // A stopping execution refuses joiners: wait out its cleanup, then run fresh. if (execution.stopping) return restore(Deferred.await(execution.done).pipe(Effect.andThen(run(key)))) return restore(Deferred.await(execution.done)) } diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 5a5eff7ee8..85aa8a2c43 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -125,13 +125,10 @@ const layer = Layer.effect( return session }) - const getContext = Effect.fn("SessionRunner.getContext")(function* (sessionID: SessionSchema.ID) { - return yield* store.context(sessionID) - }) const failInterruptedTools = Effect.fn("SessionRunner.failInterruptedTools")(function* ( sessionID: SessionSchema.ID, ) { - for (const message of yield* getContext(sessionID)) { + 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 !== "pending" && tool.state.status !== "running")) continue @@ -235,8 +232,11 @@ const layer = Layer.effect( snapshot: startSnapshot, }) const publication = Semaphore.makeUnsafe(1) + // Durable publishes are serialized so tool fibers and turn settlement never interleave + // mid-event. + const serialized = (effect: Effect.Effect) => publication.withPermit(effect) const publish = (event: LLMEvent, outputPaths: ReadonlyArray = []) => - publication.withPermit(publisher.publish(event, outputPaths)) + serialized(publisher.publish(event, outputPaths)) let overflowFailure: ProviderErrorEvent | undefined const providerStream = llm.stream(request).pipe( Stream.runForEach((event) => @@ -251,9 +251,7 @@ const layer = Layer.effect( yield* publish(event) if (event.type !== "tool-call" || event.providerExecuted) return if (!toolMaterialization) { - yield* publication.withPermit( - publisher.failUnsettledTools("Tools are disabled after the maximum agent steps"), - ) + yield* serialized(publisher.failUnsettledTools("Tools are disabled after the maximum agent steps")) return } needsContinuation = true @@ -282,9 +280,34 @@ const layer = Layer.effect( ).pipe(FiberSet.run(toolFibers)) }), ), - Effect.ensuring(publication.withPermit(publisher.flush())), + Effect.ensuring(serialized(publisher.flush())), ) + // Captures the end snapshot, diffs it against the turn's start, and durably ends the + // assistant step. + const publishStepEnd = (settlement: NonNullable>) => + Effect.gen(function* () { + const endSnapshot = yield* snapshots.capture() + const files = + startSnapshot && endSnapshot + ? yield* snapshots + .files({ from: startSnapshot, to: endSnapshot }) + .pipe(Effect.catch(() => Effect.succeed(undefined))) + : undefined + yield* serialized( + events.publish(SessionEvent.Step.Ended, { + sessionID: session.id, + timestamp: yield* DateTime.now, + assistantMessageID: yield* publisher.startAssistant(), + finish: settlement.finish, + cost: 0, + tokens: settlement.tokens, + snapshot: endSnapshot, + files, + }), + ) + }) + return yield* Effect.uninterruptibleMask((restore) => Effect.gen(function* () { // Gather the evidence: how did the provider stream end? @@ -310,8 +333,8 @@ const layer = Layer.effect( if (overflowFailure) yield* publish(overflowFailure) const llmFailure = streamFailure instanceof LLMError ? streamFailure : undefined if (llmFailure && !publisher.hasProviderError()) { - yield* publication.withPermit(publisher.failUnsettledTools("Provider did not return a tool result", true)) - yield* publication.withPermit(publisher.failAssistant(llmFailure.reason.message)) + yield* serialized(publisher.failUnsettledTools("Provider did not return a tool result", true)) + yield* serialized(publisher.failAssistant(llmFailure.reason.message)) } // Provider error events only arrive from the stream, so the flag is final here. const providerFailed = publisher.hasProviderError() @@ -324,8 +347,8 @@ const layer = Layer.effect( if (questionDismissed || streamInterrupted || toolsInterrupted) { yield* FiberSet.clear(toolFibers) - yield* publication.withPermit(publisher.failUnsettledTools("Tool execution interrupted")) - yield* publication.withPermit(publisher.failAssistant("Provider turn interrupted")) + yield* serialized(publisher.failUnsettledTools("Tool execution interrupted")) + yield* serialized(publisher.failAssistant("Provider turn interrupted")) // Match V1: dismissing a question halts the loop like an interruption. if (questionDismissed) return yield* Effect.interrupt } @@ -339,44 +362,20 @@ const layer = Layer.effect( if (settledFailure !== undefined) { const failure = infraError ?? Cause.squash(settledFailure) const message = failure instanceof Error ? failure.message : String(failure) - yield* publication.withPermit(publisher.failUnsettledTools(`Tool execution failed: ${message}`)) + yield* serialized(publisher.failUnsettledTools(`Tool execution failed: ${message}`)) if (infraError !== undefined) - yield* publication.withPermit(publisher.failAssistant(`Tool execution failed: ${message}`)) + yield* serialized(publisher.failAssistant(`Tool execution failed: ${message}`)) } const stepSettlement = publisher.stepSettlement() - if ( - stepSettlement && - !streamInterrupted && - !toolsInterrupted && - infraError === undefined && - !providerFailed - ) { - const endSnapshot = yield* snapshots.capture() - const files = - startSnapshot && endSnapshot - ? yield* snapshots - .files({ from: startSnapshot, to: endSnapshot }) - .pipe(Effect.catch(() => Effect.succeed(undefined))) - : undefined - yield* publication.withPermit( - events.publish(SessionEvent.Step.Ended, { - sessionID: session.id, - timestamp: yield* DateTime.now, - assistantMessageID: yield* publisher.startAssistant(), - finish: stepSettlement.finish, - cost: 0, - tokens: stepSettlement.tokens, - snapshot: endSnapshot, - files, - }), - ) - } + const stepEndedCleanly = + !streamInterrupted && !toolsInterrupted && infraError === undefined && !providerFailed + if (stepSettlement && stepEndedCleanly) yield* publishStepEnd(stepSettlement) // A provider error orphans recorded local calls; a clean stream can still leave // hosted calls without results. - if (providerFailed) yield* publication.withPermit(publisher.failUnsettledTools("Tool execution interrupted")) + if (providerFailed) yield* serialized(publisher.failUnsettledTools("Tool execution interrupted")) if (stream._tag === "Success" && !providerFailed) - yield* publication.withPermit(publisher.failUnsettledTools("Provider did not return a tool result", true)) + yield* serialized(publisher.failUnsettledTools("Provider did not return a tool result", true)) if (stream._tag === "Failure") return yield* Effect.failCause(stream.cause) if (settled._tag === "Failure" && (toolsInterrupted || infraError !== undefined)) @@ -411,7 +410,9 @@ const layer = Layer.effect( } }) - const drain = Effect.fnUntraced(function* (input: { + // ExecutionSettled is published per execution (busy period) by SessionExecution, not per + // drain here. + const run = Effect.fn("SessionRunner.run")(function* (input: { readonly sessionID: SessionSchema.ID readonly force: boolean }) { @@ -442,11 +443,7 @@ const layer = Layer.effect( } }) - return Service.of({ - // ExecutionSettled is published per execution (busy period) by SessionExecution, - // not per drain here. - run: Effect.fn("SessionRunner.run")(drain), - }) + return Service.of({ run }) }), ) diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 4232112094..3014f231f8 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -2863,14 +2863,12 @@ describe("SessionRunnerLLM", () => { input: Schema.Struct({}), output: Schema.Struct({}), execute: (_, context) => - forms - .ask({ sessionID: context.sessionID, mode: "form", fields: [] }) - .pipe( - Effect.orDie, - Effect.flatMap((state) => - state.status === "answered" ? Effect.succeed({}) : Effect.die(new QuestionTool.RejectedError()), - ), + forms.ask({ sessionID: context.sessionID, mode: "form", fields: [] }).pipe( + Effect.orDie, + Effect.flatMap((state) => + state.status === "answered" ? Effect.succeed({}) : Effect.die(new QuestionTool.RejectedError()), ), + ), }), }) yield* session.prompt({ sessionID, prompt: Prompt.make({ text: "Ask then stop" }), resume: false }) @@ -3033,6 +3031,8 @@ describe("SessionRunnerLLM", () => { yield* session.prompt({ sessionID, prompt: Prompt.make({ text: "Interrupt tool settlement" }), resume: false }) executions.length = 0 toolExecutionGate = yield* Deferred.make() + toolExecutionsStarted = yield* Deferred.make() + toolExecutionsReady = 1 response = [ LLMEvent.stepStart({ index: 0 }), LLMEvent.toolCall({ id: "call-await-interrupt", name: "echo", input: { text: "blocked" } }), @@ -3042,7 +3042,7 @@ describe("SessionRunnerLLM", () => { const runner = yield* SessionRunner.Service const run = yield* runner.run({ sessionID, force: true }).pipe(Effect.forkChild) - while (executions.length === 0) yield* Effect.yieldNow + yield* Deferred.await(toolExecutionsStarted) yield* Fiber.interrupt(run) toolExecutionGate = undefined