diff --git a/packages/agent/docs/assistant-durability.md b/packages/agent/docs/assistant-durability.md index 1c9edb2f4..d79f97458 100644 --- a/packages/agent/docs/assistant-durability.md +++ b/packages/agent/docs/assistant-durability.md @@ -192,7 +192,7 @@ A deferred poll that returns an `AssistantMessageEventStream` uses the same `pen - when a poll permit replaces the unknown poll with fresh reserved response/usage IDs, that intent transaction deletes the abandoned old frame-list address; - the replacement poll starts a fresh list under its fresh response ID. -Frames never change deferred-handle equality or poll-number rules. +Frames never change poll-number rules. ## Structural generation scope diff --git a/packages/agent/docs/harness.md b/packages/agent/docs/harness.md index ab3ea9a4a..938267a1e 100644 --- a/packages/agent/docs/harness.md +++ b/packages/agent/docs/harness.md @@ -1427,9 +1427,9 @@ type Deferred = configuration: LaneConfiguration; streamOptions: AgentHarnessStreamOptions }; ``` -One `drive({ pollDeferred: true }, context)` performs at most one `fetchDeferred(handle, { wait: 0 })`; `resume(context)` is the convenience call that supplies that permit. Suspended `poll` is the number of completed polls; a fresh intent uses `poll + 1`, and that 1-based value is `before_request.attempt` and the poll turn-id suffix. A poll starts from the original generation's copied base stream options, forces `deferred:false`, runs `before_request`, mounts `before_payload`/`after_response`, then commits its fresh intent and dispatches like assistant generation. Current global stream settings do not affect it. There is no polling retry cap, backoff, or internal loop. A `deferred` response has a valid handle only when `message.deferred` exists, its id is non-empty, and its `{ provider, modelId, api }` equals the captured request model identity/API. An invalid initial handle is normalized to a durable `error` response and response-provenance `failure_drain`. A later pending response must carry a handle completely equal to its source handle and becomes the next source; a mismatch is normalized the same way. Response, usage, `latestAssistantEntryId`, and failure state commit atomically. +One `drive({ pollDeferred: true }, context)` performs at most one `Models.streamDeferred(handle, { wait: 0 })` poll; `resume(context)` is the convenience call that supplies that permit. `Models.streamDeferred` preserves the provider event stream required for frame persistence, while `Models.fetchDeferred` remains its result-only convenience and is not the harness runtime path. Suspended `poll` is the number of completed polls; a fresh intent uses `poll + 1`, and that 1-based value is `before_request.attempt` and the poll turn-id suffix. A poll starts from the original generation's copied base stream options, forces `deferred:false`, runs `before_request`, mounts `before_payload`/`after_response`, then commits its fresh intent and dispatches like assistant generation. Current global stream settings do not affect it. There is no polling retry cap, backoff, or internal loop. A `deferred` response has a valid handle only when `message.deferred` exists, its id is non-empty, and its `{ provider, modelId, api }` equals the captured request model identity/API. An invalid initial handle is normalized to a durable `error` response and response-provenance `failure_drain`. A later pending response becomes the next source. Response, usage, `latestAssistantEntryId`, and failure state commit atomically. -A deferred poll that returns an event stream persists frames exactly like assistant generation, at `pendingAssistantFrames(operationId, responseEntryId)` for the poll's reserved response id (§3.7). Every row below that inserts a response entry additionally deletes that response's frame list in the same transaction. Frames never change deferred-handle equality or poll-number rules. +A deferred poll that returns an event stream persists frames exactly like assistant generation, at `pendingAssistantFrames(operationId, responseEntryId)` for the poll's reserved response id (§3.7). Every row below that inserts a response entry additionally deletes that response's frame list in the same transaction. Frames never change poll-number rules. The complete transition table — every row is one `commit()`; classification order (§3.7) applies to every poll settlement, cancellation first: @@ -1438,8 +1438,7 @@ The complete transition table — every row is one `commit()`; classification or | assistant `effect_pending` | settlement classifies `deferred` with a valid handle | §3.7's deferred row | suspended, `poll: 0`, `sourceEntryId: R` | | suspended, poll _k_ | captured model is unavailable when a drive supplies `pollDeferred: true` | `TX[ S(failure_drain{error:{code:"model_unavailable", ...}, provenance:{kind:"configuration"}}) ]` — the handle is abandoned without fabricating a fetch response or usage | failure_drain | | suspended, poll _k_ | captured model resolves and a drive supplies `pollDeferred: true`: the poll's `before_request` settlement commits its intent, consuming the pass's single poll permit | mint fresh R′ and U′, then `TX[ S(deferred{effect_pending, poll k+1, responseEntryId R′, usageId U′}) ]` | effect_pending, poll _k_+1 | -| effect_pending, poll _k_+1 | fetch returns **pending** with a completely equal handle | `TX[ insert response entry R′, insert usage U′, upsert pi.branch.tip = R′, S(latestAssistantEntryId=R′, deferred{suspended, sourceEntryId R′, poll k+1}) ]` — the pending response becomes the next source and the operation re-suspends; no second poll this invocation | suspended, poll _k_+1 | -| effect_pending | fetch returns **pending** with a mismatched handle | normalize to a durable `error` response explaining the mismatch: `TX[ insert normalized response R′, insert usage U′, upsert pi.branch.tip = R′, S(latestAssistantEntryId=R′, failure_drain{error, provenance:response R′}) ]` | failure_drain | +| effect_pending, poll _k_+1 | fetch returns **pending** | `TX[ insert response entry R′, insert usage U′, upsert pi.branch.tip = R′, S(latestAssistantEntryId=R′, deferred{suspended, sourceEntryId R′, poll k+1}) ]` — the pending response becomes the next source and the operation re-suspends; no second poll this invocation | suspended, poll _k_+1 | | effect_pending | fetch returns **ready** with tool calls | `TX[ insert response R′, insert usage U′, upsert pi.branch.tip = R′, S(latestAssistantEntryId=R′, tools{plan with reserved result ids}) ]` — result ids minted as followers of R′ (§1.2) | tools | | effect_pending | fetch returns **ready** without tool calls | `TX[ insert response R′, insert usage U′, upsert pi.branch.tip = R′, S(latestAssistantEntryId=R′, checkpoint{may_finish, includeFinalAssistant:true}) ]` | checkpoint | | effect_pending | fetch settles as a provider `error` | `TX[ insert response R′, insert usage U′, upsert pi.branch.tip = R′, S(latestAssistantEntryId=R′, failure_drain{error, provenance:response R′}) ]` — polls have no retry path | failure_drain | @@ -2142,7 +2141,7 @@ The complete admission catalog is: **Accepted hook aggregates:** `before_drive`, `before_run`, `before_run_end`, `transform_context`, `before_request`, `before_payload`, `after_response`, `before_tool`, `after_tool`, `before_compaction`, `before_navigation`. One `admit` wraps the complete registered pipeline, not each handler. -**Provider operations:** one assistant `Models.streamSimple`, each individual structural-summary request, and one explicit `Models.fetchDeferred` poll. Best-effort `cancelDeferred` is cancellation cleanup and uses its separate close-only signal. +**Provider operations:** one assistant `Models.streamSimple`, each individual structural-summary request, and one explicit `Models.streamDeferred` poll. Best-effort `cancelDeferred` is cancellation cleanup and uses its separate close-only signal. **Other ordinary operations:** one real `tool.execute` and creation of each assistant/structural retry timer. Unknown, invalid, blocked, and synthetic tool outcomes start no tool and use no gate. @@ -3531,7 +3530,7 @@ These rows preserve requirements, not frozen package boundaries. Create their ha | R4 | **Tools** | Add the harness-specific tool-update checkpoint options/callback while leaving pi-ai's `AgentToolUpdateCallback` unchanged, then compose §5.7's neutral tool phases with bound context, invocation-scoped memos, `pi.op.tool_args`, scalar `pendingToolOutput(operationId, invocationId)` checkpoints, `outcome_ready` staging in `pi.pending.entry`, completion-order outcome durability, source-ordered materialization, orphan recovery/replay, sequential/parallel promises, hook-blocked/genuine-length outcomes, hooks/events/usage, and deterministic transition coverage. Missing requested tools stage direct `ToolResultMessage` errors with no `details`; missing safe-replay implementations synthesize interruption rather than suspending. The built-in bash tool emits live updates at 100 ms and requests distinct checkpoints at most every two seconds. | Stable invocation id/memos across safe replay; every planned/effect-pending/outcome-ready/completed state; parallel out-of-order staging without replay; synchronous update callback with the latest delivery and latest checkpoint write awaited before `after_tool`; `tool_end` before staging; every requested checkpoint enqueued synchronously with only the latest promise reference retained and no write dropped or coalesced; tool-selected cadence/duplicate suppression; late-write fencing after settlement; JSONL checkpoint growth/compaction; safe replay deletes old checkpoint; unsafe recovery with/without checkpoint; missing planned tools become source-ordered error results; absent or no-longer-safe replay implementations synthesize interruption without invented typed details; source-ordered materialization; termination and abort-ready states. | | R5 | **Inbox, configuration, and writes** | `nextRun`/steer/follow-up via `pendingEntry(id)` values, pending-payload consumption-time dereference, `cancelQueued` triage (`not_found`), durable drain markers, checkpoint consumption with pending-value deletion, immediate total config setters, deferred tree writes, adjustments, and deterministic tests for the checkpoint boundaries introduced here. | Capture/cancel/consume races, repeated cancellation answering `not_found`, invalid queue payload kinds and missing pending values rejected at consumption, one-at-a-time crash after one drain, pending-value/entry exclusivity at every boundary, custom-write continuation, config-step race, writes surviving reopen. | | R6 | **Abort, close, and failure drain** | Add expected-id `requestAbort` as commit/signalling only. Compose `abort()` from inspection, marker, and same-id drive. Implement orthogonal control, per-phase reconciliation including planned/effect-pending/outcome-ready tools, best-effort deferred cancellation, waiters/run-when-idle, controlled-crash close, complete terminal deletion of memos/checkpoints/staged outcomes, external-finalization stop, and deterministic transition coverage. | Cancellation with/without live pass; idempotency/fencing; abort at every tool state; restored interruption incorporates latest checkpoint without replay; restored cancellation synthesizes aborted responses preserving frame-reduced partial content and deleting the list; outcome-ready values survive and materialize; checkpoint/memo/frame cleanup; gate ordering; deferred cancellation; close races including a non-cooperative started tool that cannot block close and parallel retained promises that are all observed without later starts or unhandled rejection; external finalization fences late scalar/list writes; failure drain. | -| R7 | **Deferred provider redemption** | One poll per drive pass carrying `pollDeferred:true`; `resume()` composes inspection and that drive option. Add deferred consumption-time dereference and captured model resolution; unavailable models durably abandon redemption through configuration-provenance failure rather than waiting. Preserve copied configuration/options, per-poll hooks, exact source lineage/equality, poll frame persistence with fresh intent after unknown poll deleting the abandoned frame list, mismatch-to-error, ready tools, advancement of R6 cancellation, and deterministic poll-transition coverage. | `pollDeferred:false` returns waiting without work; one permit performs at most one poll; invalid/missing/mismatched durable handles; repeated pending, ready/error/aborted/mismatch, crash positions, poll frame cleanup on every settlement, replacement, and configuration abandonment with no fabricated response/usage, unknown-poll snapshot without permit, convenience/primitive equivalence, no cap/backoff/loop, newest-handle cancellation. | +| R7 | **Deferred provider redemption** | One poll per drive pass carrying `pollDeferred:true`; `resume()` composes inspection and that drive option. Add deferred consumption-time dereference and captured model resolution; unavailable models durably abandon redemption through configuration-provenance failure rather than waiting. Preserve copied configuration/options, per-poll hooks, source lineage, poll frame persistence with fresh intent after unknown poll deleting the abandoned frame list, ready tools, advancement of R6 cancellation, and deterministic poll-transition coverage. | `pollDeferred:false` returns waiting without work; one permit performs at most one poll; invalid or missing durable source handles; repeated pending, ready/error/aborted, crash positions, poll frame cleanup on every settlement, replacement, and configuration abandonment with no fabricated response/usage, unknown-poll snapshot without permit, convenience/primitive equivalence, no cap/backoff/loop, newest-handle cancellation. | | R8 | **Manual compaction** | Add compaction `OperationRequest` acceptance without task installation and compose `compact()` from accept+drive. Implement unowned preparation with final-command source revalidation, `pi.op.preparation/{opId}:{taskId}`, total structural state with consumption-time dereference, in-band structural model-unavailability failure, hook/generated sources, nested request intents/usage, retained tail, retry/recovery/abort, and deterministic structural-transition coverage. | Explicit accept+drive versus convenience equivalence; empty/stale-preparation race; accept/close/reopen/drive; invalid task/preparation/source relationships; hook decline/result; crash after request one of split-turn generation; every state/crash; no public summary-stream messages and no structural frame persistence. | | R9 | **Threshold and overflow compaction** | In-run structural decision, durable once-per-trigger threshold marker, continuation preservation, all overflow predicates, atomic response/preparation publication, specified normalization/projection, one overflow recovery flag, bounded second failure, and deterministic tests for added boundaries. | Threshold decline/empty across reopen, all overflow classifier/preparation inputs, no overflow tool plan, genuine length, crash/reopen at every transition. | | R10 | **Navigation** | Add navigation `OperationRequest` acceptance without task installation and compose `navigateTree()` from accept+drive. Implement navigation source/target/summary consumption-time dereference, validation, summarized decision/generation, one final transaction combining move/summary/tip/label with terminal writes, summary-only navigation hook, and deterministic transition coverage. | Explicit accept+drive versus convenience equivalence; root/current/unknown rejection; invalid durable source/target/summary combinations; accept/close/reopen/drive; summarized/unsummarized paths; final tip at summary; abort race; exact atomic publication including cleanup. | diff --git a/packages/agent/docs/work-packages/05-direct-durable-drive.md b/packages/agent/docs/work-packages/05-direct-durable-drive.md index 3bd5f2a6f..a1a93d653 100644 --- a/packages/agent/docs/work-packages/05-direct-durable-drive.md +++ b/packages/agent/docs/work-packages/05-direct-durable-drive.md @@ -1,6 +1,6 @@ # WP05 — Direct durable drive -**Status: M3 complete; WP06 foundation landed before M4.** Independent Opus, GPT-5.6-Sol, and Fable reviews approved the authoritative-`Lane.state` design and M3 generation implementation with no blockers. +**Status: M4 complete; WP06 foundation landed before M4.** Independent Opus, GPT-5.6-Sol, and Fable reviews approved the authoritative-`Lane.state` design and M3 generation implementation with no blockers. M4 retry/deferred implementation also passed independent Fable review with no blockers. WP06 separates Session/Branch/AgentLane/AgentHarness, replaces keyed lane lines with one keyless Session mutation line, removes implicit main, and renames durable/public leaf fields to tip. Every M4–M8 instruction below is interpreted on that foundation; no compatibility aliases or Harness-as-main assumptions remain. @@ -1353,6 +1353,10 @@ Only `packages/coding-agent/src/core/cache-stats.ts` (comment) and vendored `hig |---|---| | `src/harness/runtime/drive/generation.ts` | retry classification → `retry_wait`; `waitForRetry` policy | | `src/harness/runtime/drive/deferred.ts` | `drive.deferredPermits` consumption lives with the deferred procedure. **Do not touch `drive.ts`** — it does not exist until M7 | +| `packages/ai/src/models.ts` | add the event-stream-preserving `Models.streamDeferred`; retain `fetchDeferred` as `streamDeferred(...).result()` | +| `packages/coding-agent/src/core/model-runtime.ts` | implement the same `Models.streamDeferred` auth-preserving runtime path | + +The two `Models` rows are a necessary implementation prerequisite discovered in M4: the existing promise-returning `fetchDeferred()` discarded provider events, making deferred frame persistence and partial recovery impossible. The additive stream method preserves Models auth/request transforms and avoids bypassing the registry. **Behavior/invariants landed.** No durable armed poll state: `suspended` never becomes a durable pollable `ready`; the permit is consumed by the **fresh intent commit** that reserves new ids; a crash before that commit leaves `suspended` at the same poll number; an unknown-outcome poll does not consume a poll, deletes the old frame list, and mints fresh ids next poll; an unknown-outcome structural attempt *does* consume an attempt. @@ -1741,6 +1745,8 @@ Milestone in parentheses where creation order matters. - `packages/agent/src/harness/runtime/harness.ts` - `packages/agent/src/harness/runtime/types.ts` - `packages/agent/src/harness/runtime/restore.ts` +- `packages/ai/src/models.ts` +- `packages/coding-agent/src/core/model-runtime.ts` - `packages/agent/src/harness/session/session.ts` - `packages/agent/src/harness/session/types.ts` - `packages/agent/src/harness/session/commit.ts` diff --git a/packages/agent/src/harness/runtime/drive/deferred.ts b/packages/agent/src/harness/runtime/drive/deferred.ts new file mode 100644 index 000000000..c2cacaf4b --- /dev/null +++ b/packages/agent/src/harness/runtime/drive/deferred.ts @@ -0,0 +1,627 @@ +import { + type Api, + type AssistantMessageEvent, + AssistantMessageFrameEncoder, + type DeferredHandle, + type Model, +} from "@earendil-works/pi-ai"; +import type { HarnessEvent } from "../../agent-harness.ts"; +import { withAbortSignal } from "../../context.ts"; +import { AbortRequested } from "../../execution/effect-gate.ts"; +import { applyStreamOptionsPatch } from "../../hooks.ts"; +import { insertEntry, insertUsage } from "../../session/commit.ts"; +import { SessionInvariantError } from "../../session/session.ts"; +import type { + Deferred, + MessageEntry, + NewEntry, + OperationError, + RunState, + SettledAssistantMessage, + ToolCall, + UsageRow, +} from "../../session/types.ts"; +import { + branchTip, + deleteList, + operationState as operationStateValue, + pendingAssistantFrames, + setValue, +} from "../../session/values.ts"; +import type { AgentHarnessStreamOptions } from "../../types.ts"; +import type { Lane } from "../lane.ts"; +import { openFrameProgress } from "../progress.ts"; +import type { Drive, LaneCommand, LostOwnership, ProcedureResult } from "../types.ts"; + +type LocalResult = { kind: "committed" } | { kind: "state_changed" } | LostOwnership; +type SourceResult = { kind: "value"; handle: DeferredHandle } | { kind: "state_changed" } | LostOwnership; + +function sameDeferred(current: Deferred, expected: Deferred): boolean { + if ( + current.status !== expected.status || + current.stepId !== expected.stepId || + current.sourceEntryId !== expected.sourceEntryId || + current.poll !== expected.poll + ) { + return false; + } + return ( + current.status === "suspended" || + (expected.status === "effect_pending" && + current.responseEntryId === expected.responseEntryId && + current.usageId === expected.usageId) + ); +} + +function configurationError(identity: Deferred["configuration"]["model"]): OperationError { + return { + code: "model_unavailable", + message: "The configured model is unavailable in this process", + details: identity, + }; +} + +function providerError(message: SettledAssistantMessage): OperationError { + return { + code: "assistant_error", + message: message.errorMessage ?? `Deferred request ended with ${message.stopReason}`, + }; +} + +function normalizeError(message: SettledAssistantMessage, errorMessage: string): SettledAssistantMessage { + return { ...message, stopReason: "error", errorMessage }; +} + +function normalizeAborted(message: SettledAssistantMessage): SettledAssistantMessage { + return { + ...message, + stopReason: "aborted", + errorMessage: message.errorMessage ?? "Deferred request was cancelled", + }; +} + +function uuidV7Timestamp(id: string): number { + const timestamp = Number.parseInt(id.slice(0, 8) + id.slice(9, 13), 16); + if (!Number.isSafeInteger(timestamp)) throw new SessionInvariantError(`Invalid reserved UUIDv7 ${id}`); + return timestamp; +} + +function isUpdateEvent( + event: AssistantMessageEvent, +): event is Exclude { + return event.type !== "start" && event.type !== "done" && event.type !== "error"; +} + +async function readSourceHandle( + lane: Lane, + drive: Drive, + expected: Deferred, +): Promise { + return lane.commandDriveOwned( + drive, + async (state, reader): Promise> => { + const operation = state.operation; + if (operation === null || operation.state.kind !== "run") { + throw new SessionInvariantError("Deferred source reached a non-run operation"); + } + const run = operation.state; + if (run.control.status === "cancel_requested") { + return { kind: "return", result: { kind: "state_changed" } as const }; + } + if (run.phase.kind !== "deferred" || !sameDeferred(run.phase.deferred, expected)) { + throw new SessionInvariantError("Deferred source lost its durable boundary"); + } + const source = (await reader.getEntries([expected.sourceEntryId], drive.context)).get(expected.sourceEntryId); + if ( + source?.type !== "message" || + source.message.role !== "assistant" || + source.message.stopReason !== "deferred" || + source.message.deferred === undefined + ) { + throw new SessionInvariantError( + `Deferred source ${expected.sourceEntryId} is missing its assistant handle`, + ); + } + const handle = source.message.deferred; + const identity = expected.configuration.model; + if ( + handle.id.length === 0 || + handle.provider !== identity.provider || + handle.modelId !== identity.modelId || + handle.api !== source.message.api + ) { + throw new SessionInvariantError(`Deferred source ${expected.sourceEntryId} has an invalid handle`); + } + return { kind: "return", result: { kind: "value", handle } as const }; + }, + drive.context, + ); +} + +function enterDeferredConfigurationFailure( + lane: Lane, + drive: Drive, + expected: Deferred, +): Promise { + return lane.commandDriveOwned( + drive, + (state) => { + const operation = state.operation; + if (operation === null || operation.state.kind !== "run") { + throw new SessionInvariantError("Deferred configuration failure reached a non-run operation"); + } + const run = operation.state; + if (run.control.status === "cancel_requested") { + return { kind: "return", result: { kind: "continue" } as const }; + } + if (run.phase.kind !== "deferred" || !sameDeferred(run.phase.deferred, expected)) { + throw new SessionInvariantError("Deferred configuration failure lost its durable boundary"); + } + const nextState: RunState = { + ...run, + phase: { + kind: "failure_drain", + error: configurationError(expected.configuration.model), + provenance: { kind: "configuration" }, + }, + }; + return { + kind: "commit", + writes: [ + ...(expected.status === "effect_pending" + ? [deleteList(pendingAssistantFrames(drive.operationId, expected.responseEntryId))] + : []), + setValue(operationStateValue(drive.operationId), nextState), + ], + next: { ...state, operation: { meta: operation.meta, state: nextState } }, + materialize: () => ({ kind: "continue" }) as const, + }; + }, + drive.context, + ); +} + +async function commitPollIntent( + lane: Lane, + drive: Drive, + expected: Deferred, + next: Extract, + recovery: boolean, +): Promise { + return lane.commandDriveOwned( + drive, + (state) => { + const operation = state.operation; + if (operation === null || operation.state.kind !== "run") { + throw new SessionInvariantError("Deferred poll intent reached a non-run operation"); + } + const run = operation.state; + if (run.control.status === "cancel_requested") { + return { kind: "return", result: { kind: "state_changed" } as const }; + } + if (drive.deferredPermits === 0) { + throw new SessionInvariantError("Deferred poll intent lost its pass-local permit"); + } + if (run.phase.kind !== "deferred" || !sameDeferred(run.phase.deferred, expected)) { + throw new SessionInvariantError("Deferred poll intent lost its durable boundary"); + } + const nextState: RunState = { ...run, phase: { kind: "deferred", deferred: next } }; + return { + kind: "commit", + writes: [ + ...(expected.status === "effect_pending" + ? [deleteList(pendingAssistantFrames(drive.operationId, expected.responseEntryId))] + : []), + setValue(operationStateValue(drive.operationId), nextState), + ], + next: { ...state, operation: { meta: operation.meta, state: nextState } }, + materialize: () => { + drive.deferredPermits--; + return { kind: "committed" } as const; + }, + events: () => [ + { + type: "turn_start", + lane: lane.name, + runId: drive.operationId, + turnId: `${next.stepId}:poll:${next.poll}`, + ...(recovery ? { recovery: true as const } : {}), + }, + ], + }; + }, + drive.context, + ); +} + +async function streamDeferredResponse( + lane: Lane, + drive: Drive, + model: Model, + handle: DeferredHandle, + streamOptions: AgentHarnessStreamOptions, + responseEntryId: string, + recovery: boolean, +): Promise { + const progress = openFrameProgress(lane, drive, responseEntryId); + const frameEncoder = new AssistantMessageFrameEncoder(); + let metadata: { status?: number; headers?: Record } = {}; + let drained = false; + const drainProgress = async () => { + if (drained) return; + progress.seal(); + await progress.drain(); + drained = true; + }; + const admitted = withAbortSignal(drive.gate.signal, drive.context); + const stream = drive.gate.admit(() => + lane.models.streamDeferred(model, handle, { + wait: 0, + signal: admitted.abortSignal, + telemetryContext: admitted.telemetryContext, + timeoutMs: streamOptions.timeoutMs, + maxRetries: streamOptions.maxRetries, + maxRetryDelayMs: streamOptions.maxRetryDelayMs, + headers: streamOptions.headers, + onPayload: async (payload, requestModel) => { + const result = await lane.hooks.runWithGate( + "before_payload", + { lane: lane.name, runId: drive.operationId, model: requestModel, payload }, + drive.gate, + drive.context, + ); + return result?.payload; + }, + onResponse: (response) => { + metadata = { status: response.status, headers: response.headers }; + }, + }), + ); + + let started = false; + let response: SettledAssistantMessage; + try { + for await (const event of stream) { + if (event.type === "start") { + if (started) throw new Error("Assistant message stream emitted more than one start event"); + started = true; + const frame = frameEncoder.encode(event); + if (frame !== undefined) progress.write(frame); + await lane.emitBatch( + [ + { + type: "message_start", + lane: lane.name, + runId: drive.operationId, + message: { ...event.partial }, + ...(recovery ? { recovery: true as const } : {}), + }, + ], + drive.context, + ); + } else if (isUpdateEvent(event)) { + if (!started) throw new Error(`Assistant message stream emitted ${event.type} before start`); + const frame = frameEncoder.encode(event); + if (frame !== undefined) progress.write(frame); + await lane.emitBatch( + [ + { + type: "message_update", + lane: lane.name, + runId: drive.operationId, + message: { ...event.partial }, + event, + ...(frame === undefined ? {} : { frame }), + ...(recovery ? { recovery: true as const } : {}), + }, + ], + drive.context, + ); + } else if (event.type === "done" && !started) { + throw new Error("Assistant message stream emitted done before start"); + } + } + response = (await stream.result()) as SettledAssistantMessage; + } finally { + await drainProgress(); + } + + let final = response; + try { + const transformed = await lane.hooks.runWithGate( + "after_response", + { lane: lane.name, runId: drive.operationId, ...metadata, message: response }, + drive.gate, + drive.context, + ); + final = transformed?.message ?? response; + } catch (error) { + if (!(error instanceof AbortRequested)) throw error; + await error.cancellation; + } + await lane.emitBatch( + [ + { + type: "message_end", + lane: lane.name, + runId: drive.operationId, + message: final, + entryId: responseEntryId, + ...(recovery ? { recovery: true as const } : {}), + }, + ], + drive.context, + ); + return final; +} + +async function settleDeferredResponse( + lane: Lane, + drive: Drive, + expected: Extract, + response: SettledAssistantMessage, + recovery: boolean, +): Promise { + return lane.commandDriveOwned( + drive, + (state) => { + const operation = state.operation; + if (operation === null || operation.state.kind !== "run") { + throw new SessionInvariantError("Deferred settlement reached a non-run operation"); + } + const run = operation.state; + if (run.phase.kind !== "deferred" || !sameDeferred(run.phase.deferred, expected)) { + throw new SessionInvariantError("Deferred settlement lost its effect-pending boundary"); + } + + let committed = response; + let nextPhase: RunState["phase"]; + if (run.control.status === "cancel_requested") { + committed = normalizeAborted(response); + nextPhase = { + kind: "checkpoint", + continuation: { kind: "may_finish", includeFinalAssistant: true }, + triggerEntryId: expected.responseEntryId, + }; + } else if (response.stopReason === "aborted") { + throw new SessionInvariantError("Deferred response is aborted while durable control is running"); + } else if (response.stopReason === "deferred") { + nextPhase = { + kind: "deferred", + deferred: { + status: "suspended", + stepId: expected.stepId, + sourceEntryId: expected.responseEntryId, + poll: expected.poll, + configuration: expected.configuration, + streamOptions: expected.streamOptions, + }, + }; + } else if (response.stopReason === "error") { + nextPhase = { + kind: "failure_drain", + error: providerError(response), + provenance: { kind: "response", entryId: expected.responseEntryId }, + }; + } else { + const calls = response.content.flatMap((content, sourceIndex) => + content.type === "toolCall" ? [{ content, sourceIndex }] : [], + ); + if (calls.length !== 0) { + const timestamp = uuidV7Timestamp(expected.responseEntryId); + const planned: ToolCall[] = calls.map(({ sourceIndex }) => ({ + status: "planned", + sourceIndex, + resultEntryId: lane.session.idGenerator.next(timestamp), + })); + nextPhase = { + kind: "tools", + batch: { + assistantEntryId: expected.responseEntryId, + configuration: expected.configuration, + turnId: `${expected.stepId}:poll:${expected.poll}`, + calls: planned, + }, + }; + } else if (response.stopReason === "toolUse") { + committed = normalizeError(response, "Provider reported tool use without any tool calls"); + nextPhase = { + kind: "failure_drain", + error: providerError(committed), + provenance: { kind: "response", entryId: expected.responseEntryId }, + }; + } else { + nextPhase = { + kind: "checkpoint", + continuation: { kind: "may_finish", includeFinalAssistant: true }, + triggerEntryId: expected.responseEntryId, + }; + } + } + + const nextState: RunState = { + ...run, + phase: nextPhase, + latestAssistantEntryId: expected.responseEntryId, + }; + const responseEntry: NewEntry = { + id: expected.responseEntryId, + parentId: state.tipId, + type: "message", + message: committed, + }; + const usageRow: Omit = { + id: expected.usageId, + usage: committed.usage, + entryId: expected.responseEntryId, + adjustment: false, + }; + return { + kind: "commit", + writes: [ + insertEntry(responseEntry), + insertUsage(usageRow), + setValue(branchTip(lane.name), expected.responseEntryId), + deleteList(pendingAssistantFrames(drive.operationId, expected.responseEntryId)), + setValue(operationStateValue(drive.operationId), nextState), + ], + next: { + ...state, + tipId: expected.responseEntryId, + operation: { meta: operation.meta, state: nextState }, + }, + materialize: () => ({ kind: "continue" }) as const, + events: (commit) => { + const entry: MessageEntry = { + ...responseEntry, + seq: commit.seqs[0]!, + timestamp: commit.timestamp, + }; + const events: HarnessEvent[] = [ + { type: "entry_added", lane: lane.name, entry, ...(recovery ? { recovery: true } : {}) }, + { + type: "usage", + lane: lane.name, + row: { ...usageRow, seq: commit.seqs[1]! }, + totals: commit.stats.usage, + }, + ]; + if (nextPhase.kind !== "tools") { + events.push({ + type: "turn_end", + lane: lane.name, + runId: drive.operationId, + turnId: `${expected.stepId}:poll:${expected.poll}`, + message: committed, + toolResults: [], + ...(recovery ? { recovery: true } : {}), + }); + } + if (nextPhase.kind === "deferred" && committed.deferred !== undefined) { + events.push({ + type: "run_suspend", + lane: lane.name, + runId: drive.operationId, + reason: "deferred", + deferred: committed.deferred, + ...(recovery ? { recovery: true } : {}), + }); + } + return events; + }, + }; + }, + drive.context, + ); +} + +async function pollDeferred( + lane: Lane, + drive: Drive, + run: RunState, + expected: Deferred, + recovery: boolean, +): Promise { + if (run.phase.kind !== "deferred" || !sameDeferred(run.phase.deferred, expected)) { + throw new SessionInvariantError("Deferred procedure requires its current durable phase"); + } + if (run.control.status === "cancel_requested") return { kind: "continue" }; + const source = await readSourceHandle(lane, drive, expected); + if (source.kind === "lost_ownership") return source; + if (source.kind === "state_changed") return { kind: "continue" }; + if (drive.deferredPermits === 0) { + return { + kind: "waiting", + outcome: { + kind: "waiting", + operationId: drive.operationId, + reason: "deferred", + deferred: source.handle, + }, + }; + } + + const identity = expected.configuration.model; + const model = lane.models.getModel(identity.provider, identity.modelId); + if (model === undefined) return enterDeferredConfigurationFailure(lane, drive, expected); + const baseOptions: AgentHarnessStreamOptions = { ...expected.streamOptions, deferred: false }; + const beforeRequest = await lane.hooks.runWithGate( + "before_request", + { + lane: lane.name, + runId: drive.operationId, + model, + step: "deferred", + attempt: expected.status === "suspended" ? expected.poll + 1 : expected.poll, + streamOptions: baseOptions, + }, + drive.gate, + drive.context, + ); + const streamOptions: AgentHarnessStreamOptions = { + ...(beforeRequest?.streamOptions === undefined + ? baseOptions + : applyStreamOptionsPatch(baseOptions, beforeRequest.streamOptions)), + deferred: false, + }; + const poll = expected.status === "suspended" ? expected.poll + 1 : expected.poll; + const at = Date.now(); + const responseEntryId = lane.session.idGenerator.next(at); + const usageId = lane.session.idGenerator.next(at); + const effectPending: Extract = { + status: "effect_pending", + stepId: expected.stepId, + sourceEntryId: expected.sourceEntryId, + poll, + responseEntryId, + usageId, + configuration: expected.configuration, + streamOptions: expected.streamOptions, + }; + const intent = await commitPollIntent(lane, drive, expected, effectPending, recovery); + if (intent.kind === "lost_ownership") return intent; + if (intent.kind === "state_changed") return { kind: "continue" }; + + const response = await streamDeferredResponse( + lane, + drive, + model, + source.handle, + streamOptions, + responseEntryId, + recovery, + ); + return settleDeferredResponse(lane, drive, effectPending, response, recovery); +} + +/** Poll one durably suspended deferred response when this pass carries a permit. */ +export function runDeferredSuspended( + lane: Lane, + drive: Drive, + run: RunState, + deferred: Extract, +): Promise { + return pollDeferred(lane, drive, run, deferred, false); +} + +/** Replace one orphaned unknown-outcome poll under fresh ids when this pass carries a permit. */ +export function recoverDeferredPoll( + lane: Lane, + drive: Drive, + run: RunState, + deferred: Extract, +): Promise { + return pollDeferred(lane, drive, run, deferred, true); +} + +/** Advance or report the wait for one deferred run phase. */ +export function runDeferred( + lane: Lane, + drive: Drive, + run: RunState, + deferred: Deferred, +): Promise { + return deferred.status === "suspended" + ? runDeferredSuspended(lane, drive, run, deferred) + : recoverDeferredPoll(lane, drive, run, deferred); +} diff --git a/packages/agent/src/harness/runtime/drive/generation.ts b/packages/agent/src/harness/runtime/drive/generation.ts index 0fa541350..2ebe7294c 100644 --- a/packages/agent/src/harness/runtime/drive/generation.ts +++ b/packages/agent/src/harness/runtime/drive/generation.ts @@ -107,6 +107,41 @@ function saturatingAdd(left: number, right: number): number { return Number.isSafeInteger(sum) ? sum : Number.MAX_SAFE_INTEGER; } +function waitUntil(notBefore: number, signal: AbortSignal): Promise { + return new Promise((resolve, reject) => { + let timer: ReturnType | undefined; + const cleanup = () => { + if (timer !== undefined) clearTimeout(timer); + signal.removeEventListener("abort", onAbort); + }; + const onAbort = () => { + cleanup(); + reject(signal.reason); + }; + const check = () => { + const remaining = notBefore - Date.now(); + if (remaining <= 0) { + cleanup(); + resolve(); + return; + } + timer = setTimeout(check, Math.min(remaining, 2_147_483_647)); + }; + signal.addEventListener("abort", onAbort, { once: true }); + if (signal.aborted) onAbort(); + else check(); + }); +} + +function sameRetryWait(current: Generation, expected: Extract): boolean { + return ( + current.status === "retry_wait" && + current.context.stepId === expected.context.stepId && + current.nextAttempt === expected.nextAttempt && + current.notBefore === expected.notBefore + ); +} + function deferredHandleIsValid( message: SettledAssistantMessage, identity: GenerationContext["configuration"]["model"], @@ -428,6 +463,27 @@ export function settleAssistant( totals: commit.stats.usage, }, ]; + if ( + options.recovery !== true && + generation.attempt > 1 && + !(nextPhase.kind === "assistant" && nextPhase.generation.status === "retry_wait") + ) { + const success = committed.stopReason !== "error" && committed.stopReason !== "aborted"; + events.push({ + type: "retry_end", + lane: lane.name, + runId: drive.operationId, + step: identity.stepId, + attempt: generation.attempt, + success, + ...(success + ? {} + : { + finalError: + committed.errorMessage ?? `Assistant request ended with ${committed.stopReason}`, + }), + }); + } if ( options.recovery !== true && nextPhase.kind === "assistant" && @@ -459,6 +515,15 @@ export function settleAssistant( ...options, }); } + if (options.recovery !== true && nextPhase.kind === "deferred" && committed.deferred !== undefined) { + events.push({ + type: "run_suspend", + lane: lane.name, + runId: drive.operationId, + reason: "deferred", + deferred: committed.deferred, + }); + } return events; }, }; @@ -467,15 +532,87 @@ export function settleAssistant( ); } -/** Execute one ready assistant generation through intent, admitted stream, and settlement. */ +/** Advance one durable assistant retry wait according to this pass's local wait policy. */ +export async function runRetryWait( + lane: Lane, + drive: Drive, + run: RunState, + generation: Extract, +): Promise { + if (run.phase.kind !== "assistant" || !sameRetryWait(run.phase.generation, generation)) { + throw new SessionInvariantError("runRetryWait requires its current retry wait"); + } + if (run.control.status === "cancel_requested") return { kind: "continue" }; + if (Date.now() < generation.notBefore) { + if (!drive.waitForRetry) { + return { + kind: "waiting", + outcome: { + kind: "waiting", + operationId: drive.operationId, + reason: "retry", + notBefore: generation.notBefore, + }, + }; + } + await drive.gate.admit(() => waitUntil(generation.notBefore, drive.gate.signal)); + } + + return lane.commandDriveOwned( + drive, + (state) => { + const operation = state.operation; + if (operation === null || operation.state.kind !== "run") { + throw new SessionInvariantError("Retry wait reached a non-run operation"); + } + const current = operation.state; + if (current.control.status === "cancel_requested") { + return { kind: "return", result: { kind: "continue" } as const }; + } + if (current.phase.kind !== "assistant" || !sameRetryWait(current.phase.generation, generation)) { + throw new SessionInvariantError("Retry wait lost its durable boundary"); + } + const nextState: RunState = { + ...current, + phase: { + kind: "assistant", + generation: { + status: "ready", + context: generation.context, + nextAttempt: generation.nextAttempt, + }, + }, + }; + return { + kind: "commit", + writes: [setValue(operationStateValue(drive.operationId), nextState)], + next: { ...state, operation: { meta: operation.meta, state: nextState } }, + materialize: () => ({ kind: "continue" }) as const, + events: () => [ + { + type: "retry_start", + lane: lane.name, + runId: drive.operationId, + step: generation.context.stepId, + attempt: generation.nextAttempt, + }, + ], + }; + }, + drive.context, + ); +} + +/** Execute one ready assistant generation or advance its durable retry wait. */ export async function runGeneration( lane: Lane, drive: Drive, run: RunState, generation: Generation, ): Promise { + if (generation.status === "retry_wait") return runRetryWait(lane, drive, run, generation); if (generation.status !== "ready") { - throw new SessionInvariantError("M3 runGeneration requires a ready generation"); + throw new SessionInvariantError("runGeneration requires a ready generation or retry wait"); } if (run.phase.kind !== "assistant" || !sameReadyGeneration(run.phase.generation, generation)) { throw new SessionInvariantError("runGeneration requires its current ready phase"); diff --git a/packages/agent/test/harness/runtime/drive-retry-deferred.test.ts b/packages/agent/test/harness/runtime/drive-retry-deferred.test.ts new file mode 100644 index 000000000..ee3c58811 --- /dev/null +++ b/packages/agent/test/harness/runtime/drive-retry-deferred.test.ts @@ -0,0 +1,573 @@ +import { + type AssistantMessage, + type DeferredFetchOptions, + fauxAssistantMessage, + fauxProvider, + fauxToolCall, + type MutableModels, + type Provider, +} from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createModels } from "../../../../ai/src/models.ts"; +import type { HarnessEvent, WatchHandle } from "../../../src/harness/agent-harness.ts"; +import { DEFAULT_COMPACTION_SETTINGS } from "../../../src/harness/compaction/compaction.ts"; +import { BACKGROUND_CONTEXT } from "../../../src/harness/context.ts"; +import { HookRegistry } from "../../../src/harness/hooks.ts"; +import { runCheckpoint, startRun } from "../../../src/harness/runtime/drive/checkpoint.ts"; +import { recoverDeferredPoll, runDeferred, runDeferredSuspended } from "../../../src/harness/runtime/drive/deferred.ts"; +import { runGeneration } from "../../../src/harness/runtime/drive/generation.ts"; +import { Lane } from "../../../src/harness/runtime/lane.ts"; +import { restoreLane } from "../../../src/harness/runtime/restore.ts"; +import { type Config, Drive } from "../../../src/harness/runtime/types.ts"; +import { MemoryStorage } from "../../../src/harness/session/memory.ts"; +import { StorageBackedSession } from "../../../src/harness/session/session.ts"; +import { GatingStorage } from "../../../src/harness/session/testing/gating-storage.ts"; +import { InstrumentedStorage } from "../../../src/harness/session/testing/instrumented-storage.ts"; +import type { + Deferred, + Generation, + LaneConfiguration, + RunState, + Session, + Write, +} from "../../../src/harness/session/types.ts"; +import * as storedValues from "../../../src/harness/session/values.ts"; + +const sessions: Session[] = []; +const operationId = "01950000-0000-7000-8000-000000000001"; + +interface Fixture { + backend: MemoryStorage; + gating: GatingStorage | undefined; + storage: InstrumentedStorage; + session: Session; + lane: Lane; + drive: Drive; + models: MutableModels; + faux: ReturnType; + hooks: HookRegistry; + events: HarnessEvent[]; + config: Config; +} + +function unusedWatch(): WatchHandle { + throw new Error("watch is not used by retry/deferred tests"); +} + +async function createFixture( + options: { gated?: boolean; pendingFetches?: number; deferredSubmission?: boolean } = {}, +): Promise { + const backend = new MemoryStorage({ now: () => 100 }); + const gating = options.gated ? new GatingStorage(backend) : undefined; + const storage = new InstrumentedStorage(gating ?? backend); + const session = new StorageBackedSession( + { id: `retry-deferred-${sessions.length}`, createdAt: 1, storageVersion: 1 }, + storage, + ); + sessions.push(session); + const faux = fauxProvider({ + tokenSize: { min: 1, max: 1 }, + deferred: { pendingFetches: options.pendingFetches ?? 0 }, + }); + const model = faux.getModel(); + const configuration: LaneConfiguration = { + model: { provider: model.provider, modelId: model.id }, + thinkingLevel: "off", + activeToolNames: [], + }; + await session.mutate( + (mutator) => + mutator.commit( + [ + storedValues.setValue(storedValues.branchTip("main"), null), + storedValues.setValue(storedValues.laneConfig("main"), configuration), + storedValues.setValue(storedValues.laneState("main"), { + currentOperationId: null, + pendingNextRun: [], + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, + ); + const models = createModels(); + models.setProvider(faux.provider); + const hooks = new HookRegistry(() => {}); + const events: HarnessEvent[] = []; + const config: Config = { + tools: [], + resources: {}, + streamOptions: { deferred: options.deferredSubmission ?? true }, + retryPolicy: { enabled: true, maxRetries: 3, baseDelayMs: 10 }, + compaction: DEFAULT_COMPACTION_SETTINGS, + steeringMode: "all", + followUpMode: "all", + toolExecution: "parallel", + toolContext: undefined, + systemPrompt: undefined, + toProviderMessages: (messages) => + messages.filter( + (message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult", + ), + entryProjectors: {}, + }; + const lane = new Lane( + "main", + session, + models, + hooks, + await restoreLane(session, "main", BACKGROUND_CONTEXT), + (cause) => (cause instanceof Error ? cause : new Error(String(cause))), + (batch) => { + events.push(...structuredClone(batch)); + return Promise.resolve(); + }, + unusedWatch, + () => config, + ); + const admission = await lane.accept({ kind: "prompt", operationId, prompt: "question" }, BACKGROUND_CONTEXT); + if (!admission.ok) throw admission.error; + const drive = new Drive({ operationId }, BACKGROUND_CONTEXT); + lane.activeDrive = drive; + storage.clearCommitAttempts(); + return { backend, gating, storage, session, lane, drive, models, faux, hooks, events, config }; +} + +function currentRun(fixture: Fixture): RunState { + const operation = fixture.lane.state.operation; + if (operation?.state.kind !== "run") throw new Error("fixture has no run"); + return operation.state; +} + +async function advanceToReady(fixture: Fixture): Promise> { + let run = currentRun(fixture); + await startRun(fixture.lane, fixture.drive, run); + run = currentRun(fixture); + if (run.phase.kind !== "checkpoint") throw new Error("run did not reach checkpoint"); + await runCheckpoint(fixture.lane, fixture.drive, run, run.phase); + run = currentRun(fixture); + if (run.phase.kind !== "assistant" || run.phase.generation.status !== "ready") { + throw new Error("run did not reach ready generation"); + } + return run.phase.generation; +} + +async function replaceRunState(fixture: Fixture, nextState: RunState, extraWrites: Write[] = []): Promise { + await fixture.lane.command((state) => { + const operation = state.operation; + if (operation?.state.kind !== "run") throw new Error("fixture has no run"); + return { + kind: "commit", + writes: [...extraWrites, storedValues.setValue(storedValues.operationState(operationId), nextState)], + next: { ...state, operation: { meta: operation.meta, state: nextState } }, + materialize: () => undefined, + }; + }, BACKGROUND_CONTEXT); +} + +async function submitDeferred( + fixture: Fixture, + response: AssistantMessage = fauxAssistantMessage("done", { timestamp: 20 }), +): Promise> { + fixture.faux.setResponses([response]); + const ready = await advanceToReady(fixture); + await runGeneration(fixture.lane, fixture.drive, currentRun(fixture), ready); + const run = currentRun(fixture); + if (run.phase.kind !== "deferred" || run.phase.deferred.status !== "suspended") { + throw new Error("initial request did not suspend"); + } + return run.phase.deferred; +} + +async function installUnknownPoll( + fixture: Fixture, + suspended: Extract, +): Promise> { + const effectPending: Extract = { + status: "effect_pending", + stepId: suspended.stepId, + sourceEntryId: suspended.sourceEntryId, + poll: suspended.poll + 1, + responseEntryId: fixture.session.idGenerator.next(), + usageId: fixture.session.idGenerator.next(), + configuration: suspended.configuration, + streamOptions: suspended.streamOptions, + }; + const run = currentRun(fixture); + const nextState: RunState = { ...run, phase: { kind: "deferred", deferred: effectPending } }; + await replaceRunState(fixture, nextState, [ + storedValues.appendList(storedValues.pendingAssistantFrames(operationId, effectPending.responseEntryId), { + type: "text_delta", + contentIndex: 0, + delta: "old", + }), + ]); + return effectPending; +} + +function currentDeferred(fixture: Fixture): Deferred { + const run = currentRun(fixture); + if (run.phase.kind !== "deferred") throw new Error("fixture has no deferred phase"); + return run.phase.deferred; +} + +function installDrive(fixture: Fixture, options: { waitForRetry?: boolean; pollDeferred?: boolean }): Drive { + const drive = new Drive({ operationId, ...options }, BACKGROUND_CONTEXT); + fixture.lane.activeDrive = drive; + return drive; +} + +async function expectProjectionRestores(fixture: Fixture): Promise { + expect(fixture.lane.state).toEqual(await restoreLane(fixture.session, "main", BACKGROUND_CONTEXT)); +} + +afterEach(async () => { + vi.useRealTimers(); + for (const session of sessions.splice(0)) await session.close(BACKGROUND_CONTEXT); +}); + +describe("runtime assistant retry wait", () => { + it("classifies a live retryable provider error into durable retry wait", async () => { + const fixture = await createFixture({ deferredSubmission: false }); + const ready = await advanceToReady(fixture); + fixture.faux.setResponses([ + fauxAssistantMessage([], { stopReason: "error", errorMessage: "503 service unavailable", timestamp: 10 }), + ]); + + expect(await runGeneration(fixture.lane, fixture.drive, currentRun(fixture), ready)).toEqual({ + kind: "continue", + }); + expect(currentRun(fixture).phase).toMatchObject({ + kind: "assistant", + generation: { + status: "retry_wait", + nextAttempt: 2, + errorMessage: "503 service unavailable", + }, + }); + expect(fixture.events.at(-1)).toMatchObject({ type: "retry_scheduled", attempt: 2 }); + await expectProjectionRestores(fixture); + }); + + it("returns a durable waiting outcome without a timer or write when local waiting is disabled", async () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + const fixture = await createFixture({ deferredSubmission: false }); + const ready = await advanceToReady(fixture); + const retryWait: Extract = { + status: "retry_wait", + context: ready.context, + nextAttempt: 2, + notBefore: 1_100, + errorMessage: "retry", + }; + const run = currentRun(fixture); + const nextState: RunState = { ...run, phase: { kind: "assistant", generation: retryWait } }; + await replaceRunState(fixture, nextState); + fixture.storage.clearCommitAttempts(); + const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout"); + + expect(await runGeneration(fixture.lane, fixture.drive, currentRun(fixture), retryWait)).toEqual({ + kind: "waiting", + outcome: { kind: "waiting", operationId, reason: "retry", notBefore: 1_100 }, + }); + expect(setTimeoutSpy).not.toHaveBeenCalled(); + expect(fixture.storage.getCommitAttempts()).toEqual([]); + }); + + it("commits ready at the deadline and emits retry lifecycle around the next attempt", async () => { + vi.useFakeTimers(); + vi.setSystemTime(2_000); + const fixture = await createFixture({ deferredSubmission: false }); + const ready = await advanceToReady(fixture); + const retryWait: Extract = { + status: "retry_wait", + context: ready.context, + nextAttempt: 2, + notBefore: 2_000, + errorMessage: "retry", + }; + const run = currentRun(fixture); + await replaceRunState(fixture, { ...run, phase: { kind: "assistant", generation: retryWait } }); + fixture.storage.clearCommitAttempts(); + + expect(await runGeneration(fixture.lane, fixture.drive, currentRun(fixture), retryWait)).toEqual({ + kind: "continue", + }); + const next = currentRun(fixture); + if (next.phase.kind !== "assistant" || next.phase.generation.status !== "ready") { + throw new Error("retry did not become ready"); + } + expect(fixture.events.at(-1)).toMatchObject({ type: "retry_start", attempt: 2 }); + fixture.faux.setResponses([fauxAssistantMessage("retried", { timestamp: 30 })]); + await runGeneration(fixture.lane, fixture.drive, next, next.phase.generation); + expect(fixture.events).toContainEqual(expect.objectContaining({ type: "retry_end", attempt: 2, success: true })); + await expectProjectionRestores(fixture); + }); + + it("admits an abort-aware timer only for local waiting", async () => { + vi.useFakeTimers(); + vi.setSystemTime(3_000); + const fixture = await createFixture({ deferredSubmission: false }); + const ready = await advanceToReady(fixture); + const retryWait: Extract = { + status: "retry_wait", + context: ready.context, + nextAttempt: 2, + notBefore: 3_100, + errorMessage: "retry", + }; + const run = currentRun(fixture); + await replaceRunState(fixture, { ...run, phase: { kind: "assistant", generation: retryWait } }); + const drive = installDrive(fixture, { waitForRetry: true }); + fixture.storage.clearCommitAttempts(); + const waiting = runGeneration(fixture.lane, drive, currentRun(fixture), retryWait); + + await vi.advanceTimersByTimeAsync(99); + expect(fixture.storage.getCommitAttempts()).toEqual([]); + await vi.advanceTimersByTimeAsync(1); + expect(await waiting).toEqual({ kind: "continue" }); + expect(currentRun(fixture).phase).toMatchObject({ + kind: "assistant", + generation: { status: "ready", nextAttempt: 2 }, + }); + }); +}); + +describe("runtime deferred polling", () => { + it("waits without a permit, then performs at most one pending poll from captured options", async () => { + const fixture = await createFixture({ pendingFetches: 1 }); + const suspended = await submitDeferred(fixture); + fixture.storage.clearCommitAttempts(); + + expect(await runDeferredSuspended(fixture.lane, fixture.drive, currentRun(fixture), suspended)).toEqual({ + kind: "waiting", + outcome: { + kind: "waiting", + operationId, + reason: "deferred", + deferred: expect.objectContaining({ id: expect.any(String) }), + }, + }); + expect(fixture.faux.state.deferredFetchCount).toBe(0); + expect(fixture.storage.getCommitAttempts()).toEqual([]); + + let hookOptions: unknown; + fixture.hooks.on("before_request", (event) => { + if (event.step === "deferred") hookOptions = event.streamOptions; + return undefined; + }); + let fetchOptions: DeferredFetchOptions | undefined; + const provider = fixture.faux.provider; + const fetchDeferred = provider.fetchDeferred; + if (fetchDeferred === undefined) throw new Error("faux provider has no deferred fetch"); + const wrapped: Provider = { + ...provider, + fetchDeferred: (model, handle, options) => { + fetchOptions = options; + return fetchDeferred(model, handle, options); + }, + }; + fixture.models.setProvider(wrapped); + const driveOptions = { operationId, pollDeferred: true } as const; + const drive = new Drive(driveOptions, BACKGROUND_CONTEXT); + fixture.lane.activeDrive = drive; + + expect(await runDeferred(fixture.lane, drive, currentRun(fixture), currentDeferred(fixture))).toEqual({ + kind: "continue", + }); + expect(drive.deferredPermits).toBe(0); + expect(driveOptions).toEqual({ operationId, pollDeferred: true }); + expect(hookOptions).toMatchObject({ deferred: false }); + expect(fetchOptions).toMatchObject({ wait: 0 }); + expect(fetchOptions?.signal).toBe(drive.gate.signal); + expect(currentDeferred(fixture)).toMatchObject({ status: "suspended", poll: 1 }); + expect(fixture.faux.state.deferredFetchCount).toBe(1); + expect(fixture.events.slice(-2).map((event) => event.type)).toEqual(["turn_end", "run_suspend"]); + + expect(await runDeferred(fixture.lane, drive, currentRun(fixture), currentDeferred(fixture))).toMatchObject({ + kind: "waiting", + outcome: { reason: "deferred" }, + }); + expect(fixture.faux.state.deferredFetchCount).toBe(1); + await expectProjectionRestores(fixture); + }); + + it("plans ready deferred tool calls with the poll turn identity and follower result ids", async () => { + const fixture = await createFixture(); + const suspended = await submitDeferred( + fixture, + fauxAssistantMessage(fauxToolCall("lookup", { query: "value" }), { + stopReason: "toolUse", + timestamp: 20, + }), + ); + const drive = installDrive(fixture, { pollDeferred: true }); + + expect(await runDeferredSuspended(fixture.lane, drive, currentRun(fixture), suspended)).toEqual({ + kind: "continue", + }); + const run = currentRun(fixture); + if (run.phase.kind !== "tools") throw new Error("deferred tool response did not create a batch"); + expect(run.phase.batch).toMatchObject({ + turnId: `${suspended.stepId}:poll:1`, + calls: [{ status: "planned", sourceIndex: 0 }], + }); + expect(run.phase.batch.calls[0]?.resultEntryId.slice(0, 13)).toBe(run.latestAssistantEntryId?.slice(0, 13)); + expect(fixture.events.at(-1)).toMatchObject({ type: "usage" }); + await expectProjectionRestores(fixture); + }); + + it("consumes its permit only after the fresh intent commit lands", async () => { + const fixture = await createFixture({ gated: true }); + const suspended = await submitDeferred( + fixture, + fauxAssistantMessage([], { stopReason: "error", errorMessage: "failed", timestamp: 20 }), + ); + const gating = fixture.gating; + if (gating === undefined) throw new Error("fixture is not gated"); + const drive = installDrive(fixture, { pollDeferred: true }); + gating.arm(); + const polling = runDeferredSuspended(fixture.lane, drive, currentRun(fixture), suspended); + + await gating.waitPending(); + expect(drive.deferredPermits).toBe(1); + expect(fixture.faux.state.deferredFetchCount).toBe(0); + await gating.next(); + expect(drive.deferredPermits).toBe(0); + await expect.poll(() => fixture.faux.state.deferredFetchCount).toBe(1); + await gating.next(2); // start frame, then response settlement + expect(await polling).toEqual({ kind: "continue" }); + expect(currentRun(fixture).phase).toMatchObject({ kind: "failure_drain" }); + }); + + it("leaves suspended state unchanged when the fresh intent is discarded", async () => { + const fixture = await createFixture({ gated: true }); + const suspended = await submitDeferred(fixture); + const gating = fixture.gating; + if (gating === undefined) throw new Error("fixture is not gated"); + const drive = installDrive(fixture, { pollDeferred: true }); + gating.arm(); + const polling = runDeferredSuspended(fixture.lane, drive, currentRun(fixture), suspended); + void polling.catch(() => {}); + + await gating.waitPending(); + gating.discard(); + await expect(polling).rejects.toThrow("commit discarded"); + expect(drive.deferredPermits).toBe(1); + expect(fixture.faux.state.deferredFetchCount).toBe(0); + const durable = await fixture.backend.getValue(storedValues.operationState(operationId), BACKGROUND_CONTEXT); + expect(durable?.value).toMatchObject({ + kind: "run", + phase: { kind: "deferred", deferred: { status: "suspended", poll: 0 } }, + }); + }); + + it("replaces an unknown poll under fresh ids at the same poll number and deletes old frames", async () => { + const fixture = await createFixture(); + const suspended = await submitDeferred(fixture); + const unknown = await installUnknownPoll(fixture, suspended); + fixture.storage.clearCommitAttempts(); + + const noPermit = installDrive(fixture, {}); + expect(await recoverDeferredPoll(fixture.lane, noPermit, currentRun(fixture), unknown)).toMatchObject({ + kind: "waiting", + outcome: { reason: "deferred" }, + }); + expect( + await fixture.session.readList( + storedValues.pendingAssistantFrames(operationId, unknown.responseEntryId), + undefined, + BACKGROUND_CONTEXT, + ), + ).toHaveLength(1); + expect(fixture.faux.state.deferredFetchCount).toBe(0); + + const replacement = installDrive(fixture, { pollDeferred: true }); + expect(await recoverDeferredPoll(fixture.lane, replacement, currentRun(fixture), unknown)).toEqual({ + kind: "continue", + }); + const intent = fixture.storage + .getCommitAttempts() + .find((writes) => + writes.some( + (write) => + write.kind === "list" && + write.op === "delete" && + write.namespace === "pi.pending.assistant_frame" && + write.key.endsWith(unknown.responseEntryId), + ), + ); + expect(intent?.map((write) => `${write.kind}:${"op" in write ? write.op : "insert"}`)).toEqual([ + "list:delete", + "value:set", + ]); + const intentState = intent?.find( + (write) => write.kind === "value" && write.op === "set" && write.namespace === "pi.op.state", + ); + expect(intentState).toMatchObject({ + value: { + phase: { + kind: "deferred", + deferred: { + status: "effect_pending", + poll: unknown.poll, + responseEntryId: expect.not.stringMatching(unknown.responseEntryId), + usageId: expect.not.stringMatching(unknown.usageId), + }, + }, + }, + }); + expect(await fixture.session.getEntry(unknown.responseEntryId, BACKGROUND_CONTEXT)).toBeUndefined(); + expect((await fixture.storage.scanUsage({}, BACKGROUND_CONTEXT)).some((row) => row.id === unknown.usageId)).toBe( + false, + ); + expect( + await fixture.session.readList( + storedValues.pendingAssistantFrames(operationId, unknown.responseEntryId), + undefined, + BACKGROUND_CONTEXT, + ), + ).toEqual([]); + expect(currentRun(fixture).phase).toMatchObject({ kind: "checkpoint" }); + expect(fixture.events.some((event) => "recovery" in event && event.recovery === true)).toBe(true); + await expectProjectionRestores(fixture); + }); + + it("abandons unknown ids and frames into configuration failure when the captured model is unavailable", async () => { + const fixture = await createFixture(); + const suspended = await submitDeferred(fixture); + const unknown = await installUnknownPoll(fixture, suspended); + fixture.models.deleteProvider(fixture.faux.provider.id); + fixture.storage.clearCommitAttempts(); + const drive = installDrive(fixture, { pollDeferred: true }); + + expect(await recoverDeferredPoll(fixture.lane, drive, currentRun(fixture), unknown)).toEqual({ + kind: "continue", + }); + expect(drive.deferredPermits).toBe(1); + expect(fixture.faux.state.deferredFetchCount).toBe(0); + expect(currentRun(fixture).phase).toMatchObject({ + kind: "failure_drain", + error: { code: "model_unavailable" }, + provenance: { kind: "configuration" }, + }); + expect( + fixture.storage + .getCommitAttempts() + .at(-1) + ?.map((write) => + write.kind === "value" || write.kind === "list" + ? `${write.kind}:${write.op}:${write.namespace}` + : write.kind, + ), + ).toEqual(["list:delete:pi.pending.assistant_frame", "value:set:pi.op.state"]); + expect( + fixture.storage + .getCommitAttempts() + .flat() + .some((write) => write.kind === "entry" || write.kind === "usage"), + ).toBe(false); + await expectProjectionRestores(fixture); + }); +}); diff --git a/packages/ai/src/models.ts b/packages/ai/src/models.ts index dfc5202c0..ceeda8db3 100644 --- a/packages/ai/src/models.ts +++ b/packages/ai/src/models.ts @@ -214,6 +214,11 @@ export interface Models { streamSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): AssistantMessageEventStream; completeSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): Promise; + streamDeferred( + model: Model, + handle: DeferredHandle, + options?: ModelsDeferredFetchOptions, + ): AssistantMessageEventStream; fetchDeferred( model: Model, handle: DeferredHandle, @@ -703,11 +708,11 @@ class ModelsImpl implements MutableModels { return this.streamSimple(model, context, options).result(); } - async fetchDeferred( + streamDeferred( model: Model, handle: DeferredHandle, options?: ModelsDeferredFetchOptions, - ): Promise { + ): AssistantMessageEventStream { return lazyStream(model, async () => { const provider = this.requireProvider(model); if (!provider.fetchDeferred) { @@ -715,7 +720,15 @@ class ModelsImpl implements MutableModels { } const { requestModel, requestOptions } = await this.applyAuth(model, options); return provider.fetchDeferred(requestModel, handle, requestOptions as DeferredFetchOptions); - }).result(); + }); + } + + async fetchDeferred( + model: Model, + handle: DeferredHandle, + options?: ModelsDeferredFetchOptions, + ): Promise { + return this.streamDeferred(model, handle, options).result(); } async cancelDeferred( diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index 6f4071b60..a7bfec8b7 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -644,18 +644,26 @@ export class ModelRuntime implements Models { return this.streamSimple(model, context, options).result(); } - async fetchDeferred( + streamDeferred( model: Model, handle: DeferredHandle, options?: ModelsDeferredFetchOptions, - ): Promise { + ): AssistantMessageEventStream { return lazyStream(model, async () => { const prepared = await this.prepareRequest(model, options); if (!prepared.provider.fetchDeferred) { throw new ModelsError("provider", `Provider ${model.provider} does not support deferred responses`); } return prepared.provider.fetchDeferred(prepared.model, handle, prepared.options as DeferredFetchOptions); - }).result(); + }); + } + + async fetchDeferred( + model: Model, + handle: DeferredHandle, + options?: ModelsDeferredFetchOptions, + ): Promise { + return this.streamDeferred(model, handle, options).result(); } async cancelDeferred(