From 0e77e57d96e99b49f710c652d3f5985ac0c66f8f Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 26 Aug 2026 02:07:57 +0200 Subject: [PATCH] feat(agent): simplify durable runtime procedures --- packages/agent/docs/harness.md | 14 +- packages/agent/docs/runtime-simplification.md | 46 +- .../agent/src/harness/execution/assistant.ts | 71 +- .../src/harness/runtime/drive/checkpoint.ts | 149 ++-- .../src/harness/runtime/drive/deferred.ts | 293 +++---- .../src/harness/runtime/drive/generation.ts | 410 ++++------ .../src/harness/runtime/drive/recovery.ts | 64 +- .../src/harness/runtime/drive/response.ts | 105 ++- .../harness/runtime/drive/tool-placement.ts | 326 ++++++++ .../agent/src/harness/runtime/drive/tools.ts | 739 +++++------------- packages/agent/src/harness/runtime/harness.ts | 42 +- packages/agent/src/harness/runtime/lane.ts | 143 ++-- .../agent/src/harness/runtime/progress.ts | 37 +- packages/agent/src/harness/runtime/restore.ts | 129 ++- .../agent/src/harness/runtime/transcript.ts | 44 ++ packages/agent/src/harness/runtime/types.ts | 12 +- packages/agent/src/harness/session/types.ts | 11 - .../harness/runtime/drive-generation.test.ts | 48 ++ .../runtime/drive-retry-deferred.test.ts | 40 + .../agent/test/harness/runtime/lane.test.ts | 69 +- .../test/harness/runtime/progress.test.ts | 3 - .../test/harness/runtime/restore.test.ts | 8 +- 22 files changed, 1501 insertions(+), 1302 deletions(-) create mode 100644 packages/agent/src/harness/runtime/drive/tool-placement.ts create mode 100644 packages/agent/src/harness/runtime/transcript.ts diff --git a/packages/agent/docs/harness.md b/packages/agent/docs/harness.md index 955fc7be8..e4a1bbfe1 100644 --- a/packages/agent/docs/harness.md +++ b/packages/agent/docs/harness.md @@ -1856,18 +1856,20 @@ Parallel tool children are the exception: sibling call statuses genuinely race, The task runs direct async procedures. There is no graph interpreter or action scheduler. The Lane supplies two concrete mutation operations: -- `continueOperation` declines ordinary progress when current control is cancelled, pairs the next operation-state write with process-local projection publication, and otherwise exposes the procedure's writes/events; -- `settleOperation` performs effect settlement and tool-child transitions against current control, and also owns the universal terminal suffix when the procedure returns a terminal decision. +- `continueOperation` returns explicit `cancel_requested` without invoking the planner when current control is cancelled; otherwise it pairs the next operation-state write with process-local projection publication and returns the planner's result; +- `settleOperation` performs already-admitted effect settlement and tool-child transitions despite cancellation, and also owns the universal terminal suffix when the procedure returns a terminal decision. + +Intent publishers use `continueOperation`; outcome publishers use `settleOperation`. Cancellation therefore prevents a new durable intent but cannot erase the result of work already admitted. A pass ends at a terminal result or durable wait. `activeDrive` is then cleared by that pass; no live pass is replaced in-process. A process crash or close destroys/detaches the continuation, and a later attachment rebuilds `Lane.state` from durable values before another pass starts. Normal procedures remain straight line: ```text -prepare -→ commit effect intent -→ admit and await external effect -→ commit output + next state +prepare immutable inputs +→ publish durable effect intent +→ perform the external effect +→ publish one durable outcome ``` Recovery dispatches directly from the flat `state.at` leaf. Cancellation reconciliation runs before ordinary dispatch and never starts new ordinary effects. diff --git a/packages/agent/docs/runtime-simplification.md b/packages/agent/docs/runtime-simplification.md index 10a8716fc..fd24cee71 100644 --- a/packages/agent/docs/runtime-simplification.md +++ b/packages/agent/docs/runtime-simplification.md @@ -6,6 +6,37 @@ Rework the real implementation under `packages/agent/src/harness/runtime/` and t Public drive remains disabled until the execution graph is total. Format 4 is still work in progress, so the durable type replacement requires no migration or compatibility representation. +## Implementation status + +Completion measurements: + +- pre-simplification runtime at `eb1185d93`: 5,358 TypeScript lines; +- runtime after the first simplification at `417905647`: 4,667 lines; +- runtime after this pass: 4,654 lines; +- total reduction: 704 lines (13.1%); +- this pass versus `417905647`: -13 runtime lines. + +`npm run check` passes. The 13 focused assistant/runtime/type files pass with 128 tests. + +Implemented in the working tree: + +- explicit `ContinueOperationResult` with `cancel_requested` instead of implicit `undefined`; +- concrete generation, deferred-poll, and tool-call phases: prepare immutable input, publish durable intent, perform the effect, publish the durable outcome; +- shared assistant stream protocol consumption and Drive response lifecycle; +- `publishResponse` for common assistant/deferred durable publication and `publishConfigurationFailure` for their identical pre-intent failure transition; +- source-order tool publication, one-read immutable source validation, and shared tool-batch reconstruction isolated in `runtime/drive/tool-placement.ts`; +- same-call tool status defenses, unchecked `toolOperation()` casts, and their result unions removed; genuine sibling/progress/memo checks retained; +- centralized lane-storage classification without duplicate projection reads; +- checkpoint dead-check removal and exact checkpoint preservation; +- shared transcript mechanics in `runtime/transcript.ts`; +- narrower `LanePatch`, no redundant Drive promise-settlement flag, and no unused operation reject arm. + +Current largest files are `runtime/lane.ts` (~996), `runtime/drive/tools.ts` (~663), `runtime/drive/checkpoint.ts` (~452), and `runtime/drive/response.ts` (~417). + +Public drive remains disabled; M6–M8 have not started. Keep the visible durable procedure order: +`prepare → publish intent → perform effect → publish outcome`. +Do not introduce a generic Procedure interface, runner, scheduler, graph, callback plan, or dependency facade. Historical work-package documents remain unchanged. + ## Core model A Lane is one process-local actor over one durable lane projection. @@ -132,7 +163,7 @@ Used for ordinary non-terminal progress. - Enters the Session mutation line. - Receives the current authoritative Lane projection and operation state. -- If control is cancelled, returns to the drive loop without invoking the semantic planner. +- If control is cancelled, returns explicit `cancel_requested` without invoking the semantic planner; callers cannot confuse cancellation with a planner value. - The planner supplies procedure-specific writes, the next complete `OperationState`, materialization, and events. - The helper appends the `operationState` write and publishes the matching process-local operation projection. - It does not verify an expected state or `at` value. @@ -149,6 +180,19 @@ Used after an admitted provider/tool/structural effect and for genuine parallel The caller supplies the typed outcome, cleanup/publication writes, last result, and event. Terminal business decisions remain visible in the owning procedure. +## Durable procedure shape + +Effectful procedures expose four concrete phases: + +```text +prepare immutable inputs +→ publish durable effect intent +→ perform the external effect +→ publish one durable outcome +``` + +The intent phase must remain visible and precede the external effect. Otherwise a crash after the effect but before its intent would leave no recoverable unknown-outcome marker. Each procedure uses concrete functions such as `prepareGeneration`, `publishGenerationIntent`, `performGeneration`, and `publishResponse`; there is no generic Procedure abstraction. + ## Drive lifecycle Delete the installer-owned model. diff --git a/packages/agent/src/harness/execution/assistant.ts b/packages/agent/src/harness/execution/assistant.ts index 7faad5aac..62929cd3b 100644 --- a/packages/agent/src/harness/execution/assistant.ts +++ b/packages/agent/src/harness/execution/assistant.ts @@ -96,6 +96,42 @@ function isUpdateEvent( return event.type !== "start" && event.type !== "done" && event.type !== "error"; } +export async function consumeAssistantStream( + stream: AssistantMessageEventStream, + observer: AssistantStreamObserver, + afterResponse: + | ((message: SettledAssistantMessage, context: Context) => Promise) + | undefined, + context: Context, +): Promise { + let started = false; + 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; + await observer.start({ ...event.partial }, event, context); + } else if (isUpdateEvent(event)) { + if (!started) throw new Error(`Assistant message stream emitted ${event.type} before start`); + await observer.update({ ...event.partial }, event, context); + } else if (event.type === "done" && !started) { + throw new Error("Assistant message stream emitted done before start"); + } + } + + const settled = (await stream.result()) as SettledAssistantMessage; + let finalMessage = settled; + if (afterResponse !== undefined) { + try { + finalMessage = await afterResponse(settled, context); + } catch (error) { + if (!(error instanceof AbortRequested)) throw error; + await error.cancellation; + } + } + await observer.end(finalMessage, context); + return finalMessage; +} + /** Stream one assistant response without mutating the caller's message list. */ export async function streamHarnessAssistant( messages: AgentMessage[], @@ -127,30 +163,13 @@ export async function streamHarnessAssistant( context, ); - let started = false; - 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; - await config.observer.start({ ...event.partial }, event, context); - } else if (isUpdateEvent(event)) { - if (!started) throw new Error(`Assistant message stream emitted ${event.type} before start`); - await config.observer.update({ ...event.partial }, event, context); - } else if (event.type === "done" && !started) { - throw new Error("Assistant message stream emitted done before start"); - } - } - - const settled = (await stream.result()) as SettledAssistantMessage; - let finalMessage = settled; - if (config.afterResponse) { - try { - finalMessage = await config.afterResponse(settled, metadata, context); - } catch (error) { - if (!(error instanceof AbortRequested)) throw error; - await error.cancellation; - } - } - await config.observer.end(finalMessage, context); - return finalMessage; + const afterResponse = config.afterResponse; + return consumeAssistantStream( + stream, + config.observer, + afterResponse === undefined + ? undefined + : (message, afterContext) => afterResponse(message, metadata, afterContext), + context, + ); } diff --git a/packages/agent/src/harness/runtime/drive/checkpoint.ts b/packages/agent/src/harness/runtime/drive/checkpoint.ts index 3cd800463..82e8e216d 100644 --- a/packages/agent/src/harness/runtime/drive/checkpoint.ts +++ b/packages/agent/src/harness/runtime/drive/checkpoint.ts @@ -1,10 +1,9 @@ import type { AgentMessage } from "../../../types.ts"; -import type { HarnessEvent, TerminalOperationOutcome } from "../../agent-harness.ts"; +import type { TerminalOperationOutcome } from "../../agent-harness.ts"; import { insertEntry } from "../../session/commit.ts"; import { buildSessionContext } from "../../session/context.ts"; import { SessionInvariantError } from "../../session/session.ts"; import { - checkpointDataOf, type Entry, type LaneLastResult, type MessageEntry, @@ -17,28 +16,21 @@ import { } from "../../session/types.ts"; import { branchTip, deleteValue, pendingEntry, setValue } from "../../session/values.ts"; import type { Lane } from "../lane.ts"; -import type { Drive, ProcedureResult } from "../types.ts"; +import { chainEntries, entryLifecycleEvents, readPendingMessages } from "../transcript.ts"; +import type { ContinueOperationResult, Drive, ProcedureResult } from "../types.ts"; import { operationCleanupWrites } from "./terminal.ts"; +type FinishContinuation = Extract; + function isSettledAssistant(message: AgentMessage): message is SettledAssistantMessage { return message.role === "assistant" && message.stopReason !== "pending"; } -function lifecycleEvents(entry: Entry, lane: string, runId: string): HarnessEvent[] { - return entry.type === "message" - ? [ - { type: "message_start", lane, runId, message: entry.message }, - { type: "message_end", lane, runId, message: entry.message, entryId: entry.id }, - { type: "entry_added", lane, entry }, - ] - : [{ type: "entry_added", lane, entry }]; -} - async function readRunContext( lane: Lane, drive: Drive, checkpoint: RunCheckpointOperation, -): Promise { +): Promise> { const entries = await lane.continueOperation( checkpoint, async (state, _current, _meta, reader) => { @@ -53,9 +45,15 @@ async function readRunContext( }, drive.context, ); - return entries === undefined - ? undefined - : buildSessionContext(entries, { entryProjectors: lane.readConfig().entryProjectors }, drive.context); + if (entries.kind === "cancel_requested") return entries; + return { + kind: "result", + value: await buildSessionContext( + entries.value, + { entryProjectors: lane.readConfig().entryProjectors }, + drive.context, + ), + }; } /** Consume before_run and commit the initial checkpoint. */ @@ -80,11 +78,11 @@ export async function startRun( }, drive.context, ); - if (prompt === undefined) return { kind: "continue" }; + if (prompt.kind === "cancel_requested") return { kind: "continue" }; const hook = await lane.hooks.runWithGate( "before_run", - { lane: lane.name, runId: drive.operationId, prompt, resources: lane.readConfig().resources }, + { lane: lane.name, runId: drive.operationId, prompt: prompt.value, resources: lane.readConfig().resources }, drive.gate, drive.context, ); @@ -99,13 +97,11 @@ export async function startRun( const result = await lane.continueOperation( run, (state, current) => { - let parentId = state.tipId; - const entries = reserved.map(({ id, message }) => { - const entry: NewEntry = { id, parentId, type: "message", message }; - parentId = id; - return entry; - }); - const triggerEntryId = parentId; + const entries: NewEntry[] = chainEntries( + state.tipId, + reserved.map(({ id, message }) => ({ id, type: "message" as const, message })), + ); + const triggerEntryId = entries.at(-1)?.id ?? state.tipId; if (triggerEntryId === null) throw new SessionInvariantError("Run start has no trigger entry"); const nextState: RunCheckpointOperation = { ...runScopeOf(current), @@ -124,7 +120,7 @@ export async function startRun( materialize: () => ({ kind: "continue" }) as const, events: (commit) => entries.flatMap((entry, index) => - lifecycleEvents( + entryLifecycleEvents( { ...entry, seq: commit.seqs[index]!, timestamp: commit.timestamp }, lane.name, drive.operationId, @@ -134,7 +130,7 @@ export async function startRun( }, drive.context, ); - return result ?? { kind: "continue" }; + return result.kind === "cancel_requested" ? { kind: "continue" } : result.value; } async function applyPendingWrites( @@ -173,12 +169,13 @@ async function applyPendingWrites( } return entry; }); - const scope = { ...runScopeOf(current), inbox: { ...current.inbox, writes: [] } }; + const inbox = { ...current.inbox, writes: [] }; const nextState: RunCheckpointOperation = triggerEntryId === undefined - ? { ...scope, at: "run.checkpoint", ...checkpointDataOf(current) } + ? { ...current, inbox } : { - ...scope, + ...runScopeOf(current), + inbox, at: "run.checkpoint", continuation: { kind: "need_assistant", overflowRecoveryUsed: false }, triggerEntryId, @@ -196,7 +193,7 @@ async function applyPendingWrites( materialize: () => ({ kind: "continue" }) as const, events: (commit) => entries.flatMap((entry, index) => - lifecycleEvents( + entryLifecycleEvents( { ...entry, seq: commit.seqs[index]!, timestamp: commit.timestamp } as Entry, lane.name, drive.operationId, @@ -206,7 +203,7 @@ async function applyPendingWrites( }, drive.context, ); - return result ?? { kind: "continue" }; + return result.kind === "cancel_requested" ? { kind: "continue" } : result.value; } async function consumeQueuedMessages( @@ -222,35 +219,21 @@ async function consumeQueuedMessages( const mode = queue === "steer" ? current.settings.steeringMode : current.settings.followUpMode; const ids = mode === "all" ? allIds : allIds.slice(0, 1); if (ids.length === 0) return { kind: "return", result: { kind: "continue" } as const }; - const readMessages = (queuedIds: readonly string[], description: string) => - Promise.all( - queuedIds.map(async (id) => { - const value = await reader.getValue(pendingEntry(id), drive.context); - if (value?.value.type !== "message") { - throw new SessionInvariantError(`${description} ${id} is missing its message payload`); - } - return { entryId: id, message: value.value.payload }; - }), - ); - const messages = (await readMessages(ids, `Queued ${queue} entry`)).map(({ entryId, message }) => ({ - id: entryId, - message, - })); + const messages = (await readPendingMessages(reader, ids, `Queued ${queue} entry`, drive.context)).map( + ({ entryId, message }) => ({ id: entryId, message }), + ); const remainingSteerIds = queue === "steer" ? allIds.slice(ids.length) : current.inbox.steer; const remainingFollowUpIds = queue === "followUp" ? allIds.slice(ids.length) : current.inbox.followUp; const [remainingSteer, remainingFollowUp, nextRun] = await Promise.all([ - readMessages(remainingSteerIds, "Steer entry"), - readMessages(remainingFollowUpIds, "Follow-up entry"), - readMessages(state.pendingNextRun, "Pending next-run entry"), + readPendingMessages(reader, remainingSteerIds, "Steer entry", drive.context), + readPendingMessages(reader, remainingFollowUpIds, "Follow-up entry", drive.context), + readPendingMessages(reader, state.pendingNextRun, "Pending next-run entry", drive.context), ]); - let parentId = state.tipId; - const entries = messages.map(({ id, message }) => { - const entry: NewEntry = { id, parentId, type: "message", message }; - parentId = id; - return entry; - }); - const triggerEntryId = entries.at(-1)?.id; - if (triggerEntryId === undefined) throw new SessionInvariantError("Queue drain produced no trigger entry"); + const entries: NewEntry[] = chainEntries( + state.tipId, + messages.map(({ id, message }) => ({ id, type: "message" as const, message })), + ); + const triggerEntryId = messages[messages.length - 1]!.id; const nextState: RunCheckpointOperation = { ...runScopeOf(current), at: "run.checkpoint", @@ -271,7 +254,7 @@ async function consumeQueuedMessages( materialize: () => ({ kind: "continue" }) as const, events: (commit) => [ ...entries.flatMap((entry, index) => - lifecycleEvents( + entryLifecycleEvents( { ...entry, seq: commit.seqs[index]!, timestamp: commit.timestamp }, lane.name, drive.operationId, @@ -289,59 +272,62 @@ async function consumeQueuedMessages( }, drive.context, ); - return result ?? { kind: "continue" }; + return result.kind === "cancel_requested" ? { kind: "continue" } : result.value; } async function finishRun( lane: Lane, drive: Drive, checkpoint: RunCheckpointOperation, + continuation: FinishContinuation, ): Promise { const context = await readRunContext(lane, drive, checkpoint); - if (context === undefined) return { kind: "continue" }; + if (context.kind === "cancel_requested") return { kind: "continue" }; const hook = await lane.hooks.runWithGate( "before_run_end", - { lane: lane.name, runId: drive.operationId, messages: context }, + { lane: lane.name, runId: drive.operationId, messages: context.value }, drive.gate, drive.context, ); - const followUp = hook?.followUp; - const followUpId = followUp === undefined ? undefined : lane.session.idGenerator.next(); - const followUpMessage: AgentMessage | undefined = - followUp === undefined ? undefined : { role: "user", content: followUp, timestamp: Date.now() }; + const followUp = + hook?.followUp === undefined + ? undefined + : { + id: lane.session.idGenerator.next(), + message: { role: "user" as const, content: hook.followUp, timestamp: Date.now() }, + }; const result = await lane.continueOperation( checkpoint, async (state, current, _meta, reader) => { if ( - current.continuation.kind !== "may_finish" || current.inbox.steer.length !== 0 || current.inbox.followUp.length !== 0 || current.inbox.writes.length !== 0 ) { return { kind: "return", result: { kind: "continue" } as const }; } - if (followUpId !== undefined && followUpMessage !== undefined) { + if (followUp !== undefined) { const nextState: RunCheckpointOperation = { ...runScopeOf(current), at: "run.checkpoint", continuation: { kind: "need_assistant", overflowRecoveryUsed: false }, - triggerEntryId: followUpId, + triggerEntryId: followUp.id, }; const entry: NewEntry = { - id: followUpId, + id: followUp.id, parentId: state.tipId, type: "message", - message: followUpMessage, + message: followUp.message, }; return { kind: "commit", - writes: [insertEntry(entry), setValue(branchTip(lane.name), followUpId)], + writes: [insertEntry(entry), setValue(branchTip(lane.name), followUp.id)], operationState: nextState, - lane: { tipId: followUpId }, + lane: { tipId: followUp.id }, materialize: () => ({ kind: "continue" }) as const, events: (commit) => - lifecycleEvents( + entryLifecycleEvents( { ...entry, seq: commit.seqs[0]!, timestamp: commit.timestamp }, lane.name, drive.operationId, @@ -350,8 +336,8 @@ async function finishRun( } if (state.tipId === null) throw new SessionInvariantError("Completed run has no tip"); - const finalId = current.continuation.includeFinalAssistant ? current.latestAssistantEntryId : null; - if (current.continuation.includeFinalAssistant && finalId === null) { + const finalId = continuation.includeFinalAssistant ? current.latestAssistantEntryId : null; + if (continuation.includeFinalAssistant && finalId === null) { throw new SessionInvariantError("Completed run is missing its final assistant"); } let finalEntry: { id: string; message: SettledAssistantMessage } | undefined; @@ -403,7 +389,7 @@ async function finishRun( }, drive.context, ); - return result ?? { kind: "continue" }; + return result.kind === "cancel_requested" ? { kind: "continue" } : result.value; } /** Advance one durable run checkpoint by one visible transition. */ @@ -418,8 +404,9 @@ export async function runCheckpoint( if (run.skipInboxOnce !== true && run.inbox.steer.length !== 0) { return consumeQueuedMessages(lane, drive, run, "steer"); } - if (run.continuation.kind === "need_assistant") { - const overflowRecoveryUsed = run.continuation.overflowRecoveryUsed; + const continuation = run.continuation; + if (continuation.kind === "need_assistant") { + const overflowRecoveryUsed = continuation.overflowRecoveryUsed; const config = lane.readConfig(); const retryPolicy = config.retryPolicy.enabled ? { maxAttempts: config.retryPolicy.maxRetries + 1, baseDelayMs: config.retryPolicy.baseDelayMs } @@ -456,10 +443,10 @@ export async function runCheckpoint( }, drive.context, ); - return result ?? { kind: "continue" }; + return result.kind === "cancel_requested" ? { kind: "continue" } : result.value; } if (run.inbox.followUp.length !== 0) { return consumeQueuedMessages(lane, drive, run, "followUp"); } - return finishRun(lane, drive, run); + return finishRun(lane, drive, run, continuation); } diff --git a/packages/agent/src/harness/runtime/drive/deferred.ts b/packages/agent/src/harness/runtime/drive/deferred.ts index 6101ba5a7..060fdb737 100644 --- a/packages/agent/src/harness/runtime/drive/deferred.ts +++ b/packages/agent/src/harness/runtime/drive/deferred.ts @@ -1,18 +1,11 @@ -import { - type Api, - type AssistantMessageEvent, - AssistantMessageFrameEncoder, - type DeferredHandle, - type Model, -} from "@earendil-works/pi-ai"; +import type { Api, DeferredHandle, Model } from "@earendil-works/pi-ai"; import { withAbortSignal } from "../../context.ts"; -import { AbortRequested } from "../../execution/effect-gate.ts"; +import { consumeAssistantStream } from "../../execution/assistant.ts"; import { applyStreamOptionsPatch } from "../../hooks.ts"; import { SessionInvariantError } from "../../session/session.ts"; import { type LaneConfiguration, type OperationError, - type OperationState, type RunDeferredEffectPendingOperation, type RunDeferredSuspendedOperation, type RunScope, @@ -22,13 +15,27 @@ import { import { deleteList, pendingAssistantFrames } from "../../session/values.ts"; import type { AgentHarnessStreamOptions } from "../../types.ts"; import type { Lane } from "../lane.ts"; -import { openFrameProgress } from "../progress.ts"; -import type { Drive, ProcedureResult } from "../types.ts"; -import { settleResponse } from "./response.ts"; +import type { ContinueOperationResult, Drive, ProcedureResult } from "../types.ts"; +import { openAssistantResponse, publishConfigurationFailure, publishResponse } from "./response.ts"; type DeferredLeaf = RunDeferredSuspendedOperation | RunDeferredEffectPendingOperation; /** Deferred effect payload without the run-wide scope fields. */ type EffectPendingFields = Omit; + +type PreparedDeferredPoll = { + kind: "ready"; + source: DeferredHandle; + model: Model; + poll: number; + streamOptions: AgentHarnessStreamOptions; +}; + +type DeferredPreparation = + | PreparedDeferredPoll + | { kind: "cancel_requested" } + | { kind: "waiting"; source: DeferredHandle } + | { kind: "configuration_failure" }; + function configurationError(identity: LaneConfiguration["model"]): OperationError { return { code: "model_unavailable", @@ -37,17 +44,11 @@ function configurationError(identity: LaneConfiguration["model"]): OperationErro }; } -function isUpdateEvent( - event: AssistantMessageEvent, -): event is Exclude { - return event.type !== "start" && event.type !== "done" && event.type !== "error"; -} - async function readSourceHandle( lane: Lane, drive: Drive, deferred: DeferredLeaf, -): Promise { +): Promise> { return lane.continueOperation( deferred, async (_state, _current, _meta, reader) => { @@ -78,48 +79,68 @@ async function readSourceHandle( ); } -async function enterDeferredConfigurationFailure( +async function prepareDeferredPoll( lane: Lane, drive: Drive, - deferred: DeferredLeaf, -): Promise { - const result = await lane.continueOperation( - deferred, - (_state, current) => { - const nextState: OperationState = { - ...runScopeOf(current), - at: "run.failure_drain", - error: configurationError(deferred.configuration.model), - provenance: { kind: "configuration" }, - }; - return { - kind: "commit", - writes: - deferred.at === "run.deferred.effect_pending" - ? [deleteList(pendingAssistantFrames(drive.operationId, deferred.responseEntryId))] - : [], - operationState: nextState, - materialize: () => ({ kind: "continue" }) as const, - }; + expected: DeferredLeaf, +): Promise { + const source = await readSourceHandle(lane, drive, expected); + if (source.kind === "cancel_requested") return source; + if (drive.deferredPermits === 0) return { kind: "waiting", source: source.value }; + + const identity = expected.configuration.model; + const model = lane.models.getModel(identity.provider, identity.modelId); + if (model === undefined) return { kind: "configuration_failure" }; + const baseOptions: AgentHarnessStreamOptions = { ...expected.streamOptions, deferred: false }; + const poll = expected.at === "run.deferred.suspended" ? expected.poll + 1 : expected.poll; + const beforeRequest = await lane.hooks.runWithGate( + "before_request", + { + lane: lane.name, + runId: drive.operationId, + model, + step: "deferred", + attempt: poll, + streamOptions: baseOptions, }, + drive.gate, drive.context, ); - return result ?? { kind: "continue" }; + return { + kind: "ready", + source: source.value, + model, + poll, + streamOptions: { + ...(beforeRequest?.streamOptions === undefined + ? baseOptions + : applyStreamOptionsPatch(baseOptions, beforeRequest.streamOptions)), + deferred: false, + }, + }; } -async function commitPollIntent( +async function publishPollIntent( lane: Lane, drive: Drive, deferred: DeferredLeaf, - next: EffectPendingFields, + prepared: PreparedDeferredPoll, recovery: boolean, -): Promise { +): Promise> { + const at = Date.now(); + const next: EffectPendingFields = { + at: "run.deferred.effect_pending", + stepId: deferred.stepId, + sourceEntryId: deferred.sourceEntryId, + poll: prepared.poll, + responseEntryId: lane.session.idGenerator.next(at), + usageId: lane.session.idGenerator.next(at), + configuration: deferred.configuration, + streamOptions: deferred.streamOptions, + }; return lane.continueOperation( deferred, (_state, current) => { - if (drive.deferredPermits === 0) { - throw new SessionInvariantError("Deferred poll intent lost its pass-local permit"); - } const nextState: RunDeferredEffectPendingOperation = { ...runScopeOf(current), ...next }; return { kind: "commit", @@ -147,35 +168,25 @@ async function commitPollIntent( ); } -async function streamDeferredResponse( +async function performDeferredPoll( lane: Lane, drive: Drive, - model: Model, - handle: DeferredHandle, - streamOptions: AgentHarnessStreamOptions, - responseEntryId: string, + prepared: PreparedDeferredPoll, + intent: RunDeferredEffectPendingOperation, recovery: boolean, ): Promise { - const progress = openFrameProgress(lane, drive, responseEntryId); - const frameEncoder = new AssistantMessageFrameEncoder(); + const response = openAssistantResponse(lane, drive, intent.responseEntryId, recovery); 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, { + lane.models.streamDeferred(prepared.model, prepared.source, { wait: 0, signal: admitted.abortSignal, telemetryContext: admitted.telemetryContext, - timeoutMs: streamOptions.timeoutMs, - maxRetries: streamOptions.maxRetries, - maxRetryDelayMs: streamOptions.maxRetryDelayMs, - headers: streamOptions.headers, + timeoutMs: prepared.streamOptions.timeoutMs, + maxRetries: prepared.streamOptions.maxRetries, + maxRetryDelayMs: prepared.streamOptions.maxRetryDelayMs, + headers: prepared.streamOptions.headers, onPayload: async (payload, requestModel) => { const result = await lane.hooks.runWithGate( "before_payload", @@ -191,81 +202,16 @@ async function streamDeferredResponse( }), ); - 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, + return await consumeAssistantStream( + stream, + response.observer, + (message, context) => response.afterResponse(message, metadata, context), drive.context, ); - final = transformed?.message ?? response; - } catch (error) { - if (!(error instanceof AbortRequested)) throw error; - await error.cancellation; + } finally { + await response.close(); } - 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 pollDeferred( @@ -274,70 +220,35 @@ async function pollDeferred( expected: DeferredLeaf, recovery: boolean, ): Promise { - const source = await readSourceHandle(lane, drive, expected); - if (source === undefined) return { kind: "continue" }; - if (drive.deferredPermits === 0) { + const prepared = await prepareDeferredPoll(lane, drive, expected); + if (prepared.kind === "cancel_requested") return { kind: "continue" }; + if (prepared.kind === "waiting") { return { kind: "waiting", outcome: { kind: "waiting", operationId: drive.operationId, reason: "deferred", - deferred: source, + deferred: prepared.source, }, }; } + if (prepared.kind === "configuration_failure") { + return publishConfigurationFailure( + lane, + drive, + expected, + configurationError(expected.configuration.model), + expected.at === "run.deferred.effect_pending" + ? [deleteList(pendingAssistantFrames(drive.operationId, expected.responseEntryId))] + : [], + ); + } - 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.at === "run.deferred.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.at === "run.deferred.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: EffectPendingFields = { - at: "run.deferred.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 === undefined) return { kind: "continue" }; - - const response = await streamDeferredResponse( - lane, - drive, - model, - source, - streamOptions, - intent.responseEntryId, - recovery, - ); - return settleResponse(lane, drive, intent, response, recovery ? { recovery: true } : {}); + const intent = await publishPollIntent(lane, drive, expected, prepared, recovery); + if (intent.kind === "cancel_requested") return { kind: "continue" }; + const response = await performDeferredPoll(lane, drive, prepared, intent.value, recovery); + return publishResponse(lane, drive, intent.value, response, recovery ? { recovery: true } : {}); } /** Poll one durably suspended deferred response when this pass carries a permit. */ diff --git a/packages/agent/src/harness/runtime/drive/generation.ts b/packages/agent/src/harness/runtime/drive/generation.ts index 325fa4cbf..beec25b2b 100644 --- a/packages/agent/src/harness/runtime/drive/generation.ts +++ b/packages/agent/src/harness/runtime/drive/generation.ts @@ -1,19 +1,13 @@ -import { - type AssistantMessage, - type AssistantMessageEvent, - AssistantMessageFrameEncoder, - type Tool, -} from "@earendil-works/pi-ai"; +import type { Api, Model, Tool } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../../types.ts"; import { type Context, withAbortSignal } from "../../context.ts"; -import { streamHarnessAssistant } from "../../execution/assistant.ts"; +import { type HarnessAssistantStreamConfig, streamHarnessAssistant } from "../../execution/assistant.ts"; import { applyStreamOptionsPatch } from "../../hooks.ts"; import { buildSessionContext } from "../../session/context.ts"; import { SessionInvariantError } from "../../session/session.ts"; import { type JsonValue, type OperationError, - type OperationState, type RunAssistantEffectPendingOperation, type RunAssistantReadyOperation, type RunAssistantRetryWaitOperation, @@ -23,22 +17,27 @@ import { } from "../../session/types.ts"; import type { AgentHarnessStreamOptions } from "../../types.ts"; import type { Lane } from "../lane.ts"; -import { openFrameProgress } from "../progress.ts"; -import type { Drive, ProcedureResult } from "../types.ts"; -import { settleResponse } from "./response.ts"; +import type { ContinueOperationResult, Drive, ProcedureResult } from "../types.ts"; +import { openAssistantResponse, publishConfigurationFailure, publishResponse } from "./response.ts"; /** Assistant effect-pending payload without the run-wide scope fields. */ type AssistantEffectPending = Omit; -const ZERO_USAGE = { - input: 0, - output: 0, - cacheRead: 0, - cacheWrite: 0, - totalTokens: 0, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +type PreparedGeneration = { + kind: "ready"; + model: Model; + tools: Tool[]; + messages: AgentMessage[]; + systemPrompt: string; + streamOptions: AgentHarnessStreamOptions; + toProviderMessages: HarnessAssistantStreamConfig["toProviderMessages"]; }; +type GenerationPreparation = + | PreparedGeneration + | { kind: "configuration_failure"; error: OperationError } + | { kind: "cancel_requested" }; + function configurationError( code: "model_unavailable" | "configured_tools_unavailable", details: JsonValue, @@ -95,7 +94,7 @@ async function readRequestMessages( lane: Lane, drive: Drive, generation: RunAssistantReadyOperation, -): Promise { +): Promise> { const path = await lane.continueOperation( generation, async (state, _current, _meta, reader) => { @@ -110,17 +109,97 @@ async function readRequestMessages( }, drive.context, ); - return path === undefined - ? undefined - : buildSessionContext(path, { entryProjectors: lane.readConfig().entryProjectors }, drive.context); + if (path.kind === "cancel_requested") return path; + return { + kind: "result", + value: await buildSessionContext( + path.value, + { entryProjectors: lane.readConfig().entryProjectors }, + drive.context, + ), + }; } -async function commitGenerationIntent( +async function prepareGeneration( + lane: Lane, + drive: Drive, + generation: RunAssistantReadyOperation, +): Promise { + const identity = generation.generationContext.configuration.model; + const model = lane.models.getModel(identity.provider, identity.modelId); + if (model === undefined) { + return { kind: "configuration_failure", error: configurationError("model_unavailable", identity) }; + } + + const config = lane.readConfig(); + const toolsByName = new Map(config.tools.map((tool) => [tool.name, tool])); + const missingTools = generation.generationContext.configuration.activeToolNames.filter( + (name) => !toolsByName.has(name), + ); + if (missingTools.length !== 0) { + return { + kind: "configuration_failure", + error: configurationError("configured_tools_unavailable", { tools: missingTools }), + }; + } + const tools: Tool[] = generation.generationContext.configuration.activeToolNames.map((name) => { + const tool = toolsByName.get(name); + if (tool === undefined) throw new SessionInvariantError(`Configured tool ${name} disappeared during resolution`); + return { + name: tool.name, + description: tool.description, + parameters: tool.parameters, + ...(tool.constrainedSampling === undefined ? {} : { constrainedSampling: tool.constrainedSampling }), + }; + }); + + const messages = await readRequestMessages(lane, drive, generation); + if (messages.kind === "cancel_requested") return messages; + const systemPrompt = await resolveSystemPrompt(lane, drive.context); + const beforeRequest = await lane.hooks.runWithGate( + "before_request", + { + lane: lane.name, + runId: drive.operationId, + model, + step: "assistant", + attempt: generation.nextAttempt, + streamOptions: generation.generationContext.streamOptions, + }, + drive.gate, + drive.context, + ); + const streamOptions = + beforeRequest?.streamOptions === undefined + ? generation.generationContext.streamOptions + : applyStreamOptionsPatch(generation.generationContext.streamOptions, beforeRequest.streamOptions); + return { + kind: "ready", + model, + tools, + messages: messages.value, + systemPrompt, + streamOptions, + toProviderMessages: config.toProviderMessages, + }; +} + +async function publishGenerationIntent( lane: Lane, drive: Drive, ready: RunAssistantReadyOperation, - pending: AssistantEffectPending, -): Promise { + prepared: PreparedGeneration, +): Promise> { + const at = Date.now(); + const pending: AssistantEffectPending = { + at: "run.assistant.effect_pending", + generationContext: ready.generationContext, + attempt: ready.nextAttempt, + responseEntryId: lane.session.idGenerator.next(at), + usageId: lane.session.idGenerator.next(at), + intendedOutputLimit: prepared.model.maxTokens, + contextWindow: prepared.model.contextWindow, + }; return lane.continueOperation( ready, (_state, current) => { @@ -147,32 +226,62 @@ async function commitGenerationIntent( ); } -/** Commit a non-retryable request-configuration failure before reserving response ids. */ -export async function enterConfigurationFailure( +async function performGeneration( lane: Lane, drive: Drive, - ready: RunAssistantReadyOperation, - error: OperationError, -): Promise { - const result = await lane.continueOperation( - ready, - (_state, current) => { - const nextState: OperationState = { - ...runScopeOf(current), - at: "run.failure_drain", - error, - provenance: { kind: "configuration" }, - }; - return { - kind: "commit", - writes: [], - operationState: nextState, - materialize: () => ({ kind: "continue" }) as const, - }; - }, - drive.context, - ); - return result ?? { kind: "continue" }; + intent: RunAssistantEffectPendingOperation, + prepared: PreparedGeneration, +): Promise { + const response = openAssistantResponse(lane, drive, intent.responseEntryId); + try { + return await streamHarnessAssistant( + prepared.messages, + { + model: prepared.model, + systemPrompt: prepared.systemPrompt, + tools: prepared.tools, + thinkingLevel: intent.generationContext.configuration.thinkingLevel, + streamOptions: prepared.streamOptions, + transformContext: async (requestContext, context) => { + const result = await lane.hooks.runWithGate( + "transform_context", + { lane: lane.name, runId: drive.operationId, ...requestContext }, + drive.gate, + context, + ); + return { + messages: result?.messages ?? requestContext.messages, + systemPrompt: result?.systemPrompt ?? requestContext.systemPrompt, + }; + }, + toProviderMessages: prepared.toProviderMessages, + beforePayload: async (payload, requestModel, context) => { + const result = await lane.hooks.runWithGate( + "before_payload", + { lane: lane.name, runId: drive.operationId, model: requestModel, payload }, + drive.gate, + context, + ); + return result?.payload; + }, + afterResponse: response.afterResponse, + request: (aiContext, options, context) => { + const admitted = withAbortSignal(drive.gate.signal, context); + return drive.gate.admit(() => + lane.models.streamSimple(prepared.model, aiContext, { + ...options, + signal: admitted.abortSignal, + telemetryContext: admitted.telemetryContext, + }), + ); + }, + observer: response.observer, + }, + drive.context, + ); + } finally { + await response.close(); + } } /** Advance one durable assistant retry wait according to this pass's local wait policy. */ @@ -223,7 +332,7 @@ export async function runRetryWait( }, drive.context, ); - return result ?? { kind: "continue" }; + return result.kind === "cancel_requested" ? { kind: "continue" } : result.value; } /** Execute one ready assistant generation or advance its durable retry wait. */ @@ -234,199 +343,14 @@ export async function runGeneration( ): Promise { if (generation.at === "run.assistant.retry_wait") return runRetryWait(lane, drive, generation); - const identity = generation.generationContext.configuration.model; - const model = lane.models.getModel(identity.provider, identity.modelId); - if (model === undefined) { - return enterConfigurationFailure(lane, drive, generation, configurationError("model_unavailable", identity)); + const prepared = await prepareGeneration(lane, drive, generation); + if (prepared.kind === "configuration_failure") { + return publishConfigurationFailure(lane, drive, generation, prepared.error); } - const config = lane.readConfig(); - const toolsByName = new Map(config.tools.map((tool) => [tool.name, tool])); - const missingTools = generation.generationContext.configuration.activeToolNames.filter( - (name) => !toolsByName.has(name), - ); - if (missingTools.length !== 0) { - return enterConfigurationFailure( - lane, - drive, - generation, - configurationError("configured_tools_unavailable", { tools: missingTools }), - ); - } - const tools: Tool[] = generation.generationContext.configuration.activeToolNames.map((name) => { - const tool = toolsByName.get(name); - if (tool === undefined) throw new SessionInvariantError(`Configured tool ${name} disappeared during resolution`); - return { - name: tool.name, - description: tool.description, - parameters: tool.parameters, - ...(tool.constrainedSampling === undefined ? {} : { constrainedSampling: tool.constrainedSampling }), - }; - }); + if (prepared.kind === "cancel_requested") return { kind: "continue" }; - const messages = await readRequestMessages(lane, drive, generation); - if (messages === undefined) return { kind: "continue" }; - const systemPrompt = await resolveSystemPrompt(lane, drive.context); - const beforeRequest = await lane.hooks.runWithGate( - "before_request", - { - lane: lane.name, - runId: drive.operationId, - model, - step: "assistant", - attempt: generation.nextAttempt, - streamOptions: generation.generationContext.streamOptions, - }, - drive.gate, - drive.context, - ); - const streamOptions: AgentHarnessStreamOptions = - beforeRequest?.streamOptions === undefined - ? generation.generationContext.streamOptions - : applyStreamOptionsPatch(generation.generationContext.streamOptions, beforeRequest.streamOptions); - const at = Date.now(); - const responseEntryId = lane.session.idGenerator.next(at); - const usageId = lane.session.idGenerator.next(at); - const effectPending: AssistantEffectPending = { - at: "run.assistant.effect_pending", - generationContext: generation.generationContext, - attempt: generation.nextAttempt, - responseEntryId, - usageId, - intendedOutputLimit: model.maxTokens, - contextWindow: model.contextWindow, - }; - const intent = await commitGenerationIntent(lane, drive, generation, effectPending); - if (intent === undefined) return { kind: "continue" }; - - const progress = openFrameProgress(lane, drive, intent.responseEntryId); - const frameEncoder = new AssistantMessageFrameEncoder(); - let drained = false; - const drainProgress = async () => { - if (drained) return; - progress.seal(); - await progress.drain(); - drained = true; - }; - let response: SettledAssistantMessage; - try { - response = await streamHarnessAssistant( - messages, - { - model, - systemPrompt, - tools, - thinkingLevel: generation.generationContext.configuration.thinkingLevel, - streamOptions, - transformContext: async (requestContext, context) => { - const result = await lane.hooks.runWithGate( - "transform_context", - { lane: lane.name, runId: drive.operationId, ...requestContext }, - drive.gate, - context, - ); - return { - messages: result?.messages ?? requestContext.messages, - systemPrompt: result?.systemPrompt ?? requestContext.systemPrompt, - }; - }, - toProviderMessages: config.toProviderMessages, - beforePayload: async (payload, requestModel, context) => { - const result = await lane.hooks.runWithGate( - "before_payload", - { lane: lane.name, runId: drive.operationId, model: requestModel, payload }, - drive.gate, - context, - ); - return result?.payload; - }, - afterResponse: async (message, metadata, context) => { - await drainProgress(); - const result = await lane.hooks.runWithGate( - "after_response", - { lane: lane.name, runId: drive.operationId, ...metadata, message }, - drive.gate, - context, - ); - return result?.message ?? message; - }, - request: (aiContext, options, context) => { - const admitted = withAbortSignal(drive.gate.signal, context); - return drive.gate.admit(() => - lane.models.streamSimple(model, aiContext, { - ...options, - signal: admitted.abortSignal, - telemetryContext: admitted.telemetryContext, - }), - ); - }, - observer: { - start(message, event, context) { - const frame = frameEncoder.encode(event); - if (frame !== undefined) progress.write(frame); - return lane.emitBatch( - [{ type: "message_start", lane: lane.name, runId: drive.operationId, message }], - context, - ); - }, - update(message, event: AssistantMessageEvent, context) { - const frame = frameEncoder.encode(event); - if (frame !== undefined) progress.write(frame); - return lane.emitBatch( - [ - { - type: "message_update", - lane: lane.name, - runId: drive.operationId, - message, - event, - ...(frame === undefined ? {} : { frame }), - }, - ], - context, - ); - }, - end(message, context) { - return lane.emitBatch( - [ - { - type: "message_end", - lane: lane.name, - runId: drive.operationId, - message, - entryId: responseEntryId, - }, - ], - context, - ); - }, - }, - }, - drive.context, - ); - } finally { - await drainProgress(); - } - - return settleResponse(lane, drive, intent, response); -} - -export function interruptedAssistantMessage( - generation: RunAssistantEffectPendingOperation, - partial: AssistantMessage | undefined, -): SettledAssistantMessage { - const warning = - "Assistant request was interrupted. The preceding content is the latest committed partial; newer live output may be missing and the external outcome is unknown."; - return partial === undefined - ? { - role: "assistant", - content: [], - api: "unknown", - provider: generation.generationContext.configuration.model.provider, - model: generation.generationContext.configuration.model.modelId, - usage: ZERO_USAGE, - stopReason: "error", - errorMessage: warning, - timestamp: Date.now(), - } - : { ...partial, usage: ZERO_USAGE, stopReason: "error", errorMessage: warning }; + const intent = await publishGenerationIntent(lane, drive, generation, prepared); + if (intent.kind === "cancel_requested") return { kind: "continue" }; + const response = await performGeneration(lane, drive, intent.value, prepared); + return publishResponse(lane, drive, intent.value, response); } diff --git a/packages/agent/src/harness/runtime/drive/recovery.ts b/packages/agent/src/harness/runtime/drive/recovery.ts index ce0f38eae..8afa40f96 100644 --- a/packages/agent/src/harness/runtime/drive/recovery.ts +++ b/packages/agent/src/harness/runtime/drive/recovery.ts @@ -1,10 +1,39 @@ -import { type AssistantMessageFrame, reduceAssistantMessageFrames } from "@earendil-works/pi-ai"; -import type { RunAssistantEffectPendingOperation } from "../../session/types.ts"; -import { pendingAssistantFrames } from "../../session/values.ts"; +import { type AssistantMessage, reduceAssistantMessageFrames } from "@earendil-works/pi-ai"; +import type { RunAssistantEffectPendingOperation, SettledAssistantMessage } from "../../session/types.ts"; import type { Lane } from "../lane.ts"; +import { readAssistantFrames } from "../progress.ts"; import type { Drive, ProcedureResult } from "../types.ts"; -import { interruptedAssistantMessage } from "./generation.ts"; -import { settleResponse } from "./response.ts"; +import { publishResponse } from "./response.ts"; + +const ZERO_USAGE = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +function interruptedAssistantMessage( + generation: RunAssistantEffectPendingOperation, + partial: AssistantMessage | undefined, +): SettledAssistantMessage { + const warning = + "Assistant request was interrupted. The preceding content is the latest committed partial; newer live output may be missing and the external outcome is unknown."; + return partial === undefined + ? { + role: "assistant", + content: [], + api: "unknown", + provider: generation.generationContext.configuration.model.provider, + model: generation.generationContext.configuration.model.modelId, + usage: ZERO_USAGE, + stopReason: "error", + errorMessage: warning, + timestamp: Date.now(), + } + : { ...partial, usage: ZERO_USAGE, stopReason: "error", errorMessage: warning }; +} /** Settle an orphaned assistant request from its bounded committed frame prefix without another provider call. */ export async function recoverAssistantGeneration( @@ -14,26 +43,15 @@ export async function recoverAssistantGeneration { const frames = await lane.continueOperation( generation, - async (_state, _current, _meta, reader) => { - const result: AssistantMessageFrame[] = []; - let cursor: { seq: number } | undefined; - for (;;) { - const page = await reader.readList( - pendingAssistantFrames(drive.operationId, generation.responseEntryId), - { order: "asc", limit: 1_000, ...(cursor === undefined ? {} : { cursor }) }, - drive.context, - ); - result.push(...page.map(({ value }) => value)); - if (page.length < 1_000) break; - cursor = { seq: page[page.length - 1]!.seq }; - } - return { kind: "return", result }; - }, + async (_state, _current, _meta, reader) => ({ + kind: "return", + result: await readAssistantFrames(reader, drive.operationId, generation.responseEntryId, drive.context), + }), drive.context, ); - if (frames === undefined) return { kind: "continue" }; + if (frames.kind === "cancel_requested") return { kind: "continue" }; - const message = interruptedAssistantMessage(generation, reduceAssistantMessageFrames(frames)); + const message = interruptedAssistantMessage(generation, reduceAssistantMessageFrames(frames.value)); await lane.emitBatch( [ { @@ -54,5 +72,5 @@ export async function recoverAssistantGeneration; + close(): Promise; +}; + +export function openAssistantResponse( + lane: Lane, + drive: Drive, + responseEntryId: string, + recovery = false, +): AssistantResponseLifecycle { + const progress = openFrameProgress(lane, drive, responseEntryId); + const frameEncoder = new AssistantMessageFrameEncoder(); + const eventContext = { + lane: lane.name, + runId: drive.operationId, + ...(recovery ? { recovery: true as const } : {}), + }; + const close = async () => { + progress.seal(); + await progress.drain(); + }; + return { + observer: { + start(message, event, context) { + const frame = frameEncoder.encode(event); + if (frame !== undefined) progress.write(frame); + return lane.emitBatch([{ type: "message_start", ...eventContext, message }], context); + }, + update(message, event: AssistantMessageEvent, context) { + const frame = frameEncoder.encode(event); + if (frame !== undefined) progress.write(frame); + return lane.emitBatch( + [{ type: "message_update", ...eventContext, message, event, ...(frame === undefined ? {} : { frame }) }], + context, + ); + }, + end(message, context) { + return lane.emitBatch( + [{ type: "message_end", ...eventContext, message, entryId: responseEntryId }], + context, + ); + }, + }, + async afterResponse(message, metadata, context) { + await close(); + const result = await lane.hooks.runWithGate( + "after_response", + { lane: lane.name, runId: drive.operationId, ...metadata, message }, + drive.gate, + context, + ); + return result?.message ?? message; + }, + close, + }; +} + type ResponseIntent = RunAssistantEffectPendingOperation | RunDeferredEffectPendingOperation; +/** Publish a non-retryable request-configuration failure before reserving response ids. */ +export async function publishConfigurationFailure< + TContext extends object | undefined, + TState extends RunOperationState, +>( + lane: Lane, + drive: Drive, + capability: TState, + error: OperationError, + writes: Write[] = [], +): Promise { + const result = await lane.continueOperation( + capability, + (_state, current) => ({ + kind: "commit", + writes, + operationState: { + ...runScopeOf(current), + at: "run.failure_drain", + error, + provenance: { kind: "configuration" }, + }, + materialize: () => ({ kind: "continue" }) as const, + }), + drive.context, + ); + return result.kind === "cancel_requested" ? { kind: "continue" } : result.value; +} + 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}`); @@ -75,7 +176,7 @@ function deferredHandleIsValid( } /** Classify and atomically settle one assistant-generation or deferred-poll response. */ -export function settleResponse( +export function publishResponse( lane: Lane, drive: Drive, intent: ResponseIntent, diff --git a/packages/agent/src/harness/runtime/drive/tool-placement.ts b/packages/agent/src/harness/runtime/drive/tool-placement.ts new file mode 100644 index 000000000..267852f16 --- /dev/null +++ b/packages/agent/src/harness/runtime/drive/tool-placement.ts @@ -0,0 +1,326 @@ +import type { AssistantMessage, ToolResultMessage } from "@earendil-works/pi-ai"; +import type { AgentToolCall } from "../../../types.ts"; +import type { HarnessEvent } from "../../agent-harness.ts"; +import { insertEntry, insertUsage } from "../../session/commit.ts"; +import { SessionInvariantError } from "../../session/session.ts"; +import { + type MessageEntry, + type NewEntry, + type OperationState, + type RunCheckpointOperation, + type RunToolsOperation, + runScopeOf, + type ToolBatch, + type ToolCall, + type UsageRow, + type Write, +} from "../../session/types.ts"; +import { + branchTip, + deleteValue, + laneConfig, + operationToolArgsPrefix, + pendingEntry, + setValue, +} from "../../session/values.ts"; +import type { Lane } from "../lane.ts"; +import type { Drive } from "../types.ts"; + +export type ToolBatchSource = { + assistant: AssistantMessage; + calls: Map; +}; + +export async function readToolBatchSource( + lane: Lane, + drive: Drive, + batch: ToolBatch, +): Promise { + return lane.command(async (_state, reader) => { + const entry = (await reader.getEntries([batch.assistantEntryId], drive.context)).get(batch.assistantEntryId); + if (entry?.type !== "message" || entry.message.role !== "assistant") { + throw new SessionInvariantError("Tool batch assistant entry is invalid"); + } + const calls = new Map(); + for (const call of batch.calls) { + const block = entry.message.content[call.sourceIndex]; + if (block?.type !== "toolCall") { + throw new SessionInvariantError( + `Tool call source index ${call.sourceIndex} does not name a tool-call block`, + ); + } + calls.set(call.sourceIndex, block); + } + return { kind: "return", result: { assistant: entry.message, calls } }; + }, drive.context); +} + +type PlacementItem = { + call: Extract; + message: ToolResultMessage; +}; + +type PlacementRead = { + items: PlacementItem[]; + turnResults?: ToolResultMessage[]; +}; + +function isToolResultMessage(value: unknown): value is ToolResultMessage { + return typeof value === "object" && value !== null && "role" in value && value.role === "toolResult"; +} + +export function toolCallFor(sources: ToolBatchSource, call: ToolCall): AgentToolCall { + const source = sources.calls.get(call.sourceIndex); + if (source === undefined) { + throw new SessionInvariantError(`Tool call source index ${call.sourceIndex} is invalid`); + } + return source; +} + +export function withToolBatch(run: RunToolsOperation, batch: ToolBatch): RunToolsOperation { + return { ...runScopeOf(run), at: "run.tools", batch }; +} + +async function readPlacement( + lane: Lane, + drive: Drive, + sources: ToolBatchSource, +): Promise { + return lane.command(async (state, reader) => { + const operation = state.operation; + if (operation?.state.at !== "run.tools") return { kind: "return", result: undefined }; + const current = operation.state.batch; + let first = current.calls.findIndex((call) => call.status !== "completed"); + if (first === -1) return { kind: "return", result: undefined }; + const ready: Array> = []; + while (first < current.calls.length) { + const call = current.calls[first]!; + if (call.status !== "outcome_ready") break; + ready.push(call); + first += 1; + } + if (ready.length === 0) return { kind: "return", result: undefined }; + + const items: PlacementItem[] = []; + for (const call of ready) { + const stored = await reader.getValue(pendingEntry(call.resultEntryId), drive.context); + if (stored?.value.type !== "message" || !isToolResultMessage(stored.value.payload)) { + throw new SessionInvariantError(`Tool call ${call.resultEntryId} is missing its staged result`); + } + const source = toolCallFor(sources, call); + if (stored.value.payload.toolCallId !== source.id || stored.value.payload.toolName !== source.name) { + throw new SessionInvariantError(`Tool call ${call.resultEntryId} has a mismatched staged result`); + } + items.push({ call, message: stored.value.payload }); + } + + let turnResults: ToolResultMessage[] | undefined; + if (first === current.calls.length) { + const placedIds = current.calls + .filter((call) => call.status === "completed") + .map((call) => call.resultEntryId); + const placed = await reader.getEntries(placedIds, drive.context); + const staged = new Map(items.map((item) => [item.call.resultEntryId, item.message])); + turnResults = current.calls.map((call) => { + const message = + staged.get(call.resultEntryId) ?? + (() => { + const entry = placed.get(call.resultEntryId); + return entry?.type === "message" && isToolResultMessage(entry.message) ? entry.message : undefined; + })(); + if (message === undefined) { + throw new SessionInvariantError(`Completed tool call ${call.resultEntryId} is missing its result entry`); + } + return message; + }); + } + return { + kind: "return", + result: { items, ...(turnResults === undefined ? {} : { turnResults }) }, + }; + }, drive.context); +} + +async function commitPlacement( + lane: Lane, + drive: Drive, + capability: RunToolsOperation, + read: PlacementRead, +): Promise { + const usageIds = read.items.map((item) => + item.message.usage === undefined ? undefined : lane.session.idGenerator.next(), + ); + return lane.settleOperation( + capability, + async (state, run, _meta, reader) => { + const current = run.batch; + const writes: Write[] = []; + const eventEntries: Array<{ entry: NewEntry; seqIndex: number }> = []; + const eventUsage: Array<{ row: Omit; seqIndex: number }> = []; + let parentId = state.tipId; + for (const [index, item] of read.items.entries()) { + const entry: NewEntry = { + id: item.call.resultEntryId, + parentId, + type: "message", + message: item.message, + ...(item.call.terminate ? { terminate: true } : {}), + }; + eventEntries.push({ entry, seqIndex: writes.length }); + writes.push(insertEntry(entry), deleteValue(pendingEntry(item.call.resultEntryId))); + const usageId = usageIds[index]; + if (usageId !== undefined && item.message.usage !== undefined) { + const row: Omit = { + id: usageId, + usage: item.message.usage, + entryId: item.call.resultEntryId, + adjustment: false, + }; + eventUsage.push({ row, seqIndex: writes.length }); + writes.push(insertUsage(row)); + } + parentId = item.call.resultEntryId; + } + + const completedCalls = current.calls.map((call) => { + const item = read.items.find( + (candidate) => + candidate.call.sourceIndex === call.sourceIndex && + candidate.call.resultEntryId === call.resultEntryId, + ); + return item === undefined + ? call + : { + status: "completed" as const, + sourceIndex: call.sourceIndex, + resultEntryId: call.resultEntryId, + terminate: item.call.terminate, + }; + }); + const complete = completedCalls.every((call) => call.status === "completed"); + let nextConfiguration = state.configuration; + const addedNames: string[] = []; + for (const item of read.items) { + for (const name of item.message.addedToolNames ?? []) { + if (!nextConfiguration.activeToolNames.includes(name) && !addedNames.includes(name)) { + addedNames.push(name); + } + } + } + if (addedNames.length !== 0) { + nextConfiguration = { + ...nextConfiguration, + activeToolNames: [...nextConfiguration.activeToolNames, ...addedNames], + }; + writes.push(setValue(laneConfig(lane.name), nextConfiguration)); + } + writes.push(setValue(branchTip(lane.name), parentId)); + + let nextRun: OperationState; + if (complete) { + const allTerminate = completedCalls.every((call) => call.status === "completed" && call.terminate); + const checkpoint: RunCheckpointOperation = { + ...runScopeOf(run), + at: "run.checkpoint", + continuation: allTerminate + ? { kind: "may_finish", includeFinalAssistant: false } + : { kind: "need_assistant", overflowRecoveryUsed: false }, + triggerEntryId: parentId!, + }; + nextRun = checkpoint; + const args = await reader.scanValues( + operationToolArgsPrefix(drive.operationId, capability.batch.turnId), + drive.context, + ); + writes.push(...args.map(({ address }) => deleteValue(address))); + } else { + nextRun = withToolBatch(run, { ...current, calls: completedCalls }); + } + return { + kind: "commit", + writes, + operationState: nextRun, + lane: { tipId: parentId, configuration: nextConfiguration }, + materialize: () => complete, + events: (commit) => { + const events: HarnessEvent[] = []; + for (const { entry, seqIndex } of eventEntries) { + events.push({ + type: "entry_added", + lane: lane.name, + entry: { ...entry, seq: commit.seqs[seqIndex]!, timestamp: commit.timestamp }, + }); + const usage = eventUsage.find((candidate) => candidate.row.entryId === entry.id); + if (usage !== undefined) { + events.push({ + type: "usage", + lane: lane.name, + row: { ...usage.row, seq: commit.seqs[usage.seqIndex]! }, + totals: commit.stats.usage, + }); + } + } + if (addedNames.length !== 0) { + events.push({ + type: "config_update", + lane: lane.name, + property: "activeTools", + previous: state.configuration.activeToolNames, + value: nextConfiguration.activeToolNames, + }); + } + return events; + }, + }; + }, + drive.context, + ); +} + +export async function materializeReady( + lane: Lane, + drive: Drive, + capability: RunToolsOperation, + sources: ToolBatchSource, + recovery: boolean, +): Promise { + const read = await readPlacement(lane, drive, sources); + if (read === undefined) return; + await lane.emitBatch( + read.items.flatMap(({ call, message }) => [ + { + type: "message_start" as const, + lane: lane.name, + runId: drive.operationId, + message, + ...(recovery ? { recovery: true as const } : {}), + }, + { + type: "message_end" as const, + lane: lane.name, + runId: drive.operationId, + message, + entryId: call.resultEntryId, + ...(recovery ? { recovery: true as const } : {}), + }, + ]), + drive.context, + ); + const complete = await commitPlacement(lane, drive, capability, read); + if (complete && read.turnResults !== undefined) { + await lane.emitBatch( + [ + { + type: "turn_end", + lane: lane.name, + runId: drive.operationId, + turnId: capability.batch.turnId, + message: sources.assistant, + toolResults: read.turnResults, + ...(recovery ? { recovery: true as const } : {}), + }, + ], + drive.context, + ); + } +} diff --git a/packages/agent/src/harness/runtime/drive/tools.ts b/packages/agent/src/harness/runtime/drive/tools.ts index 4d73bc2a2..de9623640 100644 --- a/packages/agent/src/harness/runtime/drive/tools.ts +++ b/packages/agent/src/harness/runtime/drive/tools.ts @@ -1,6 +1,5 @@ -import type { AssistantMessage, ToolResultMessage } from "@earendil-works/pi-ai"; +import type { ToolResultMessage } from "@earendil-works/pi-ai"; import type { AgentToolCall, AgentToolResult } from "../../../types.ts"; -import type { HarnessEvent } from "../../agent-harness.ts"; import { AbortRequested } from "../../execution/effect-gate.ts"; import { applyBeforeToolDecision, @@ -8,30 +7,15 @@ import { createToolResultMessage, type ExecutedToolCall, executeToolCall, + type FinalizedToolCall, finalizeToolCall, prepareToolCall, } from "../../execution/tools.ts"; -import { insertEntry, insertUsage } from "../../session/commit.ts"; import { SessionInvariantError } from "../../session/session.ts"; +import type { JsonValue, RunToolsOperation, ToolBatch, ToolCall } from "../../session/types.ts"; import { - type JsonValue, - type MessageEntry, - type NewEntry, - type OperationState, - type RunCheckpointOperation, - type RunToolsOperation, - runScopeOf, - type ToolBatch, - type ToolCall, - type UsageRow, - type Write, -} from "../../session/types.ts"; -import { - branchTip, deleteValue, - laneConfig, operationToolArgs, - operationToolArgsPrefix, operationToolMemo, operationToolMemoPrefix, pendingEntry, @@ -41,28 +25,20 @@ import { import type { AgentHarnessTool, AgentHarnessToolInvocation } from "../../types.ts"; import type { Lane } from "../lane.ts"; import { openToolProgress } from "../progress.ts"; -import type { Drive, LaneState, ProcedureResult } from "../types.ts"; +import type { ContinueOperationResult, Drive, LaneState, ProcedureResult } from "../types.ts"; +import { + materializeReady, + readToolBatchSource, + type ToolBatchSource, + toolCallFor, + withToolBatch, +} from "./tool-placement.ts"; -type LocalResult = { kind: "committed" } | { kind: "state_changed" }; -type PlacementResult = { kind: "none" } | { kind: "committed"; complete: boolean }; - -type BatchSource = { - kind: "source"; - assistant: AssistantMessage; - calls: Map; -}; - -type PlacementItem = { - call: Extract; - message: ToolResultMessage; -}; - -type PlacementRead = { - kind: "placement"; - items: PlacementItem[]; - assistant: AssistantMessage; - turnResults?: ToolResultMessage[]; -}; +type ToolCallTask = { completion: Promise }; +type ToolOutcome = { message: ToolResultMessage; terminate: boolean }; +type PreparedToolInvocation = + | { kind: "ready"; cleared: ClearedToolCall } + | { kind: "outcome"; outcome: ToolOutcome }; const INTERRUPTION_MARKER = "[Tool execution was interrupted. The preceding output is the latest durable progress snapshot; newer live output may be missing, and the external outcome is unknown.]"; @@ -82,10 +58,6 @@ function currentBatch( return { run: operation.state, batch: operation.state.batch }; } -function toolOperation(lane: Lane): RunToolsOperation { - return lane.state.operation!.state as RunToolsOperation; -} - function findCall(batch: ToolBatch, sourceIndex: number, resultEntryId: string): ToolCall | undefined { return batch.calls.find((call) => call.sourceIndex === sourceIndex && call.resultEntryId === resultEntryId); } @@ -101,51 +73,11 @@ function replaceCall(batch: ToolBatch, replacement: ToolCall): ToolBatch { }; } -function withBatch(run: RunToolsOperation, batch: ToolBatch): RunToolsOperation { - return { ...runScopeOf(run), at: "run.tools", batch }; -} - function validateMemoName(name: string): void { if (name.length === 0) throw new TypeError("Tool invocation memo name must not be empty"); if (name.includes(":")) throw new TypeError("Tool invocation memo name must not contain ':'"); } -function isToolResultMessage(value: unknown): value is ToolResultMessage { - return typeof value === "object" && value !== null && "role" in value && value.role === "toolResult"; -} - -function toolCallFor(sources: BatchSource, call: ToolCall): AgentToolCall { - const source = sources.calls.get(call.sourceIndex); - if (source === undefined) { - throw new SessionInvariantError(`Tool call source index ${call.sourceIndex} is invalid`); - } - return source; -} - -async function readBatchSource( - lane: Lane, - drive: Drive, - batch: ToolBatch, -): Promise { - return lane.command(async (_state, reader) => { - const entry = (await reader.getEntries([batch.assistantEntryId], drive.context)).get(batch.assistantEntryId); - if (entry?.type !== "message" || entry.message.role !== "assistant") { - throw new SessionInvariantError("Tool batch assistant entry is invalid"); - } - const calls = new Map(); - for (const call of batch.calls) { - const block = entry.message.content[call.sourceIndex]; - if (block?.type !== "toolCall") { - throw new SessionInvariantError( - `Tool call source index ${call.sourceIndex} does not name a tool-call block`, - ); - } - calls.set(call.sourceIndex, block); - } - return { kind: "return", result: { kind: "source", assistant: entry.message, calls } }; - }, drive.context); -} - function invocationCapability( lane: Lane, drive: Drive, @@ -167,40 +99,28 @@ function invocationCapability( getMemo(name) { validateMemoName(name); if (!active) return Promise.reject(ended()); - return lane - .command<{ kind: "value"; value: JsonValue | undefined } | { kind: "ended" }>(async (state, reader) => { - if (!ownsEffect(state)) return { kind: "return", result: { kind: "ended" } as const }; - const stored = await reader.getValue( - operationToolMemo(drive.operationId, call.resultEntryId, name), - drive.context, - ); - return { - kind: "return", - result: { kind: "value", value: stored?.value } as const, - }; - }, drive.context) - .then((result) => { - if (result.kind === "ended") throw ended(); - return result.value; - }); + return lane.command(async (state, reader) => { + if (!ownsEffect(state)) return { kind: "reject", error: ended() }; + const stored = await reader.getValue( + operationToolMemo(drive.operationId, call.resultEntryId, name), + drive.context, + ); + return { kind: "return", result: stored?.value }; + }, drive.context); }, setMemo(name, value) { validateMemoName(name); if (!active) return Promise.reject(ended()); - return lane - .command<{ kind: "written" } | { kind: "ended" }>((state) => { - if (!ownsEffect(state)) return { kind: "return", result: { kind: "ended" } as const }; - const address = operationToolMemo(drive.operationId, call.resultEntryId, name); - return { - kind: "commit", - writes: [value === undefined ? deleteValue(address) : setValue(address, value)], - next: state, - materialize: () => ({ kind: "written" }) as const, - }; - }, drive.context) - .then((result) => { - if (result.kind === "ended") throw ended(); - }); + return lane.command((state) => { + if (!ownsEffect(state)) return { kind: "reject", error: ended() }; + const address = operationToolMemo(drive.operationId, call.resultEntryId, name); + return { + kind: "commit", + writes: [value === undefined ? deleteValue(address) : setValue(address, value)], + next: state, + materialize: () => undefined, + }; + }, drive.context); }, }, expire() { @@ -250,80 +170,68 @@ function truncatedMessage(toolCall: AgentToolCall): ToolResultMessage { ]); } -async function commitIntent( +function outcomeFromFinalizedCall(finalized: FinalizedToolCall): ToolOutcome { + return { message: createToolResultMessage(finalized), terminate: finalized.terminate }; +} + +async function publishToolIntent( lane: Lane, drive: Drive, - batch: ToolBatch, + run: RunToolsOperation, planned: Extract, args: Record, replay: "never" | "safe", -): Promise { - const result = await lane.continueOperation( - toolOperation(lane), +): Promise>> { + return lane.continueOperation( + run, (_state, run) => { - const currentBatchState = run.batch; - const current = findCall(currentBatchState, planned.sourceIndex, planned.resultEntryId); - if (current?.status !== "planned") { - return { kind: "return", result: { kind: "state_changed" } as const }; - } const effectPending: Extract = { status: "effect_pending", - sourceIndex: current.sourceIndex, - resultEntryId: current.resultEntryId, + sourceIndex: planned.sourceIndex, + resultEntryId: planned.resultEntryId, replay, }; - const nextRun = withBatch(run, replaceCall(currentBatchState, effectPending)); return { kind: "commit", - writes: [setValue(operationToolArgs(drive.operationId, batch.turnId, current.sourceIndex), args)], - operationState: nextRun, - materialize: () => ({ kind: "committed" }) as const, + writes: [setValue(operationToolArgs(drive.operationId, run.batch.turnId, planned.sourceIndex), args)], + operationState: withToolBatch(run, replaceCall(run.batch, effectPending)), + materialize: () => effectPending, }; }, drive.context, ); - return result ?? { kind: "state_changed" }; } -async function stageOutcome( +async function publishToolOutcome( lane: Lane, drive: Drive, + capability: RunToolsOperation, call: ToolCall, message: ToolResultMessage, terminate: boolean, -): Promise { - return lane.settleOperation( - toolOperation(lane), +): Promise { + await lane.settleOperation( + capability, async (_state, run, _meta, reader) => { - const currentBatchState = run.batch; - const current = findCall(currentBatchState, call.sourceIndex, call.resultEntryId); - if ( - current === undefined || - current.status !== call.status || - (current.status === "effect_pending" && call.status === "effect_pending" && current.replay !== call.replay) - ) { - return { kind: "return", result: { kind: "state_changed" } as const }; - } const memos = await reader.scanValues( - operationToolMemoPrefix(drive.operationId, current.resultEntryId), + operationToolMemoPrefix(drive.operationId, call.resultEntryId), drive.context, ); const outcome: Extract = { status: "outcome_ready", - sourceIndex: current.sourceIndex, - resultEntryId: current.resultEntryId, + sourceIndex: call.sourceIndex, + resultEntryId: call.resultEntryId, terminate: run.control.status === "running" && terminate, }; - const nextRun = withBatch(run, replaceCall(currentBatchState, outcome)); return { kind: "commit", writes: [ - setValue(pendingEntry(current.resultEntryId), { type: "message", payload: message }), - deleteValue(pendingToolOutput(drive.operationId, current.resultEntryId)), + setValue(pendingEntry(call.resultEntryId), { type: "message", payload: message }), + deleteValue(pendingToolOutput(drive.operationId, call.resultEntryId)), ...memos.map(({ address }) => deleteValue(address)), ], - operationState: nextRun, - materialize: () => ({ kind: "committed" }) as const, + operationState: withToolBatch(run, replaceCall(run.batch, outcome)), + materialize: () => undefined, }; }, drive.context, @@ -335,52 +243,33 @@ async function clearReplayCheckpoint( drive: Drive, batch: ToolBatch, call: Extract, -): Promise<{ kind: "args"; args: Record } | { kind: "state_changed" }> { - return lane.command<{ kind: "args"; args: Record } | { kind: "state_changed" }>( - async (state, reader) => { - const run = state.operation!.state as RunToolsOperation; - const current = findCall(run.batch, call.sourceIndex, call.resultEntryId); - if (current?.status !== "effect_pending") { - return { kind: "return", result: { kind: "state_changed" } as const }; - } - const stored = await reader.getValue( - operationToolArgs(drive.operationId, batch.turnId, call.sourceIndex), - drive.context, - ); - if (stored === undefined) { - throw new SessionInvariantError(`Tool call ${call.resultEntryId} is missing persisted arguments`); - } - return { - kind: "commit", - writes: [deleteValue(pendingToolOutput(drive.operationId, call.resultEntryId))], - next: state, - materialize: () => ({ kind: "args", args: stored.value }) as const, - }; - }, - drive.context, - ); +): Promise> { + return lane.command(async (state, reader) => { + const stored = await reader.getValue( + operationToolArgs(drive.operationId, batch.turnId, call.sourceIndex), + drive.context, + ); + if (stored === undefined) { + throw new SessionInvariantError(`Tool call ${call.resultEntryId} is missing persisted arguments`); + } + return { + kind: "commit", + writes: [deleteValue(pendingToolOutput(drive.operationId, call.resultEntryId))], + next: state, + materialize: () => stored.value, + }; + }, drive.context); } -async function readCheckpoint( +function readCheckpoint( lane: Lane, drive: Drive, call: Extract, -): Promise<{ kind: "checkpoint"; value: AgentToolResult | undefined } | { kind: "state_changed" }> { - return lane.command<{ kind: "checkpoint"; value: AgentToolResult | undefined } | { kind: "state_changed" }>( - async (state, reader) => { - const run = state.operation!.state as RunToolsOperation; - const current = findCall(run.batch, call.sourceIndex, call.resultEntryId); - if (current?.status !== "effect_pending") { - return { kind: "return", result: { kind: "state_changed" } as const }; - } - const stored = await reader.getValue(pendingToolOutput(drive.operationId, call.resultEntryId), drive.context); - return { - kind: "return", - result: { kind: "checkpoint", value: stored?.value } as const, - }; - }, - drive.context, - ); +): Promise | undefined> { + return lane.command(async (_state, reader) => { + const stored = await reader.getValue(pendingToolOutput(drive.operationId, call.resultEntryId), drive.context); + return { kind: "return", result: stored?.value }; + }, drive.context); } async function resolveToolContext( @@ -391,7 +280,7 @@ async function resolveToolContext( return (typeof source === "function" ? await source(drive.context) : source) as TContext; } -async function executeAndStage( +async function performToolInvocation( lane: Lane, drive: Drive, batch: ToolBatch, @@ -399,7 +288,7 @@ async function executeAndStage( cleared: ClearedToolCall, toolContext: TContext, recovery: boolean, -): Promise { +): Promise { const capability = invocationCapability(lane, drive, batch, call); const progress = openToolProgress(lane, drive, batch.turnId, call.sourceIndex, call.resultEntryId); let latestUpdateDelivery: Promise = Promise.resolve(); @@ -462,13 +351,10 @@ async function executeAndStage( await progress.drain(); if (!(error instanceof AbortRequested)) throw error; await error.cancellation; - return stageOutcome( - lane, - drive, - call, - recovery ? interruptedMessage(cleared.toolCall, undefined) : abortedMessage(cleared.toolCall), - false, - ); + return { + message: recovery ? interruptedMessage(cleared.toolCall, undefined) : abortedMessage(cleared.toolCall), + terminate: false, + }; } const executed = await execution.finally(() => { @@ -501,8 +387,7 @@ async function executeAndStage( if (!(error instanceof AbortRequested)) throw error; patch = undefined; } - const current = currentBatch(lane); - const cancelled = current?.run.control.status === "cancel_requested"; + const cancelled = lane.state.operation!.state.control.status === "cancel_requested"; const finalized = finalizeToolCall(cleared, executed, cancelled ? { ...patch, terminate: false } : patch); await lane.emitBatch( [ @@ -521,292 +406,22 @@ async function executeAndStage( ], drive.context, ); - return stageOutcome(lane, drive, call, createToolResultMessage(finalized), cancelled ? false : finalized.terminate); + return { message: createToolResultMessage(finalized), terminate: cancelled ? false : finalized.terminate }; } -async function readPlacement( +async function prepareToolInvocation( lane: Lane, drive: Drive, -): Promise { - return lane.command(async (state, reader) => { - const operation = state.operation; - if (operation?.state.at !== "run.tools") return { kind: "return", result: { kind: "none" } as const }; - const current = operation.state.batch; - let first = current.calls.findIndex((call) => call.status !== "completed"); - if (first === -1) return { kind: "return", result: { kind: "none" } as const }; - const ready: Array> = []; - while (first < current.calls.length) { - const call = current.calls[first]!; - if (call.status !== "outcome_ready") break; - ready.push(call); - first += 1; - } - if (ready.length === 0) return { kind: "return", result: { kind: "none" } as const }; - - const assistantEntry = (await reader.getEntries([current.assistantEntryId], drive.context)).get( - current.assistantEntryId, - ); - if (assistantEntry?.type !== "message" || assistantEntry.message.role !== "assistant") { - throw new SessionInvariantError("Tool placement assistant entry is invalid"); - } - const items: PlacementItem[] = []; - for (const call of ready) { - const stored = await reader.getValue(pendingEntry(call.resultEntryId), drive.context); - if (stored?.value.type !== "message" || !isToolResultMessage(stored.value.payload)) { - throw new SessionInvariantError(`Tool call ${call.resultEntryId} is missing its staged result`); - } - const block = assistantEntry.message.content[call.sourceIndex]; - if ( - block?.type !== "toolCall" || - stored.value.payload.toolCallId !== block.id || - stored.value.payload.toolName !== block.name - ) { - throw new SessionInvariantError(`Tool call ${call.resultEntryId} has a mismatched staged result`); - } - items.push({ call, message: stored.value.payload }); - } - - let turnResults: ToolResultMessage[] | undefined; - if (first === current.calls.length) { - const placedIds = current.calls - .filter((call) => call.status === "completed") - .map((call) => call.resultEntryId); - const placed = await reader.getEntries(placedIds, drive.context); - const staged = new Map(items.map((item) => [item.call.resultEntryId, item.message])); - turnResults = current.calls.map((call) => { - const message = - staged.get(call.resultEntryId) ?? - (() => { - const entry = placed.get(call.resultEntryId); - return entry?.type === "message" && isToolResultMessage(entry.message) ? entry.message : undefined; - })(); - if (message === undefined) { - throw new SessionInvariantError(`Completed tool call ${call.resultEntryId} is missing its result entry`); - } - return message; - }); - } - return { - kind: "return", - result: { - kind: "placement", - items, - assistant: assistantEntry.message, - ...(turnResults === undefined ? {} : { turnResults }), - }, - }; - }, drive.context); -} - -async function commitPlacement( - lane: Lane, - drive: Drive, - batch: ToolBatch, - read: PlacementRead, -): Promise { - const usageIds = read.items.map((item) => - item.message.usage === undefined ? undefined : lane.session.idGenerator.next(), - ); - return lane.settleOperation( - toolOperation(lane), - async (state, run, _meta, reader) => { - const current = run.batch; - const writes: Write[] = []; - const eventEntries: Array<{ entry: NewEntry; seqIndex: number }> = []; - const eventUsage: Array<{ row: Omit; seqIndex: number }> = []; - let parentId = state.tipId; - for (const [index, item] of read.items.entries()) { - const entry: NewEntry = { - id: item.call.resultEntryId, - parentId, - type: "message", - message: item.message, - ...(item.call.terminate ? { terminate: true } : {}), - }; - eventEntries.push({ entry, seqIndex: writes.length }); - writes.push(insertEntry(entry), deleteValue(pendingEntry(item.call.resultEntryId))); - const usageId = usageIds[index]; - if (usageId !== undefined && item.message.usage !== undefined) { - const row: Omit = { - id: usageId, - usage: item.message.usage, - entryId: item.call.resultEntryId, - adjustment: false, - }; - eventUsage.push({ row, seqIndex: writes.length }); - writes.push(insertUsage(row)); - } - parentId = item.call.resultEntryId; - } - - const completedCalls = current.calls.map((call) => { - const item = read.items.find( - (candidate) => - candidate.call.sourceIndex === call.sourceIndex && - candidate.call.resultEntryId === call.resultEntryId, - ); - return item === undefined - ? call - : { - status: "completed" as const, - sourceIndex: call.sourceIndex, - resultEntryId: call.resultEntryId, - terminate: item.call.terminate, - }; - }); - const complete = completedCalls.every((call) => call.status === "completed"); - let nextConfiguration = state.configuration; - const addedNames: string[] = []; - for (const item of read.items) { - for (const name of item.message.addedToolNames ?? []) { - if (!nextConfiguration.activeToolNames.includes(name) && !addedNames.includes(name)) - addedNames.push(name); - } - } - if (addedNames.length !== 0) { - nextConfiguration = { - ...nextConfiguration, - activeToolNames: [...nextConfiguration.activeToolNames, ...addedNames], - }; - writes.push(setValue(laneConfig(lane.name), nextConfiguration)); - } - writes.push(setValue(branchTip(lane.name), parentId)); - - let nextRun: OperationState; - if (complete) { - const allTerminate = completedCalls.every((call) => call.status === "completed" && call.terminate); - const checkpoint: RunCheckpointOperation = { - ...runScopeOf(run), - at: "run.checkpoint", - continuation: allTerminate - ? { kind: "may_finish", includeFinalAssistant: false } - : { kind: "need_assistant", overflowRecoveryUsed: false }, - triggerEntryId: parentId!, - }; - nextRun = checkpoint; - const args = await reader.scanValues( - operationToolArgsPrefix(drive.operationId, batch.turnId), - drive.context, - ); - writes.push(...args.map(({ address }) => deleteValue(address))); - } else { - nextRun = withBatch(run, { ...current, calls: completedCalls }); - } - return { - kind: "commit", - writes, - operationState: nextRun, - lane: { tipId: parentId, configuration: nextConfiguration }, - materialize: () => ({ kind: "committed", complete }) as const, - events: (commit) => { - const events: HarnessEvent[] = []; - for (const { entry, seqIndex } of eventEntries) { - events.push({ - type: "entry_added", - lane: lane.name, - entry: { ...entry, seq: commit.seqs[seqIndex]!, timestamp: commit.timestamp }, - }); - const usage = eventUsage.find((candidate) => candidate.row.entryId === entry.id); - if (usage !== undefined) { - events.push({ - type: "usage", - lane: lane.name, - row: { ...usage.row, seq: commit.seqs[usage.seqIndex]! }, - totals: commit.stats.usage, - }); - } - } - if (addedNames.length !== 0) { - events.push({ - type: "config_update", - lane: lane.name, - property: "activeTools", - previous: state.configuration.activeToolNames, - value: nextConfiguration.activeToolNames, - }); - } - return events; - }, - }; - }, - drive.context, - ); -} - -async function materializeReady( - lane: Lane, - drive: Drive, - batch: ToolBatch, - recovery: boolean, -): Promise { - const read = await readPlacement(lane, drive); - if (read.kind === "none") return read; - await lane.emitBatch( - read.items.flatMap(({ call, message }) => [ - { - type: "message_start" as const, - lane: lane.name, - runId: drive.operationId, - message, - ...(recovery ? { recovery: true as const } : {}), - }, - { - type: "message_end" as const, - lane: lane.name, - runId: drive.operationId, - message, - entryId: call.resultEntryId, - ...(recovery ? { recovery: true as const } : {}), - }, - ]), - drive.context, - ); - const committed = await commitPlacement(lane, drive, batch, read); - if (committed.kind === "committed" && committed.complete && read.turnResults !== undefined) { - await lane.emitBatch( - [ - { - type: "turn_end", - lane: lane.name, - runId: drive.operationId, - turnId: batch.turnId, - message: read.assistant, - toolResults: read.turnResults, - ...(recovery ? { recovery: true as const } : {}), - }, - ], - drive.context, - ); - } - return committed; -} - -async function beginFreshCall( - lane: Lane, - drive: Drive, - batch: ToolBatch, - sources: BatchSource, + sources: ToolBatchSource, call: Extract, tools: AgentHarnessTool[], - toolContext: TContext, - recovery: boolean, -): Promise<{ job: Promise }> { +): Promise> { const toolCall = toolCallFor(sources, call); if (sources.assistant.stopReason === "length") { - return { job: stageOutcome(lane, drive, call, truncatedMessage(toolCall), false) }; + return { kind: "outcome", outcome: { message: truncatedMessage(toolCall), terminate: false } }; } const prepared = prepareToolCall(toolCall, tools); - if ("kind" in prepared) { - const finalized = { - toolCall: prepared.toolCall, - result: prepared.result, - isError: prepared.isError, - terminate: prepared.terminate, - }; - return { - job: stageOutcome(lane, drive, call, createToolResultMessage(finalized), finalized.terminate), - }; - } + if ("kind" in prepared) return { kind: "outcome", outcome: outcomeFromFinalizedCall(prepared) }; let decision: Awaited>>; try { @@ -825,65 +440,88 @@ async function beginFreshCall( } catch (error) { if (!(error instanceof AbortRequested)) throw error; await error.cancellation; - return { job: stageOutcome(lane, drive, call, abortedMessage(toolCall), false) }; + return { kind: "outcome", outcome: { message: abortedMessage(toolCall), terminate: false } }; } const cleared = applyBeforeToolDecision(prepared, decision); - if ("kind" in cleared) { - const finalized = { - toolCall: cleared.toolCall, - result: cleared.result, - isError: cleared.isError, - terminate: cleared.terminate, - }; + return "kind" in cleared + ? { kind: "outcome", outcome: outcomeFromFinalizedCall(cleared) } + : { kind: "ready", cleared }; +} + +async function startToolInvocation( + lane: Lane, + drive: Drive, + run: RunToolsOperation, + sources: ToolBatchSource, + call: Extract, + tools: AgentHarnessTool[], + toolContext: TContext, + recovery: boolean, +): Promise { + const prepared = await prepareToolInvocation(lane, drive, sources, call, tools); + if (prepared.kind === "outcome") { return { - job: stageOutcome(lane, drive, call, createToolResultMessage(finalized), finalized.terminate), + completion: publishToolOutcome(lane, drive, run, call, prepared.outcome.message, prepared.outcome.terminate), }; } - const intent = await commitIntent(lane, drive, batch, call, cleared.args, cleared.tool.replay ?? "never"); - if (intent.kind !== "committed") { - return { job: stageOutcome(lane, drive, call, abortedMessage(toolCall), false) }; - } - const effectPending: Extract = { - status: "effect_pending", - sourceIndex: call.sourceIndex, - resultEntryId: call.resultEntryId, - replay: cleared.tool.replay ?? "never", - }; + const effectPending = await publishToolIntent( + lane, + drive, + run, + call, + prepared.cleared.args, + prepared.cleared.tool.replay ?? "never", + ); return { - job: executeAndStage(lane, drive, batch, effectPending, cleared, toolContext, recovery), + completion: + effectPending.kind === "cancel_requested" + ? publishToolOutcome(lane, drive, run, call, abortedMessage(prepared.cleared.toolCall), false) + : performToolInvocation( + lane, + drive, + run.batch, + effectPending.value, + prepared.cleared, + toolContext, + recovery, + ).then((outcome) => + publishToolOutcome(lane, drive, run, effectPending.value, outcome.message, outcome.terminate), + ), }; } -async function beginRecoveryCall( +async function recoverToolInvocation( lane: Lane, drive: Drive, - batch: ToolBatch, - sources: BatchSource, + run: RunToolsOperation, + sources: ToolBatchSource, call: Extract, toolsByName: Map>, toolContext: TContext, cancelled: boolean, -): Promise<{ job: Promise }> { +): Promise { const toolCall = toolCallFor(sources, call); const tool = toolsByName.get(toolCall.name); if (!cancelled && call.replay === "safe" && tool?.replay === "safe") { - const replay = await clearReplayCheckpoint(lane, drive, batch, call); - if (replay.kind !== "args") return { job: Promise.resolve(replay) }; - const cleared: ClearedToolCall = { toolCall, tool, args: replay.args }; - return { job: executeAndStage(lane, drive, batch, call, cleared, toolContext, true) }; + const args = await clearReplayCheckpoint(lane, drive, run.batch, call); + const cleared: ClearedToolCall = { toolCall, tool, args }; + return { + completion: performToolInvocation(lane, drive, run.batch, call, cleared, toolContext, true).then((outcome) => + publishToolOutcome(lane, drive, run, call, outcome.message, outcome.terminate), + ), + }; } const checkpoint = await readCheckpoint(lane, drive, call); - if (checkpoint.kind !== "checkpoint") return { job: Promise.resolve(checkpoint) }; return { - job: stageOutcome(lane, drive, call, interruptedMessage(toolCall, checkpoint.value), false), + completion: publishToolOutcome(lane, drive, run, call, interruptedMessage(toolCall, checkpoint), false), }; } async function runSequential( lane: Lane, drive: Drive, - batch: ToolBatch, - sources: BatchSource, + run: RunToolsOperation, + sources: ToolBatchSource, execution: | { tools: AgentHarnessTool[]; @@ -893,8 +531,9 @@ async function runSequential( | undefined, recovery: boolean, ): Promise { + const { batch } = run; for (let transition = 0; transition <= batch.calls.length * 2 + 1; transition += 1) { - await materializeReady(lane, drive, batch, recovery); + await materializeReady(lane, drive, run, sources, recovery); const current = currentBatch(lane); if (current === undefined) return { kind: "continue" }; const call = current.batch.calls.find((candidate) => candidate.status !== "completed"); @@ -904,11 +543,10 @@ async function runSequential( if (current.run.control.status === "cancel_requested") { const toolCall = toolCallFor(sources, call); if (call.status === "planned") { - await stageOutcome(lane, drive, call, abortedMessage(toolCall), false); + await publishToolOutcome(lane, drive, current.run, call, abortedMessage(toolCall), false); } else { const checkpoint = await readCheckpoint(lane, drive, call); - if (checkpoint.kind === "state_changed") continue; - await stageOutcome(lane, drive, call, interruptedMessage(toolCall, checkpoint.value), false); + await publishToolOutcome(lane, drive, current.run, call, interruptedMessage(toolCall, checkpoint), false); } continue; } @@ -916,18 +554,27 @@ async function runSequential( if (execution === undefined) throw new SessionInvariantError("Running tool batch is missing execution context"); const started = call.status === "planned" - ? await beginFreshCall(lane, drive, batch, sources, call, execution.tools, execution.toolContext, recovery) - : await beginRecoveryCall( + ? await startToolInvocation( lane, drive, - batch, + current.run, + sources, + call, + execution.tools, + execution.toolContext, + recovery, + ) + : await recoverToolInvocation( + lane, + drive, + current.run, sources, call, execution.toolsByName, execution.toolContext, false, ); - await started.job; + await started.completion; } throw new SessionInvariantError("Sequential tool batch exceeded its bounded transition count"); } @@ -935,42 +582,38 @@ async function runSequential( async function runParallel( lane: Lane, drive: Drive, - batch: ToolBatch, - sources: BatchSource, + run: RunToolsOperation, + sources: ToolBatchSource, tools: AgentHarnessTool[], toolsByName: Map>, toolContext: TContext, recovery: boolean, ): Promise { - let materialization: Promise = Promise.resolve({ kind: "none" }); - const scheduleMaterialization = (): Promise => { - const scheduled = materialization.then(() => materializeReady(lane, drive, batch, recovery)); - materialization = scheduled.catch(() => ({ kind: "none" })); + const { batch } = run; + let materialization = Promise.resolve(); + const scheduleMaterialization = (): Promise => { + const scheduled = materialization.then(() => materializeReady(lane, drive, run, sources, recovery)); + materialization = scheduled.catch(() => {}); return scheduled; }; - const jobs: Promise[] = []; - const snapshot = currentBatch(lane); - if (snapshot === undefined) return { kind: "continue" }; - for (const call of snapshot.batch.calls) { + const jobs: Promise[] = []; + for (const call of batch.calls) { if (call.status === "completed" || call.status === "outcome_ready") continue; - const current = currentBatch(lane); - if (current === undefined) break; const started = call.status === "planned" - ? await beginFreshCall(lane, drive, batch, sources, call, tools, toolContext, recovery) - : await beginRecoveryCall( + ? await startToolInvocation(lane, drive, run, sources, call, tools, toolContext, recovery) + : await recoverToolInvocation( lane, drive, - batch, + run, sources, call, toolsByName, toolContext, - current.run.control.status === "cancel_requested", + lane.state.operation!.state.control.status === "cancel_requested", ); - const job = started.job.then(async (result) => { - if (result.kind === "committed") await scheduleMaterialization(); - return result; + const job = started.completion.then(async () => { + await scheduleMaterialization(); }); void job.catch(() => {}); jobs.push(job); @@ -1002,19 +645,19 @@ export async function runTools( drive.context, ); } - await materializeReady(lane, drive, batch, recovery); - if (currentBatch(lane) === undefined) return { kind: "continue" }; - - const sources = await readBatchSource(lane, drive, batch); - if (toolOperation(lane).control.status === "cancel_requested") { - return runSequential(lane, drive, batch, sources, undefined, recovery); + const sources = await readToolBatchSource(lane, drive, batch); + await materializeReady(lane, drive, run, sources, recovery); + const current = currentBatch(lane); + if (current === undefined) return { kind: "continue" }; + if (current.run.control.status === "cancel_requested") { + return runSequential(lane, drive, current.run, sources, undefined, recovery); } const config = lane.readConfig(); const active = new Set(batch.configuration.activeToolNames); const tools = config.tools.filter((tool) => active.has(tool.name)); const toolsByName = new Map(tools.map((tool) => [tool.name, tool])); const toolContext = await resolveToolContext(lane, drive); - return toolOperation(lane).settings.toolExecution === "sequential" - ? runSequential(lane, drive, batch, sources, { tools, toolsByName, toolContext }, recovery) - : runParallel(lane, drive, batch, sources, tools, toolsByName, toolContext, recovery); + return run.settings.toolExecution === "sequential" + ? runSequential(lane, drive, current.run, sources, { tools, toolsByName, toolContext }, recovery) + : runParallel(lane, drive, current.run, sources, tools, toolsByName, toolContext, recovery); } diff --git a/packages/agent/src/harness/runtime/harness.ts b/packages/agent/src/harness/runtime/harness.ts index d99180f4f..846e9cebf 100644 --- a/packages/agent/src/harness/runtime/harness.ts +++ b/packages/agent/src/harness/runtime/harness.ts @@ -19,19 +19,10 @@ import { convertToLlm } from "../messages.ts"; import { HarnessClosed, HarnessFault, InvalidLane, UnknownTarget } from "../result.ts"; import { SessionInvariantError } from "../session/session.ts"; import type { LaneConfiguration, Session } from "../session/types.ts"; -import { - branchTip, - deleteValue, - entryLabel, - laneConfig, - laneLastResult, - laneState, - sessionName, - setValue, -} from "../session/values.ts"; +import { branchTip, deleteValue, entryLabel, laneConfig, laneState, sessionName, setValue } from "../session/values.ts"; import type { AgentHarnessStreamOptions, AgentHarnessTool } from "../types.ts"; import { Lane } from "./lane.ts"; -import { restoreLaneState, restoreSession } from "./restore.ts"; +import { readLaneStorage, restoreLaneState, restoreSession } from "./restore.ts"; import { type Config, type LaneState, SliceNotImplemented } from "./types.ts"; type GlobalConfigProperty = GlobalConfigEventPayload["property"]; @@ -112,29 +103,20 @@ export class Harness implements AgentHarnes lane = this.lanesByName.get(name); if (lane !== undefined) return; - const [tip, configuration, storedState, lastResult] = await Promise.all([ - mutator.getValue(branchTip(name), context), - mutator.getValue(laneConfig(name), context), - mutator.getValue(laneState(name), context), - mutator.getValue(laneLastResult(name), context), - ]); - - if (tip !== undefined && configuration !== undefined && storedState !== undefined) { - const restored = await restoreLaneState(mutator, name, context); + const stored = await readLaneStorage(mutator, name, context); + if (stored.kind === "lane") { + const restored = await restoreLaneState(mutator, name, stored, context); lane = this.buildLane(name, restored); this.lanesByName.set(name, lane); return; } - if ((configuration === undefined) !== (storedState === undefined) || lastResult !== undefined) { - throw new SessionInvariantError(`Lane ${JSON.stringify(name)} has incomplete durable state`); - } - if (configuration !== undefined || storedState !== undefined) { - throw new SessionInvariantError(`Lane ${JSON.stringify(name)} has durable state without a branch`); - } - - const tipId = tip?.value ?? options.createAt ?? null; - if (tip === undefined && tipId !== null && !(await mutator.getEntries([tipId], context)).has(tipId)) { + const tipId = stored.kind === "branch" ? stored.tip.value : (options.createAt ?? null); + if ( + stored.kind === "absent" && + tipId !== null && + !(await mutator.getEntries([tipId], context)).has(tipId) + ) { throw new UnknownTarget({ targetId: tipId, message: `Unknown target: ${tipId}`, @@ -152,7 +134,7 @@ export class Harness implements AgentHarnes operation: null, }; const writes = [ - ...(tip === undefined ? [setValue(branchTip(name), tipId)] : []), + ...(stored.kind === "absent" ? [setValue(branchTip(name), tipId)] : []), setValue(laneConfig(name), attachedConfiguration), setValue(laneState(name), { currentOperationId: null, diff --git a/packages/agent/src/harness/runtime/lane.ts b/packages/agent/src/harness/runtime/lane.ts index 92745b206..6227eb814 100644 --- a/packages/agent/src/harness/runtime/lane.ts +++ b/packages/agent/src/harness/runtime/lane.ts @@ -15,7 +15,6 @@ import type { ModelIdentity, OperationAdmissionResult, OperationRequest, - QueuedItem, WatchHandle, } from "../agent-harness.ts"; import type { Context } from "../context.ts"; @@ -56,14 +55,16 @@ import { operationMeta as operationMetaValue, operationState as operationStateValue, operationToolArgs, - pendingAssistantFrames, pendingEntry, pendingToolOutput, setValue, } from "../session/values.ts"; import { formatSkillInvocation } from "../skills.ts"; +import { readAssistantFrames } from "./progress.ts"; +import { chainEntries, entryLifecycleEvents, readPendingMessages } from "./transcript.ts"; import { type Config, + type ContinueOperationResult, type Drive, type LaneCommand, type LaneState, @@ -286,7 +287,7 @@ export class Lane implements AgentLane { /** * Run an ordinary operation command only while durable control is running. Use this before starting new hooks, - * effects, or forward progress. Returns `undefined` without invoking the planner once cancellation is requested. + * effects, or forward progress. Returns `cancel_requested` without invoking the planner once cancellation is requested. */ continueOperation( capability: TState, @@ -297,13 +298,22 @@ export class Lane implements AgentLane { reader: SessionReader, ) => OperationCommand | Promise>, context: Context, - ): Promise { - return this.settleOperation( + ): Promise> { + return this.settleOperation>( capability, - (state, latest, meta, reader) => - latest.control.status === "cancel_requested" - ? { kind: "return", result: undefined } - : plan(state, latest, meta, reader), + async (state, latest, meta, reader) => { + if (latest.control.status === "cancel_requested") { + return { kind: "return", result: { kind: "cancel_requested" } }; + } + const decision = await plan(state, latest, meta, reader); + if (decision.kind === "return") { + return { kind: "return", result: { kind: "result", value: decision.result } }; + } + return { + ...decision, + materialize: (commit) => ({ kind: "result", value: decision.materialize(commit) }), + }; + }, context, ); } @@ -404,18 +414,14 @@ export class Lane implements AgentLane { }; } const capturedIds = [...state.pendingNextRun]; - const captured = await Promise.all( - capturedIds.map(async (id) => { - const stored = await reader.getValue(pendingEntry(id), context); - if (stored?.value.type !== "message") { - throw new SessionInvariantError(`Pending next-run entry ${id} is missing a message payload`); - } - if (stored.value.payload.role === "assistant" && stored.value.payload.stopReason === "pending") { - throw new SessionInvariantError(`Pending next-run entry ${id} contains a pending assistant message`); - } - return { id, message: stored.value.payload }; - }), + const captured = (await readPendingMessages(reader, capturedIds, "Pending next-run entry", context)).map( + ({ entryId, message }) => ({ id: entryId, message }), ); + for (const { id, message } of captured) { + if (message.role === "assistant" && message.stopReason === "pending") { + throw new SessionInvariantError(`Pending next-run entry ${id} contains a pending assistant message`); + } + } const placed = [...captured, ...prompt]; if (placed.length === 0) { return { @@ -430,12 +436,11 @@ export class Lane implements AgentLane { }; } - let parentId = state.tipId; - const entryWrites = placed.map(({ id, message }) => { - const write = insertEntry({ id, parentId, type: "message", message }); - parentId = id; - return write; - }); + const entries = chainEntries( + state.tipId, + placed.map(({ id, message }) => ({ id, type: "message" as const, message })), + ); + const parentId = entries[entries.length - 1]!.id; const meta = { operationId, lane: this.name, @@ -464,7 +469,7 @@ export class Lane implements AgentLane { return { kind: "commit", writes: [ - ...entryWrites, + ...entries.map((entry) => insertEntry(entry)), ...capturedIds.map((id) => deleteValue(pendingEntry(id))), setValue(branchTip(this.name), parentId), setValue(operationMetaValue(operationId), meta), @@ -475,26 +480,13 @@ export class Lane implements AgentLane { materialize: () => Result.ok({ operationId, kind: "run", startedAt }), events: (commit) => { const events: HarnessEvent[] = [{ type: "run_start", runId: operationId, lane: this.name }]; - for (const [index, item] of placed.entries()) { - const parent = index === 0 ? state.tipId : placed[index - 1]!.id; - const entry: Entry = { - id: item.id, - parentId: parent, - type: "message", - message: item.message, - seq: commit.seqs[index]!, - timestamp: commit.timestamp, - }; + for (const [index, entry] of entries.entries()) { events.push( - { type: "message_start", runId: operationId, message: item.message, lane: this.name }, - { - type: "message_end", - runId: operationId, - message: item.message, - entryId: item.id, - lane: this.name, - }, - { type: "entry_added", entry, lane: this.name }, + ...entryLifecycleEvents( + { ...entry, seq: commit.seqs[index]!, timestamp: commit.timestamp }, + this.name, + operationId, + ), ); } if (captured.length !== 0) { @@ -685,22 +677,22 @@ export class Lane implements AgentLane { context, ) ).reverse(); - const queuedItems = async (ids: readonly string[], description: string): Promise => - Promise.all( - ids.map(async (entryId) => { - const stored = await reader.getValue(pendingEntry(entryId), context); - if (stored?.value.type !== "message") { - throw new SessionInvariantError(`${description} ${entryId} is missing a message payload`); - } - return { entryId, message: stored.value.payload }; - }), - ); - const nextRun = await queuedItems(captured.pendingNextRun, "Pending next-run entry"); + const nextRun = await readPendingMessages( + reader, + captured.pendingNextRun, + "Pending next-run entry", + context, + ); const operation = captured.operation; const runState = operation !== null && isRunOperationState(operation.state) ? operation.state : undefined; - const steer = runState === undefined ? [] : await queuedItems(runState.inbox.steer, "Steer entry"); + const steer = + runState === undefined + ? [] + : await readPendingMessages(reader, runState.inbox.steer, "Steer entry", context); const followUp = - runState === undefined ? [] : await queuedItems(runState.inbox.followUp, "Follow-up entry"); + runState === undefined + ? [] + : await readPendingMessages(reader, runState.inbox.followUp, "Follow-up entry", context); // TODO does a client really need pending writes? We wouldn't visualize those, any other uses for them? const pendingWrites = runState !== undefined @@ -728,24 +720,10 @@ export class Lane implements AgentLane { let streamingMessage: NonNullable["streamingMessage"]; let retry: NonNullable["retry"]; let deferred: NonNullable["deferred"]; - // TODO we should always read the full frame list here, assistant messages - // are never humongous - const readStreamingMessage = async (responseEntryId: string) => { - const frames = []; - let cursor: { seq: number } | undefined; - while (frames.length < 10_000) { - const limit = Math.min(1_000, 10_000 - frames.length); - const page = await reader.readList( - pendingAssistantFrames(operation.meta.operationId, responseEntryId), - { order: "asc", limit, ...(cursor === undefined ? {} : { cursor }) }, - context, - ); - frames.push(...page.map(({ value }) => value)); - if (page.length < limit) break; - cursor = { seq: page[page.length - 1]!.seq }; - } - return reduceAssistantMessageFrames(frames); - }; + const readStreamingMessage = async (responseEntryId: string) => + reduceAssistantMessageFrames( + await readAssistantFrames(reader, operation.meta.operationId, responseEntryId, context), + ); const state = operation.state; switch (state.at) { case "run.assistant.retry_wait": @@ -826,10 +804,17 @@ export class Lane implements AgentLane { const drained = operation.state.control.status === "cancel_requested" ? { - steer: await queuedItems(operation.state.control.drainedSteer, "Drained steer entry"), - followUp: await queuedItems( + steer: await readPendingMessages( + reader, + operation.state.control.drainedSteer, + "Drained steer entry", + context, + ), + followUp: await readPendingMessages( + reader, operation.state.control.drainedFollowUp, "Drained follow-up entry", + context, ), } : undefined; diff --git a/packages/agent/src/harness/runtime/progress.ts b/packages/agent/src/harness/runtime/progress.ts index 23eff85d9..5c6858ac2 100644 --- a/packages/agent/src/harness/runtime/progress.ts +++ b/packages/agent/src/harness/runtime/progress.ts @@ -1,14 +1,8 @@ import type { AssistantMessageFrame } from "@earendil-works/pi-ai"; import type { AgentToolResult } from "../../types.ts"; -import type { Write } from "../session/types.ts"; -import { - appendList, - deleteList, - deleteValue, - pendingAssistantFrames, - pendingToolOutput, - setValue, -} from "../session/values.ts"; +import type { Context } from "../context.ts"; +import type { SessionReader, Write } from "../session/types.ts"; +import { appendList, pendingAssistantFrames, pendingToolOutput, setValue } from "../session/values.ts"; import type { Lane } from "./lane.ts"; import type { Drive, LaneState } from "./types.ts"; @@ -16,14 +10,32 @@ export interface ProgressChannel { write(item: T): void; seal(): void; drain(): Promise; - clearWrite(): Write; +} + +export async function readAssistantFrames( + reader: SessionReader, + operationId: string, + responseEntryId: string, + context: Context, +): Promise { + const frames: AssistantMessageFrame[] = []; + let cursor: { seq: number } | undefined; + for (;;) { + const page = await reader.readList( + pendingAssistantFrames(operationId, responseEntryId), + { order: "asc", limit: 1_000, ...(cursor === undefined ? {} : { cursor }) }, + context, + ); + frames.push(...page.map(({ value }) => value)); + if (page.length < 1_000) return frames; + cursor = { seq: page[page.length - 1]!.seq }; + } } function openProgress( lane: Lane, drive: Drive, commitWrite: (item: T) => Write, - clear: () => Write, stillOwns: (state: LaneState) => boolean, ): ProgressChannel { let sealed = false; @@ -51,7 +63,6 @@ function openProgress( async drain() { await latest; }, - clearWrite: clear, }; } @@ -65,7 +76,6 @@ export function openFrameProgress( lane, drive, (frame) => appendList(address, frame), - () => deleteList(address), (state) => { const run = state.operation?.state; if (run === undefined) return false; @@ -89,7 +99,6 @@ export function openToolProgress( lane, drive, (snapshot) => setValue(address, snapshot), - () => deleteValue(address), (state) => { const operation = state.operation; if (operation?.state.at !== "run.tools") return false; diff --git a/packages/agent/src/harness/runtime/restore.ts b/packages/agent/src/harness/runtime/restore.ts index 658b60b8d..92168e04b 100644 --- a/packages/agent/src/harness/runtime/restore.ts +++ b/packages/agent/src/harness/runtime/restore.ts @@ -1,6 +1,12 @@ import type { Context } from "../context.ts"; import { SessionInvariantError } from "../session/session.ts"; -import type { Session, SessionReader } from "../session/types.ts"; +import type { + LaneState as DurableLaneState, + LaneConfiguration, + LaneLastResult, + Session, + SessionReader, +} from "../session/types.ts"; import { branchTip, branchTipInventoryPrefix, @@ -9,34 +15,86 @@ import { laneState as laneStateValue, operationMeta, operationState, + type StoredValue, } from "../session/values.ts"; import type { LaneState } from "./types.ts"; +type LaneValues = { + tip: StoredValue | undefined; + configuration: StoredValue | undefined; + laneState: StoredValue | undefined; + lastResult: StoredValue | undefined; +}; + +export type ClassifiedLaneStorage = + | { kind: "absent" } + | { kind: "branch"; tip: StoredValue } + | { + kind: "lane"; + tip: StoredValue; + configuration: StoredValue; + laneState: StoredValue; + lastResult: StoredValue | undefined; + }; + +function classifyLaneStorage(lane: string, values: LaneValues): ClassifiedLaneStorage { + const { tip, configuration, laneState, lastResult } = values; + if (tip === undefined && configuration === undefined && laneState === undefined && lastResult === undefined) { + return { kind: "absent" }; + } + if (tip !== undefined && configuration === undefined && laneState === undefined && lastResult === undefined) { + return { kind: "branch", tip }; + } + if (tip === undefined) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing branch.tip`); + if (configuration === undefined) + throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.config`); + if (laneState === undefined) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.state`); + return { kind: "lane", tip, configuration, laneState, lastResult }; +} + +export async function readLaneStorage( + reader: SessionReader, + lane: string, + context: Context, +): Promise { + const [tip, configuration, laneState, lastResult] = await Promise.all([ + reader.getValue(branchTip(lane), context), + reader.getValue(laneConfig(lane), context), + reader.getValue(laneStateValue(lane), context), + reader.getValue(laneLastResult(lane), context), + ]); + return classifyLaneStorage(lane, { tip, configuration, laneState, lastResult }); +} + /** Restore every complete configured AgentLane in one coherent Session read. */ export function restoreSession(session: Session, context: Context): Promise> { return session.mutate(async (reader) => { - const inventories = await Promise.all([ + const [tips, configurations, states, lastResults] = await Promise.all([ reader.scanValues(branchTipInventoryPrefix(), context), reader.scanValues(laneConfig(""), context), reader.scanValues(laneStateValue(""), context), reader.scanValues(laneLastResult(""), context), ]); - const names = new Set(inventories.flatMap((values) => values.map(({ address }) => address.key))); + const tipByLane = new Map(tips.map((value) => [value.address.key, value])); + const configurationByLane = new Map(configurations.map((value) => [value.address.key, value])); + const stateByLane = new Map(states.map((value) => [value.address.key, value])); + const lastResultByLane = new Map(lastResults.map((value) => [value.address.key, value])); + const names = new Set([ + ...tipByLane.keys(), + ...configurationByLane.keys(), + ...stateByLane.keys(), + ...lastResultByLane.keys(), + ]); const restored = new Map(); for (const lane of names) { - const [configuration, state, lastResult] = await Promise.all([ - reader.getValue(laneConfig(lane), context), - reader.getValue(laneStateValue(lane), context), - reader.getValue(laneLastResult(lane), context), - ]); - const tip = await reader.getValue(branchTip(lane), context); - if (tip !== undefined && configuration === undefined && state === undefined && lastResult === undefined) { - continue; - } - if (tip === undefined || configuration === undefined || state === undefined) { - throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} has incomplete durable state`); - } - restored.set(lane, await restoreLaneState(reader, lane, context)); + const stored = classifyLaneStorage(lane, { + tip: tipByLane.get(lane), + configuration: configurationByLane.get(lane), + laneState: stateByLane.get(lane), + lastResult: lastResultByLane.get(lane), + }); + if (stored.kind !== "lane") continue; + restored.set(lane, await restoreLaneState(reader, lane, stored, context)); } return restored; }, context); @@ -44,22 +102,25 @@ export function restoreSession(session: Session, context: Context): Promise { - return session.mutate((reader) => restoreLaneState(reader, lane, context), context); + return session.mutate(async (reader) => { + const stored = await readLaneStorage(reader, lane, context); + if (stored.kind === "absent") { + throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing branch.tip`); + } + if (stored.kind === "branch") { + throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.config`); + } + return restoreLaneState(reader, lane, stored, context); + }, context); } -export async function restoreLaneState(reader: SessionReader, lane: string, context: Context): Promise { - const [tip, configuration, laneState, lastResult] = await Promise.all([ - reader.getValue(branchTip(lane), context), - reader.getValue(laneConfig(lane), context), - reader.getValue(laneStateValue(lane), context), - reader.getValue(laneLastResult(lane), context), - ]); - if (tip === undefined) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing branch.tip`); - if (configuration === undefined) - throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.config`); - if (laneState === undefined) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.state`); - - const operationId = laneState.value.currentOperationId; +export async function restoreLaneState( + reader: SessionReader, + lane: string, + stored: Extract, + context: Context, +): Promise { + const operationId = stored.laneState.value.currentOperationId; let operation: LaneState["operation"] = null; if (operationId !== null) { const [meta, state] = await Promise.all([ @@ -87,10 +148,10 @@ export async function restoreLaneState(reader: SessionReader, lane: string, cont } return { - tipId: tip.value, - configuration: configuration.value, - pendingNextRun: laneState.value.pendingNextRun, - ...(lastResult === undefined ? {} : { lastResult: lastResult.value }), + tipId: stored.tip.value, + configuration: stored.configuration.value, + pendingNextRun: stored.laneState.value.pendingNextRun, + ...(stored.lastResult === undefined ? {} : { lastResult: stored.lastResult.value }), operation, }; } diff --git a/packages/agent/src/harness/runtime/transcript.ts b/packages/agent/src/harness/runtime/transcript.ts new file mode 100644 index 000000000..b64a7b140 --- /dev/null +++ b/packages/agent/src/harness/runtime/transcript.ts @@ -0,0 +1,44 @@ +import type { AgentMessage } from "../../types.ts"; +import type { HarnessEvent } from "../agent-harness.ts"; +import type { Context } from "../context.ts"; +import { SessionInvariantError } from "../session/session.ts"; +import type { Entry, SessionReader } from "../session/types.ts"; +import { pendingEntry } from "../session/values.ts"; + +export function chainEntries( + parentId: string | null, + items: readonly T[], +): Array { + return items.map((item) => { + const entry = { ...item, parentId }; + parentId = item.id; + return entry; + }); +} + +export function entryLifecycleEvents(entry: Entry, lane: string, runId: string): HarnessEvent[] { + return entry.type === "message" + ? [ + { type: "message_start", lane, runId, message: entry.message }, + { type: "message_end", lane, runId, message: entry.message, entryId: entry.id }, + { type: "entry_added", lane, entry }, + ] + : [{ type: "entry_added", lane, entry }]; +} + +export function readPendingMessages( + reader: SessionReader, + ids: readonly string[], + description: string, + context: Context, +): Promise> { + return Promise.all( + ids.map(async (entryId) => { + const value = await reader.getValue(pendingEntry(entryId), context); + if (value?.value.type !== "message") { + throw new SessionInvariantError(`${description} ${entryId} is missing its message payload`); + } + return { entryId, message: value.value.payload }; + }), + ); +} diff --git a/packages/agent/src/harness/runtime/types.ts b/packages/agent/src/harness/runtime/types.ts index 4dd082192..4bee6a94e 100644 --- a/packages/agent/src/harness/runtime/types.ts +++ b/packages/agent/src/harness/runtime/types.ts @@ -68,8 +68,10 @@ export type LaneCommand = | { kind: "return"; result: TResult } | { kind: "reject"; error: Error }; +export type ContinueOperationResult = { kind: "cancel_requested" } | { kind: "result"; value: TResult }; + /** A durable operation transition. The Lane pairs the state write with projection publication. */ -type LanePatch = Partial>; +type LanePatch = Partial>; interface FinishDecision { kind: "finish"; @@ -83,8 +85,7 @@ interface FinishDecision { export type OperationCommand = | (CommitDecision & { operationState: OperationState; lane?: LanePatch }) | FinishDecision - | { kind: "return"; result: TResult } - | { kind: "reject"; error: Error }; + | { kind: "return"; result: TResult }; /** One installed process-local drive pass. */ export class Drive { @@ -98,7 +99,6 @@ export class Drive { private readonly control: GateControl; private readonly resolveCompletion: (outcome: DriveOutcome) => void; private readonly rejectCompletion: (error: unknown) => void; - private settled = false; constructor(options: DriveOptions, context: Context) { this.operationId = options.operationId; @@ -120,14 +120,10 @@ export class Drive { } settle(outcome: DriveOutcome): void { - if (this.settled) return; - this.settled = true; this.resolveCompletion(outcome); } fail(error: unknown): void { - if (this.settled) return; - this.settled = true; this.rejectCompletion(error); } diff --git a/packages/agent/src/harness/session/types.ts b/packages/agent/src/harness/session/types.ts index 4ae3adef6..c3d984768 100644 --- a/packages/agent/src/harness/session/types.ts +++ b/packages/agent/src/harness/session/types.ts @@ -416,17 +416,6 @@ export function runScopeOf(state: RunOperationState): RunScope { }; } -/** Copy only the checkpoint payload fields from a checkpoint-bearing state. */ -export function checkpointDataOf(state: CheckpointData): CheckpointData { - return { - continuation: state.continuation, - triggerEntryId: state.triggerEntryId, - ...(state.thresholdCheckedTriggerEntryId === undefined - ? {} - : { thresholdCheckedTriggerEntryId: state.thresholdCheckedTriggerEntryId }), - ...(state.skipInboxOnce === undefined ? {} : { skipInboxOnce: state.skipInboxOnce }), - }; -} export type Operation = { meta: OperationMeta; state: OperationState }; export interface LaneState { diff --git a/packages/agent/test/harness/runtime/drive-generation.test.ts b/packages/agent/test/harness/runtime/drive-generation.test.ts index c24b032b6..cbf57304f 100644 --- a/packages/agent/test/harness/runtime/drive-generation.test.ts +++ b/packages/agent/test/harness/runtime/drive-generation.test.ts @@ -27,6 +27,7 @@ import { type MessageEntry, type RunAssistantEffectPendingOperation, type RunAssistantReadyOperation, + type RunCheckpointOperation, type RunOperationState, runScopeOf, type Session, @@ -267,6 +268,53 @@ describe("runtime generation checkpoint", () => { expect(drained.inbox.writes).toEqual([]); await expectProjectionRestores(fixture); }); + + it("preserves checkpoint routing when a custom-only pending write does not project", async () => { + const fixture = await createFixture(); + await startFixtureRun(fixture); + const current = currentRun(fixture.lane); + if (current.at !== "run.checkpoint") throw new Error("missing checkpoint"); + const customId = "01950000-0000-7000-8000-000000000012"; + const checkpoint: RunCheckpointOperation = { + ...current, + continuation: { kind: "may_finish", includeFinalAssistant: false }, + thresholdCheckedTriggerEntryId: current.triggerEntryId, + skipInboxOnce: false, + inbox: { ...current.inbox, writes: [customId] }, + }; + await fixture.lane.command((state) => { + const operation = state.operation; + if (operation === null) throw new Error("missing operation"); + return { + kind: "commit", + writes: [ + storedValues.setValue(storedValues.pendingEntry(customId), { + type: "custom", + customType: "display-only", + payload: { value: true }, + }), + storedValues.setValue(storedValues.operationState(operationId), checkpoint), + ], + next: { ...state, operation: { meta: operation.meta, state: checkpoint } }, + materialize: () => undefined, + }; + }, BACKGROUND_CONTEXT); + + await runCheckpoint(fixture.lane, fixture.drive, checkpoint); + + const drained = currentRun(fixture.lane); + if (drained.at !== "run.checkpoint") throw new Error("missing drained checkpoint"); + expect(fixture.lane.state.tipId).toBe(customId); + expect(drained).toMatchObject({ + continuation: checkpoint.continuation, + triggerEntryId: checkpoint.triggerEntryId, + thresholdCheckedTriggerEntryId: checkpoint.triggerEntryId, + skipInboxOnce: false, + inbox: { writes: [] }, + }); + expect(await fixture.session.getValue(storedValues.pendingEntry(customId), BACKGROUND_CONTEXT)).toBeUndefined(); + await expectProjectionRestores(fixture); + }); }); describe("runtime assistant generation", () => { diff --git a/packages/agent/test/harness/runtime/drive-retry-deferred.test.ts b/packages/agent/test/harness/runtime/drive-retry-deferred.test.ts index 12e9c6243..9b0860d43 100644 --- a/packages/agent/test/harness/runtime/drive-retry-deferred.test.ts +++ b/packages/agent/test/harness/runtime/drive-retry-deferred.test.ts @@ -37,6 +37,7 @@ import { type Write, } from "../../../src/harness/session/types.ts"; import * as storedValues from "../../../src/harness/session/values.ts"; +import { deferred } from "./test-utils.ts"; const sessions: Session[] = []; const operationId = "01950000-0000-7000-8000-000000000001"; @@ -396,6 +397,45 @@ describe("runtime deferred polling", () => { await expectProjectionRestores(fixture); }); + it("declines poll intent when cancellation wins preparation", async () => { + const fixture = await createFixture(); + const suspended = await submitDeferred(fixture); + const drive = installDrive(fixture, { pollDeferred: true }); + const hookStarted = deferred(); + const releaseHook = deferred(); + fixture.hooks.on("before_request", async (event) => { + if (event.step !== "deferred") return undefined; + hookStarted.resolve(); + await releaseHook.promise; + return undefined; + }); + + const polling = runDeferredSuspended(fixture.lane, drive, suspended); + await hookStarted.promise; + await fixture.lane.command((state) => { + const operation = state.operation; + if (operation === null || !isRunOperationState(operation.state)) throw new Error("missing run"); + const nextState: RunOperationState = { + ...operation.state, + control: { status: "cancel_requested", requestedAt: 10, drainedSteer: [], drainedFollowUp: [] }, + }; + return { + kind: "commit", + writes: [storedValues.setValue(storedValues.operationState(operationId), nextState)], + next: { ...state, operation: { meta: operation.meta, state: nextState } }, + materialize: () => undefined, + }; + }, BACKGROUND_CONTEXT); + releaseHook.resolve(); + + expect(await polling).toEqual({ kind: "continue" }); + expect(fixture.faux.state.deferredFetchCount).toBe(0); + expect(currentRun(fixture)).toMatchObject({ + at: "run.deferred.suspended", + control: { status: "cancel_requested" }, + }); + }); + it("plans ready deferred tool calls with the poll turn identity and follower result ids", async () => { const fixture = await createFixture(); const suspended = await submitDeferred( diff --git a/packages/agent/test/harness/runtime/lane.test.ts b/packages/agent/test/harness/runtime/lane.test.ts index c0fa58012..1d3e6e31e 100644 --- a/packages/agent/test/harness/runtime/lane.test.ts +++ b/packages/agent/test/harness/runtime/lane.test.ts @@ -9,7 +9,12 @@ import { convertToLlm } from "../../../src/harness/messages.ts"; import { Lane } from "../../../src/harness/runtime/lane.ts"; import { restoreLane } from "../../../src/harness/runtime/restore.ts"; import { StorageBackedSession } from "../../../src/harness/session/session.ts"; -import type { LaneConfiguration, Session } from "../../../src/harness/session/types.ts"; +import type { + LaneConfiguration, + OperationState, + RunStartingOperation, + Session, +} from "../../../src/harness/session/types.ts"; import * as storedValues from "../../../src/harness/session/values.ts"; import { ControlledMemoryStorage, deferred } from "./test-utils.ts"; @@ -303,6 +308,68 @@ describe("runtime Lane commands", () => { ); }); + it("diverts ordinary work but settles against latest cancelled control", async () => { + const { lane, session } = await createLane(); + const accepted = await lane.accept({ kind: "prompt", prompt: "hello" }, BACKGROUND_CONTEXT); + expect(accepted.ok).toBe(true); + const operation = lane.state.operation; + if (operation?.state.at !== "run.starting") throw new Error("missing accepted operation"); + const cancelled: RunStartingOperation = { + ...operation.state, + control: { + status: "cancel_requested", + requestedAt: 1, + drainedSteer: [], + drainedFollowUp: [], + }, + }; + await lane.command( + (state) => ({ + kind: "commit", + writes: [storedValues.setValue(storedValues.operationState(operation.meta.operationId), cancelled)], + next: { ...state, operation: { meta: operation.meta, state: cancelled } }, + materialize: () => undefined, + }), + BACKGROUND_CONTEXT, + ); + + let continued = false; + expect( + await lane.continueOperation( + cancelled, + () => { + continued = true; + return { kind: "return", result: "continued" }; + }, + BACKGROUND_CONTEXT, + ), + ).toEqual({ kind: "cancel_requested" }); + expect(continued).toBe(false); + + const settled = await lane.settleOperation( + cancelled, + (_state, current) => { + expect(current.control.status).toBe("cancel_requested"); + const operationState: OperationState = { + ...current, + inbox: { ...current.inbox, writes: ["accepted-during-cancellation"] }, + }; + return { + kind: "commit", + writes: [], + operationState, + materialize: () => operationState, + }; + }, + BACKGROUND_CONTEXT, + ); + expect(settled.inbox.writes).toEqual(["accepted-during-cancellation"]); + expect(lane.state.operation?.state).toEqual(settled); + expect( + (await session.getValue(storedValues.operationState(operation.meta.operationId), BACKGROUND_CONTEXT))?.value, + ).toEqual(settled); + }); + it("plans queued commands from the latest committed memory", async () => { const { lane, storage } = await createLane(); const commitStarted = deferred(); diff --git a/packages/agent/test/harness/runtime/progress.test.ts b/packages/agent/test/harness/runtime/progress.test.ts index aa1292e1c..d29cbfc4a 100644 --- a/packages/agent/test/harness/runtime/progress.test.ts +++ b/packages/agent/test/harness/runtime/progress.test.ts @@ -215,9 +215,6 @@ describe("runtime progress channels", () => { ) ).map(({ value }) => value), ).toEqual(frames); - expect(progress.clearWrite()).toEqual( - storedValues.deleteList(storedValues.pendingAssistantFrames(operationId, responseEntryId)), - ); }); it("declines a queued frame after the authoritative projection leaves its phase without control reads", async () => { diff --git a/packages/agent/test/harness/runtime/restore.test.ts b/packages/agent/test/harness/runtime/restore.test.ts index 05861c5ed..ed96c7b40 100644 --- a/packages/agent/test/harness/runtime/restore.test.ts +++ b/packages/agent/test/harness/runtime/restore.test.ts @@ -262,6 +262,7 @@ describe("runtime lane restore", () => { workerState: await session.getValue(storedValues.laneState("worker"), BACKGROUND_CONTEXT), }; const mutate = vi.spyOn(session, "mutate"); + const getValue = vi.spyOn(MemoryStorage.prototype, "getValue"); const readList = vi.spyOn(MemoryStorage.prototype, "readList"); const lanes = await restoreSession(session, BACKGROUND_CONTEXT); @@ -273,7 +274,12 @@ describe("runtime lane restore", () => { operation: { meta, state }, }); expect(mutate).toHaveBeenCalledTimes(1); + expect(getValue.mock.calls.map(([address]) => ({ namespace: address.namespace, key: address.key }))).toEqual([ + { namespace: storedValues.operationMeta("").namespace, key: operationId }, + { namespace: storedValues.operationState("").namespace, key: operationId }, + ]); expect(readList).not.toHaveBeenCalled(); + getValue.mockRestore(); expect({ leaves: await session.scanValues(storedValues.branchTipInventoryPrefix(), BACKGROUND_CONTEXT), mainConfiguration: await session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT), @@ -295,6 +301,6 @@ describe("runtime lane restore", () => { (mutator) => mutator.commit([storedValues.deleteValue(storedValues.branchTip("main"))], BACKGROUND_CONTEXT), BACKGROUND_CONTEXT, ); - await expect(restoreSession(session, BACKGROUND_CONTEXT)).rejects.toThrow("incomplete durable state"); + await expect(restoreSession(session, BACKGROUND_CONTEXT)).rejects.toThrow('Lane "main" is missing branch.tip'); }); });