diff --git a/agent/extensions/delegate/manager-state.ts b/agent/extensions/delegate/manager-state.ts new file mode 100644 index 0000000..6154acf --- /dev/null +++ b/agent/extensions/delegate/manager-state.ts @@ -0,0 +1,392 @@ +import { Deferred, Effect, Fiber } from "effect"; +import type { DelegateStatus } from "./contract.ts"; +import { cancelTimer, type scheduleTimer } from "./host-timers.ts"; +import type { ChildSession } from "./runtime.ts"; + +type ExecutionTimer = ReturnType; + +type SubscribingChild = { + readonly child: ChildSession; +}; + +type OwnedChild = SubscribingChild & { + readonly unsubscribe: () => void; +}; + +type StoppingOwnership = + | { readonly kind: "subscribing"; readonly child: ChildSession } + | ({ readonly kind: "running" } & OwnedChild); + +type StopReason = + | { readonly kind: "cancel" } + | { readonly kind: "execution-ceiling"; readonly error: string }; + +type SettledOutcome = + | { readonly kind: "done" } + | { + readonly kind: "error"; + readonly error: string; + readonly checkpoint: string; + } + | { + readonly kind: "cancelled"; + readonly error: string; + readonly checkpoint: string; + }; + +type RunLifecycle = + | { readonly kind: "creating"; readonly timer: ExecutionTimer } + | { + readonly kind: "subscribing"; + readonly ownership: SubscribingChild; + readonly timer: ExecutionTimer; + } + | { + readonly kind: "running"; + readonly ownership: OwnedChild; + readonly timer: ExecutionTimer; + } + | { + readonly kind: "stopping"; + readonly task: Fiber.Fiber; + readonly reason: StopReason; + readonly ownership?: StoppingOwnership; + } + | { + readonly kind: "settled"; + readonly settledAt: number; + readonly settlementOrder: number; + readonly outcome: SettledOutcome; + }; + +type Delivery = + | { readonly kind: "foreground" } + | { readonly kind: "pending"; readonly waiters: number } + | { readonly kind: "consumed" }; + +export type RunStateView = + | { readonly status: "running" } + | { + readonly status: "done"; + readonly settledAt: number; + readonly settlementOrder: number; + } + | { + readonly status: "error" | "cancelled"; + readonly settledAt: number; + readonly settlementOrder: number; + readonly error: string; + readonly checkpoint: string; + }; + +export type SettlementTransition = + | { readonly kind: "unchanged" } + | { readonly kind: "settled"; readonly child?: ChildSession }; + +function assertNever(value: never): never { + throw new Error(`Unhandled delegate state: ${String(value)}`); +} + +export class RunState { + private lifecycle: RunLifecycle; + private delivery: Delivery; + + private constructor(lifecycle: RunLifecycle, delivery: Delivery) { + this.lifecycle = lifecycle; + this.delivery = delivery; + } + + static creating(timer: ExecutionTimer, background: boolean): RunState { + return new RunState( + { kind: "creating", timer }, + background ? { kind: "pending", waiters: 0 } : { kind: "foreground" }, + ); + } + + view(): RunStateView { + const lifecycle = this.lifecycle; + switch (lifecycle.kind) { + case "creating": + case "subscribing": + case "running": + case "stopping": + return { status: "running" }; + case "settled": + switch (lifecycle.outcome.kind) { + case "done": + return { + status: "done", + settledAt: lifecycle.settledAt, + settlementOrder: lifecycle.settlementOrder, + }; + case "error": + case "cancelled": + return { + status: lifecycle.outcome.kind, + settledAt: lifecycle.settledAt, + settlementOrder: lifecycle.settlementOrder, + error: lifecycle.outcome.error, + checkpoint: lifecycle.outcome.checkpoint, + }; + default: + return assertNever(lifecycle.outcome); + } + default: + return assertNever(lifecycle); + } + } + + status(): DelegateStatus { + return this.view().status; + } + + isActive(): boolean { + return this.lifecycle.kind !== "settled"; + } + + settlementOrder(): number { + const view = this.view(); + return view.status === "running" ? 0 : view.settlementOrder; + } + + startSubscribing(child: ChildSession): boolean { + const lifecycle = this.lifecycle; + if (lifecycle.kind !== "creating") return false; + this.lifecycle = { + kind: "subscribing", + ownership: { child }, + timer: lifecycle.timer, + }; + return true; + } + + startRunning(child: ChildSession, unsubscribe: () => void): boolean { + const lifecycle = this.lifecycle; + if (lifecycle.kind !== "subscribing" || lifecycle.ownership.child !== child) + return false; + this.lifecycle = { + kind: "running", + ownership: { child, unsubscribe }, + timer: lifecycle.timer, + }; + return true; + } + + runningChild(): ChildSession | undefined { + return this.lifecycle.kind === "running" + ? this.lifecycle.ownership.child + : undefined; + } + + ownsRunningChild(child: ChildSession): boolean { + return ( + this.lifecycle.kind === "running" && + this.lifecycle.ownership.child === child + ); + } + + isRunning(): boolean { + return this.lifecycle.kind === "running"; + } + + isStopping(): boolean { + return this.lifecycle.kind === "stopping"; + } + + stopForCancellation( + stop: () => Effect.Effect, + onStarted: () => void, + ): Effect.Effect { + return this.stop({ kind: "cancel" }, stop, onStarted); + } + + stopAtExecutionCeiling( + error: string, + stop: () => Effect.Effect, + onStarted: () => void, + ): Effect.Effect { + return this.stop({ kind: "execution-ceiling", error }, stop, onStarted); + } + + stoppingChild(): ChildSession | undefined { + return this.lifecycle.kind === "stopping" + ? this.lifecycle.ownership?.child + : undefined; + } + + releaseStoppingChild(child: ChildSession): boolean { + const lifecycle = this.lifecycle; + if (lifecycle.kind !== "stopping" || lifecycle.ownership?.child !== child) { + return false; + } + if (lifecycle.ownership.kind === "running") { + lifecycle.ownership.unsubscribe(); + } + this.lifecycle = { + kind: "stopping", + task: lifecycle.task, + reason: lifecycle.reason, + }; + return true; + } + + settleDone(settledAt: number, settlementOrder: number): SettlementTransition { + if (this.lifecycle.kind !== "running") return { kind: "unchanged" }; + return this.settle( + settledAt, + settlementOrder, + { kind: "done" }, + this.lifecycle.ownership, + this.lifecycle.timer, + ); + } + + settleError( + error: string, + checkpoint: string, + settledAt: number, + settlementOrder: number, + ): SettlementTransition { + const lifecycle = this.lifecycle; + if (lifecycle.kind === "creating") { + return this.settle( + settledAt, + settlementOrder, + { kind: "error", error, checkpoint }, + undefined, + lifecycle.timer, + ); + } + if (lifecycle.kind === "subscribing" || lifecycle.kind === "running") { + return this.settle( + settledAt, + settlementOrder, + { kind: "error", error, checkpoint }, + lifecycle.ownership, + lifecycle.timer, + ); + } + return { kind: "unchanged" }; + } + + settleCancelled( + error: string, + checkpoint: string, + settledAt: number, + settlementOrder: number, + ): SettlementTransition { + if (this.lifecycle.kind !== "running") return { kind: "unchanged" }; + return this.settle( + settledAt, + settlementOrder, + { kind: "cancelled", error, checkpoint }, + this.lifecycle.ownership, + this.lifecycle.timer, + ); + } + + settleStopping( + checkpoint: string, + settledAt: number, + settlementOrder: number, + ): SettlementTransition { + const lifecycle = this.lifecycle; + if (lifecycle.kind !== "stopping") return { kind: "unchanged" }; + const outcome: SettledOutcome = + lifecycle.reason.kind === "execution-ceiling" + ? { + kind: "error", + error: lifecycle.reason.error, + checkpoint, + } + : { + kind: "cancelled", + error: "Delegation cancelled", + checkpoint, + }; + return this.settle( + settledAt, + settlementOrder, + outcome, + lifecycle.ownership, + ); + } + + claimDelivery(): boolean { + if (this.delivery.kind !== "pending") return false; + this.delivery = { + kind: "pending", + waiters: this.delivery.waiters + 1, + }; + return true; + } + + releaseDeliveryClaim(): boolean { + if (this.delivery.kind !== "pending" || this.delivery.waiters === 0) { + return false; + } + const waiters = this.delivery.waiters - 1; + this.delivery = { kind: "pending", waiters }; + return waiters === 0 && this.lifecycle.kind === "settled"; + } + + consumeDelivery(): void { + if (this.delivery.kind === "pending") { + this.delivery = { kind: "consumed" }; + } + } + + shouldDeliverSettlement(): boolean { + return this.delivery.kind === "pending" && this.delivery.waiters === 0; + } + + private stop( + reason: StopReason, + stop: () => Effect.Effect, + onStarted: () => void, + ): Effect.Effect { + return Effect.suspend(() => { + const lifecycle = this.lifecycle; + if (lifecycle.kind === "settled") return Effect.void; + if (lifecycle.kind === "stopping") return Fiber.join(lifecycle.task); + + cancelTimer(lifecycle.timer); + const start = Deferred.makeUnsafe(); + const task = Effect.runFork( + Deferred.await(start).pipe(Effect.andThen(stop())), + ); + this.lifecycle = { + kind: "stopping", + task, + reason, + ownership: + lifecycle.kind === "running" + ? { kind: "running", ...lifecycle.ownership } + : lifecycle.kind === "subscribing" + ? { kind: "subscribing", ...lifecycle.ownership } + : undefined, + }; + onStarted(); + Effect.runSync(Deferred.succeed(start, undefined)); + return Fiber.join(task); + }); + } + + private settle( + settledAt: number, + settlementOrder: number, + outcome: SettledOutcome, + ownership?: SubscribingChild | OwnedChild | StoppingOwnership, + timer?: ExecutionTimer, + ): SettlementTransition { + if (timer !== undefined) cancelTimer(timer); + this.lifecycle = { + kind: "settled", + settledAt, + settlementOrder, + outcome, + }; + if (ownership && "unsubscribe" in ownership) ownership.unsubscribe(); + return { kind: "settled", child: ownership?.child }; + } +} diff --git a/agent/extensions/delegate/manager.ts b/agent/extensions/delegate/manager.ts index 11654cd..570667c 100644 --- a/agent/extensions/delegate/manager.ts +++ b/agent/extensions/delegate/manager.ts @@ -11,13 +11,13 @@ import { ChildState } from "./child-state.ts"; import { type DelegateEffort, type DelegateSnapshot, - type DelegateStatus, type DelegateThinking, MAX_EXECUTION_MS, MAX_EXECUTION_TOKENS, } from "./contract.ts"; import { delegateError, errorMessage } from "./errors.ts"; -import { cancelTimer, scheduleTimer } from "./host-timers.ts"; +import { scheduleTimer } from "./host-timers.ts"; +import { RunState, type SettlementTransition } from "./manager-state.ts"; import { type ChildSession, createChild, @@ -45,38 +45,26 @@ export interface DelegateRequest { ctx: ExtensionContext; } -interface Job { - id: string; - task: string; - cwd: string; - effort: DelegateEffort; - thinking: DelegateThinking; - outputFormat?: string; - ctx: ExtensionContext; - requestedModel: string; - fallbackReason?: string; - modelChoice: ExtensionContext["model"]; +interface Run { + readonly id: string; + readonly task: string; + readonly cwd: string; + readonly effort: DelegateEffort; + readonly thinking: DelegateThinking; + readonly outputFormat?: string; + readonly ctx: ExtensionContext; + readonly requestedModel: string; + readonly fallbackReason?: string; + readonly modelChoice: ExtensionContext["model"]; model?: string; - status: DelegateStatus; - createdAt: number; - settledAt?: number; - settlementOrder: number; - error?: string; - childState: ChildState; - child?: ChildSession; - unsubscribe?: () => void; - stopping?: boolean; - stopTask?: Fiber.Fiber; - completion: Deferred.Deferred; - ownership: AbortController; - sendSemaphore: Semaphore.Semaphore; + readonly createdAt: number; + readonly childState: ChildState; + readonly completion: Deferred.Deferred; + readonly ownership: AbortController; + readonly sendSemaphore: Semaphore.Semaphore; pendingSends: number; - deliveryPending: boolean; - deliveryWaiters: number; waiters: number; - hardTimer?: ReturnType; - hardLimitError?: string; - checkpoint?: string; + readonly state: RunState; } export interface DelegateManagerOptions { @@ -137,7 +125,7 @@ function waitUntil( export class DelegateManager { // The product contract deliberately admits every run immediately and retains it for the parent session; the user accepts unbounded aggregate use instead of backpressure or eviction. - private readonly jobs = new Map(); + private readonly jobs = new Map(); private readonly createSession?: DelegateManagerOptions["createSession"]; private readonly shutdownSession?: DelegateManagerOptions["shutdownSession"]; private readonly onSettled?: (snapshot: DelegateSnapshot) => void; @@ -194,7 +182,17 @@ export class DelegateManager { }), ); const effort = request.effort === "thorough" ? "thorough" : "fast"; - const job: Job = { + let job: Run; + const timer = scheduleTimer( + () => + this.stopAtHardLimit( + job, + `${MAX_EXECUTION_MS / 60_000} minutes of wall time`, + ), + MAX_EXECUTION_MS, + ); + timer.unref?.(); + job = { id: `delegate-${++this.nextId}`, task: request.task, cwd, @@ -206,29 +204,23 @@ export class DelegateManager { fallbackReason: modelChoice.fallbackReason, modelChoice: modelChoice.model, model: modelName(modelChoice.model), - status: "running", createdAt: Effect.runSync(Clock.currentTimeMillis), - settlementOrder: 0, childState: new ChildState(), completion: Deferred.makeUnsafe(), ownership: new AbortController(), sendSemaphore: Semaphore.makeUnsafe(1), pendingSends: 0, - deliveryPending: request.background === true, - deliveryWaiters: 0, waiters: 0, + state: RunState.creating(timer, request.background === true), }; this.jobs.set(job.id, job); - this.startExecutionBudget(job); const snapshot = this.snapshot(job); this.notify(snapshot); const task = Effect.runFork( this.run(job).pipe( Effect.catchCause((cause) => Effect.sync(() => { - if (job.status === "running" && !job.stopping) { - this.finalize(job, "error", errorMessage(Cause.squash(cause))); - } + this.settleError(job, errorMessage(Cause.squash(cause))); }), ), ), @@ -243,10 +235,11 @@ export class DelegateManager { return [...new Set(ids)].map((id) => this.snapshot(this.requireJob(id))); } return [...this.jobs.values()] - .sort((a, b) => { - const active = (job: Job) => (job.status === "running" ? 0 : 1); - return active(a) - active(b) || b.settlementOrder - a.settlementOrder; - }) + .sort( + (a, b) => + Number(!a.state.isActive()) - Number(!b.state.isActive()) || + b.state.settlementOrder() - a.state.settlementOrder(), + ) .map((job) => this.snapshot(job)); } @@ -271,15 +264,11 @@ export class DelegateManager { ); } for (const job of jobs) job.waiters++; - const claims = jobs.filter((job) => { - if (!job.deliveryPending) return false; - job.deliveryWaiters++; - return true; - }); + const claims = jobs.filter((job) => job.state.claimDelivery()); let completed = false; return yield* Effect.all( jobs.map((job) => - job.status === "running" + job.state.isActive() ? Deferred.await(job.completion) : Effect.succeed(this.snapshot(job)), ), @@ -294,7 +283,7 @@ export class DelegateManager { Effect.tap((snapshots) => Effect.sync(() => { completed = true; - for (const job of claims) job.deliveryPending = false; + for (const job of claims) job.state.consumeDelivery(); return snapshots; }), ), @@ -302,13 +291,7 @@ export class DelegateManager { Effect.sync(() => { for (const job of jobs) job.waiters--; for (const job of claims) { - job.deliveryWaiters--; - if ( - !completed && - job.deliveryWaiters === 0 && - job.deliveryPending && - job.status !== "running" - ) { + if (!completed && job.state.releaseDeliveryClaim()) { this.onSettled?.(this.snapshot(job)); } } @@ -325,26 +308,25 @@ export class DelegateManager { const job = this.requireJob(id); const text = message.trim(); if (!text) throw new Error("Delegate message must not be empty."); - if (job.status !== "running") { + if (!job.state.isActive()) { throw new Error( - `Delegate ${id} is ${job.status}; send requires a running child.`, + `Delegate ${id} is ${job.state.status()}; send requires a running child.`, ); } - if (!job.child) throw new Error(`Delegate ${id} has no active session.`); + const child = job.state.runningChild(); + if (!child) throw new Error(`Delegate ${id} has no active session.`); if (job.pendingSends >= MAX_PENDING_SENDS) { throw new Error( `Delegate ${id} already has ${MAX_PENDING_SENDS} pending messages.`, ); } - const child = job.child; job.pendingSends++; yield* job.sendSemaphore .withPermit( Effect.gen( function* (this: DelegateManager) { if ( - job.status !== "running" || - job.child !== child || + !job.state.ownsRunningChild(child) || job.ownership.signal.aborted ) { throw new Error( @@ -366,7 +348,7 @@ export class DelegateManager { ids: readonly string[], ) { const jobs = [...new Set(ids)].map((id) => this.requireJob(id)); - for (const job of jobs) job.deliveryPending = false; + for (const job of jobs) job.state.consumeDelivery(); yield* Effect.all( jobs.map((job) => this.stopOwned(job)), { @@ -378,8 +360,7 @@ export class DelegateManager { acknowledge(ids: readonly string[]) { for (const id of new Set(ids)) { - const job = this.jobs.get(id); - if (job) job.deliveryPending = false; + this.jobs.get(id)?.state.consumeDelivery(); } } @@ -396,7 +377,7 @@ export class DelegateManager { const jobs = [...this.jobs.values()]; yield* waitUntil( jobs.map((job) => - job.status === "running" ? this.stopOwned(job) : Effect.void, + job.state.isActive() ? this.stopOwned(job) : Effect.void, ), deadline, ); @@ -408,92 +389,89 @@ export class DelegateManager { private readonly run = Effect.fn("DelegateManager.run")(function* ( this: DelegateManager, - job: Job, + job: Run, ) { let receivedAssistantResponse = false; - if (!job.child) { - const request = { - task: job.task, - cwd: job.cwd, - effort: job.effort, - outputFormat: job.outputFormat, - ctx: job.ctx, - }; - job.child = yield* this.createSession - ? Effect.tryPromise({ - try: () => - this.createSession?.( - request, - job.modelChoice, - job.thinking, - job.ownership.signal, - ) as Promise, - catch: delegateError, - }) - : createChild(request.cwd, job.modelChoice, job.thinking).pipe( - Effect.mapError(delegateError), - ); - if (job.status !== "running") { - const child = job.child; - job.child = undefined; - yield* this.disposeOwned(child, job.id); - return; - } - job.model = modelName(job.child.model ?? job.modelChoice); - job.unsubscribe = job.child.subscribe((event) => { - if (isAssistantResponse(event)) receivedAssistantResponse = true; - this.onEvent(job, event); - }); + const request = { + task: job.task, + cwd: job.cwd, + effort: job.effort, + outputFormat: job.outputFormat, + ctx: job.ctx, + }; + const createSession = this.createSession; + const child = yield* createSession + ? Effect.tryPromise({ + try: () => + createSession( + request, + job.modelChoice, + job.thinking, + job.ownership.signal, + ), + catch: delegateError, + }) + : createChild(request.cwd, job.modelChoice, job.thinking).pipe( + Effect.mapError(delegateError), + ); + if (!job.state.isActive()) { + yield* this.disposeOwned(child, job.id); + return; + } + job.model = modelName(child.model ?? job.modelChoice); + if (!job.state.startSubscribing(child)) { + yield* this.disposeOwned(child, job.id); + return; + } + const unsubscribe = child.subscribe((event) => { + if (isAssistantResponse(event)) receivedAssistantResponse = true; + this.onEvent(job, event); + }); + if (!job.state.startRunning(child, unsubscribe)) { + unsubscribe(); + yield* this.disposeOwned(child, job.id); + return; } - const child = job.child; const outputFormat = job.outputFormat?.trim(); const instruction = outputFormat ? `${job.task}\n\nPreferred output format (advisory):\n${outputFormat}\n\nPrioritize correct and complete information over exact formatting.` : job.task; - const outcome = yield* this.untilOwnershipEnds( + const promptOutcome = yield* this.untilOwnershipEnds( job, child.prompt(instruction, { expandPromptTemplates: false, source: "extension", }), ).pipe(Effect.exit); - if (outcome._tag === "Failure") { - if (job.status === "running" && !job.stopping) { - this.finalize(job, "error", errorMessage(Cause.squash(outcome.cause))); - } + if (promptOutcome._tag === "Failure") { + this.settleError(job, errorMessage(Cause.squash(promptOutcome.cause))); return; } if (!receivedAssistantResponse) { - this.finalize( + this.settleError( job, - "error", `Delegate ${job.id} finished without an assistant response. Retry the delegation.`, ); return; } - if (job.status !== "running" || job.stopping) return; + if (!job.state.isRunning()) return; const childState = job.childState.state(); if (childState.assistantStop === "error") { - this.finalize( - job, - "error", - childState.assistantError ?? "Child agent failed.", - ); + this.settleError(job, childState.assistantError ?? "Child agent failed."); return; } if (childState.assistantStop === "aborted") { - this.finalize( + this.settleCancelled( job, - "cancelled", childState.assistantError ?? "Child agent aborted.", ); return; } - this.finalize(job, "done"); + this.settleDone(job); }); - private onEvent(job: Job, event: Parameters[0]) { + private onEvent(job: Run, event: Parameters[0]) { job.childState.capture(event); if (job.childState.state().usage.totalTokens >= MAX_EXECUTION_TOKENS) { this.stopAtHardLimit( @@ -504,52 +482,42 @@ export class DelegateManager { this.notify(this.snapshot(job)); } - private startExecutionBudget(job: Job) { - job.hardTimer = scheduleTimer( + private stopAtHardLimit(job: Run, limit: string) { + if (!job.state.isActive() || job.state.isStopping()) return; + const error = `Delegation stopped at the hard execution ceiling: ${limit}.`; + Effect.runFork(this.stopOwnedAtExecutionCeiling(job, error)); + } + + // RunState starts one root stop fiber. Interrupted observers, including an + // aborted cancel or shutdown deadline, only join it and cannot poison it. + private stopOwned(job: Run): Effect.Effect { + return job.state.stopForCancellation( + () => this.stop(job), () => - this.stopAtHardLimit( + this.endOwnership( job, - `${MAX_EXECUTION_MS / 60_000} minutes of wall time`, + new Error(`Delegate ${job.id} ownership ended.`), ), - MAX_EXECUTION_MS, ); - job.hardTimer.unref?.(); - } - - private stopAtHardLimit(job: Job, limit: string) { - if (job.status !== "running" || job.stopping) return; - job.hardLimitError = `Delegation stopped at the hard execution ceiling: ${limit}.`; - Effect.runFork(this.stopOwned(job)); } - private clearExecutionBudget(job: Job) { - if (job.hardTimer !== undefined) cancelTimer(job.hardTimer); - job.hardTimer = undefined; - } - - // The stop runs on its own root fiber so an interrupted observer (an - // aborted cancel, a shutdown deadline) cannot poison the shared stop for - // later callers; Fiber.join only attaches an observer. - private stopOwned(job: Job): Effect.Effect { - return Effect.suspend(() => { - job.stopTask ??= Effect.runFork(this.stop(job)); - return Fiber.join(job.stopTask); - }); + private stopOwnedAtExecutionCeiling( + job: Run, + error: string, + ): Effect.Effect { + return job.state.stopAtExecutionCeiling( + error, + () => this.stop(job), + () => this.endOwnership(job, new Error(error)), + ); } private readonly stop = Effect.fn("DelegateManager.stop")(function* ( this: DelegateManager, - job: Job, + job: Run, ) { - this.clearExecutionBudget(job); - if (job.status !== "running" || job.stopping) return; - job.stopping = true; - this.endOwnership( - job, - new Error(job.hardLimitError ?? `Delegate ${job.id} ownership ended.`), - ); - if (job.child) { - const child = job.child; + const child = job.state.stoppingChild(); + if (child) { let abortFailure: unknown; const stopped = yield* Effect.tryPromise({ try: () => child.abort(), @@ -574,63 +542,112 @@ export class DelegateManager { `[delegate] abort failed for ${job.id}: ${evidence}`, ); } - if (!stopped || child.isStreaming) { - job.child = undefined; - job.unsubscribe?.(); - job.unsubscribe = undefined; + if ( + (!stopped || child.isStreaming) && + job.state.releaseStoppingChild(child) + ) { yield* this.disposeOwned(child, job.id); } } - if (job.status === "running") { - const child = job.child; - this.finalize( - job, - job.hardLimitError ? "error" : "cancelled", - job.hardLimitError ?? "Delegation cancelled", - ); - if (child) yield* this.disposeOwned(child, job.id); + if (!job.state.isStopping()) return; + const settledAt = yield* Clock.currentTimeMillis; + const settlementOrder = this.nextSettlementOrder + 1; + const transition = job.state.settleStopping( + this.checkpoint(job), + settledAt, + settlementOrder, + ); + this.publishSettlement(job, settlementOrder, transition); + if (transition.kind === "settled" && transition.child) { + yield* this.disposeOwned(transition.child, job.id); } }); - private finalize(job: Job, status: DelegateStatus, error?: string) { - if (status !== "done") { - job.checkpoint = truncateUtf8Tail( - job.childState.trail().join("\n\n"), - MAX_CHECKPOINT_BYTES, - ); - } - this.clearExecutionBudget(job); + private checkpoint(job: Run) { + return truncateUtf8Tail( + job.childState.trail().join("\n\n"), + MAX_CHECKPOINT_BYTES, + ); + } + + private settleDone(job: Run) { + const settlementOrder = this.nextSettlementOrder + 1; + this.publishSettlement( + job, + settlementOrder, + job.state.settleDone( + Effect.runSync(Clock.currentTimeMillis), + settlementOrder, + ), + ); + } + + private settleError(job: Run, error: string) { + const settlementOrder = this.nextSettlementOrder + 1; + this.publishSettlement( + job, + settlementOrder, + job.state.settleError( + error, + this.checkpoint(job), + Effect.runSync(Clock.currentTimeMillis), + settlementOrder, + ), + ); + } + + private settleCancelled(job: Run, error: string) { + const settlementOrder = this.nextSettlementOrder + 1; + this.publishSettlement( + job, + settlementOrder, + job.state.settleCancelled( + error, + this.checkpoint(job), + Effect.runSync(Clock.currentTimeMillis), + settlementOrder, + ), + ); + } + + private publishSettlement( + job: Run, + settlementOrder: number, + transition: SettlementTransition, + ) { + if (transition.kind === "unchanged") return; + this.nextSettlementOrder = settlementOrder; this.endOwnership(job, new Error(`Delegate ${job.id} ownership ended.`)); - job.status = status; - job.settledAt = Effect.runSync(Clock.currentTimeMillis); - job.settlementOrder = ++this.nextSettlementOrder; - job.error = error; - job.stopping = undefined; const snapshot = this.snapshot(job); Effect.runSync(Deferred.succeed(job.completion, snapshot)); this.notify(snapshot); - if (job.deliveryPending && job.deliveryWaiters === 0) { - this.onSettled?.(snapshot); + if (job.state.shouldDeliverSettlement()) this.onSettled?.(snapshot); + if (transition.child) { + Effect.runFork(this.disposeOwned(transition.child, job.id)); } - const child = job.child; - job.child = undefined; - job.unsubscribe?.(); - job.unsubscribe = undefined; - if (child) Effect.runFork(this.disposeOwned(child, job.id)); if (this.disposed) job.childState.cleanup(); } - private snapshot(job: Job): DelegateSnapshot { + private snapshot(job: Run): DelegateSnapshot { const childState = job.childState.state(); + const state = job.state.view(); + const status = state.status; + const settledAt = status === "running" ? undefined : state.settledAt; + const error = + status === "error" || status === "cancelled" ? state.error : undefined; + const checkpoint = + status === "error" || status === "cancelled" + ? state.checkpoint || undefined + : undefined; return { id: job.id, - status: job.status, + status, createdAt: job.createdAt, - settledAt: job.settledAt, + settledAt, output: childState.output, outputTruncated: childState.outputTruncated, fullOutputFile: childState.fullOutputFile, - success: job.status === "done", + success: status === "done", assignedTask: job.task, effort: job.effort, requestedModel: job.requestedModel, @@ -638,19 +655,18 @@ export class DelegateManager { thinking: job.thinking, fallbackReason: job.fallbackReason, durationMs: - (job.settledAt ?? Effect.runSync(Clock.currentTimeMillis)) - - job.createdAt, + (settledAt ?? Effect.runSync(Clock.currentTimeMillis)) - job.createdAt, toolCalls: childState.toolCalls, failedToolCalls: childState.failedToolCalls, childUsage: childState.usage, - aborted: job.status === "cancelled", - error: job.error, - progress: job.status === "running" ? childState.progress : undefined, + aborted: status === "cancelled", + error, + progress: status === "running" ? childState.progress : undefined, idleMs: - job.status === "running" + status === "running" ? Effect.runSync(Clock.currentTimeMillis) - childState.lastActivityAt : undefined, - checkpoint: job.checkpoint || undefined, + checkpoint, }; } @@ -664,17 +680,17 @@ export class DelegateManager { } } - private requireJob(id: string): Job { + private requireJob(id: string): Run { const job = this.jobs.get(id); if (!job) throw new Error(`Unknown delegate id "${id}".`); return job; } - private endOwnership(job: Job, reason: Error) { + private endOwnership(job: Run, reason: Error) { if (!job.ownership.signal.aborted) job.ownership.abort(reason); } - private untilOwnershipEnds(job: Job, operation: Promise) { + private untilOwnershipEnds(job: Run, operation: Promise) { return Effect.tryPromise({ try: () => operation, catch: delegateError, @@ -684,9 +700,10 @@ export class DelegateManager { private disposeOwned(child: ChildSession, id: string): Effect.Effect { const existing = this.childDisposals.get(child); if (existing) return Fiber.join(existing); - const shutdown = this.shutdownSession + const shutdownSession = this.shutdownSession; + const shutdown = shutdownSession ? Effect.tryPromise({ - try: () => this.shutdownSession?.(child) as Promise, + try: () => shutdownSession(child), catch: delegateError, }) : shutdownChild(child).pipe(Effect.mapError(delegateError)); diff --git a/agent/extensions/delegate/test/manager.test.ts b/agent/extensions/delegate/test/manager.test.ts index d4771d3..7039782 100644 --- a/agent/extensions/delegate/test/manager.test.ts +++ b/agent/extensions/delegate/test/manager.test.ts @@ -12,7 +12,11 @@ import test from "node:test"; import type { AgentSessionEvent, ExtensionContext } from "@earendil-works/pi-coding-agent"; import { Cause, Deferred, Effect, Fiber } from "effect"; -import { type DelegateSnapshot, MAX_CHILD_OUTPUT_BYTES } from "../contract.ts"; +import { + type DelegateSnapshot, + MAX_CHILD_OUTPUT_BYTES, + MAX_EXECUTION_TOKENS, +} from "../contract.ts"; import { DelegateManager } from "../manager.ts"; import type { ChildSession } from "../runtime.ts"; import { deferredPromise, eventually, yieldImmediate } from "./eventually.ts"; @@ -243,6 +247,69 @@ test("all effort modes stop at sixty million reported tokens", () => Effect.runP } }))); +test("a child is disposed when subscription throws", () => Effect.runPromise(Effect.gen(function* () { + class ThrowingSubscribeChild extends FakeChild { + disposeCalls = 0; + + override subscribe( + _listener: (event: AgentSessionEvent) => void, + ): ReturnType { + throw new Error("subscription failed"); + } + + override disposeNow() { + this.disposeCalls++; + super.disposeNow(); + } + } + + const child = new ThrowingSubscribeChild(); + const manager = new DelegateManager({ + createSession: () => Promise.resolve(child as unknown as ChildSession), + shutdownSession() { + child.disposeNow(); + return Promise.resolve(); + }, + }); + const job = manager.spawn({ task: "subscription failure", ctx: context }); + + const [failed] = yield* manager.wait([job.id]); + assert.equal(failed.status, "error"); + assert.match(failed.error ?? "", /subscription failed/); + yield* eventually(() => child.disposeCalls === 1); + assert.equal(child.disposeCalls, 1); + yield* manager.shutdown(); + assert.equal(child.disposeCalls, 1); +}))); + +test("a ceiling event delivered during subscription cannot revive a stopped run", () => Effect.runPromise(Effect.gen(function* () { + class EventDuringSubscribeChild extends FakeChild { + override subscribe(listener: (event: AgentSessionEvent) => void) { + const unsubscribe = super.subscribe(listener); + this.emitAssistant("ceiling checkpoint", MAX_EXECUTION_TOKENS); + return unsubscribe; + } + } + + const child = new EventDuringSubscribeChild(); + const manager = new DelegateManager({ + createSession: () => Promise.resolve(child as unknown as ChildSession), + shutdownSession() { + child.disposeNow(); + return Promise.resolve(); + }, + }); + const job = manager.spawn({ task: "subscription race", ctx: context }); + + const [failed] = yield* manager.wait([job.id]); + assert.equal(failed.status, "error"); + assert.match(failed.error ?? "", /60,000,000 reported tokens/); + yield* yieldImmediate; + assert.deepEqual(child.prompts, []); + yield* eventually(() => child.disposed); + yield* manager.shutdown(); +}))); + test("cancellation releases prompts that ignore child abort", () => Effect.runPromise(Effect.gen(function* () { const { manager, sessions } = harness(); const jobs = Array.from({ length: 4 }, (_, index) => @@ -443,6 +510,44 @@ test("cancel consumption wins over an aborted concurrent wait", () => { })); }); +test("shutdown wins once child settlement races an owned stop", () => Effect.runPromise(Effect.gen(function* () { + const delivered: DelegateSnapshot[] = []; + const terminalNotifications: DelegateSnapshot[] = []; + let disposals = 0; + const { manager, sessions } = harness( + (snapshot) => delivered.push(snapshot), + () => Effect.sync(() => { + disposals++; + }), + ); + manager.subscribe((snapshot) => { + if (snapshot.status !== "running") terminalNotifications.push(snapshot); + }); + const job = manager.spawn({ + task: "settle during shutdown", + background: true, + ctx: context, + }); + yield* eventually(() => sessions.length === 1); + const abortGate = yield* Deferred.make(); + sessions[0].abortGate = abortGate; + + const shutdown = yield* manager.shutdown().pipe(Effect.forkChild); + yield* eventually(() => sessions[0].abortCalls === 1); + sessions[0].finish("too late"); + yield* Deferred.succeed(abortGate, undefined); + yield* Fiber.join(shutdown); + + const [snapshot] = manager.list([job.id]); + assert.equal(snapshot.status, "cancelled"); + assert.equal(snapshot.output, "too late"); + assert.equal(delivered.length, 1); + assert.equal(delivered[0].status, "cancelled"); + assert.equal(terminalNotifications.length, 1); + assert.equal(terminalNotifications[0].status, "cancelled"); + assert.equal(disposals, 1); +}))); + test("concurrent shutdown joins gated child disposal", () => Effect.runPromise(Effect.gen(function* () { const disposalGate = yield* Deferred.make(); let disposalStarted = false;