refactor(core): polish runner drain and coordinator readability (#35051)
This commit is contained in:
Kit Langton
GitHub
parent
cd0b274856
commit
e65477ab1d
3 changed files with 70 additions and 61 deletions
|
|
@ -6,11 +6,11 @@ import { Deferred, Effect, Exit, Fiber, FiberSet, Scope } from "effect"
|
|||
export interface Coordinator<Key, E> {
|
||||
/** Snapshots keys with an execution owned by this coordinator. */
|
||||
readonly active: Effect.Effect<ReadonlySet<Key>>
|
||||
/** 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<void, E>
|
||||
/** 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<void>
|
||||
/** 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<void>
|
||||
/** Resolves once no execution is active for the key. Returns immediately when already idle and never starts work. */
|
||||
readonly awaitIdle: (key: Key) => Effect.Effect<void>
|
||||
|
|
@ -30,6 +30,17 @@ type Execution<E> = {
|
|||
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 = <Key, E>(options: {
|
||||
readonly drain: (key: Key, force: boolean) => Effect.Effect<void, E>
|
||||
/**
|
||||
|
|
@ -84,6 +95,7 @@ export const make = <Key, E>(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))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 = <A, E, R>(effect: Effect.Effect<A, E, R>) => publication.withPermit(effect)
|
||||
const publish = (event: LLMEvent, outputPaths: ReadonlyArray<string> = []) =>
|
||||
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<ReturnType<typeof publisher.stepSettlement>>) =>
|
||||
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 })
|
||||
}),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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<void>()
|
||||
toolExecutionsStarted = yield* Deferred.make<void>()
|
||||
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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue