diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 8c3a0ab5c..acd6cfbb7 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -6,11 +6,13 @@ - Replaced the unfinished harness runtime and record-log session model with the shared entry/register/usage, operation-state, session, search, and public harness contracts specified for the durable harness. - Added the required replay-stable invocation metadata argument to `AgentHarnessTool.execute()`. +- Required storage backends to return exact post-commit `SessionStats` in `CommitResult.stats`. ### Added - Added scheduler-facing harness execution primitives (`accept`, fenced `drive`, `requestAbort`, and `inspectExecution`) plus replay-stable tool invocation identities. - Added isolated harness execution primitives for assistant streaming and durable-boundary tool phases. +- Added direct-drive foundation primitives for split effect admission, deterministic commit gating, invocation-fenced progress writes, and injectable session id generation. - Added durable `AgentHarness` prompt acceptance and minimal no-tool execution with recovery, lifecycle events, usage settlement, and terminal cleanup. - Added assistant request recovery, durable retry waits, deferred suspension, and missing-identity suspension to `AgentHarness`. - Added durable tool planning, context-bound sequential and parallel execution, ordered result settlement, and safe tool replay recovery to `AgentHarness`. diff --git a/packages/agent/docs/work-packages/05-direct-durable-drive.md b/packages/agent/docs/work-packages/05-direct-durable-drive.md index a91bbbd13..3749b4fc1 100644 --- a/packages/agent/docs/work-packages/05-direct-durable-drive.md +++ b/packages/agent/docs/work-packages/05-direct-durable-drive.md @@ -1284,13 +1284,13 @@ export interface CommitResult { } ``` -The seed is free: `StorageState` already folds `messageCount` and `usage` in `applyValidated` and is rebuilt from the log on open, so a reopened session's first commit reports totals that include all history. +For the deliberately materialized Memory and JSONL backends, `InMemoryStorageState` already folds `messageCount` and `usage` in `applyValidated` and is rebuilt from the log on JSONL open. It is not an architecture for database backends or long-running sessions that may not fit in memory. SQLite maintains and reads indexed durable aggregates inside its commit transaction. | File | Change | |---|---| | `src/harness/session/types.ts` | add `stats: SessionStats` to `CommitResult` | | `src/harness/session/commit.ts` | `PreparedCommit.result` becomes `Omit`; `prepareStorageCommit` is unchanged otherwise (it runs before apply and cannot know the totals) | -| `src/harness/session/storage-state.ts` | `applyValidated(writes)` returns the post-apply `SessionStats` (it already computes it into `this.stats`) | +| `src/harness/session/in-memory-storage-state.ts` | `InMemoryStorageState.applyValidated(writes)` returns the post-apply `SessionStats` for the deliberately materialized Memory and JSONL backends | | `src/harness/session/memory.ts` | `const stats = this.storageState.applyValidated(prepared.writes); return { ...prepared.result, stats };` | | `src/harness/session/jsonl/storage.ts` | same composition inside `applyCommit` | | `packages/session-backends/sqlite-node/src/sqlite/storage.ts` | **third backend — do not miss it.** `applyCommit` already runs inside `this.db.transaction(…)` and already maintains stats there via `incrementMessageCount` / `addUsageToSessionStats`. Read them back with `readSessionStats(this.db)` **inside the same transaction closure**, after `advanceNextSeq`, so the snapshot is atomic with the writes that produced it (see below) | @@ -1443,7 +1443,7 @@ Only `packages/coding-agent/src/core/cache-stats.ts` (comment) and vendored `hig **Read completely before editing** - `src/harness/execution/effect-gate.ts`, `src/harness/hooks.ts` - `src/harness/runtime2/{lane,types,harness,restore}.ts` -- `src/harness/session/{values,types,session,commit,storage-state,memory}.ts` +- `src/harness/session/{values,types,session,commit,in-memory-storage-state,memory}.ts` - `src/harness/session/jsonl/storage.ts` - `src/harness/session/testing/{storage-decorator,instrumented-storage,index}.ts` - `src/harness/session/testing/conformance/storage.ts` @@ -1468,7 +1468,7 @@ Only `packages/coding-agent/src/core/cache-stats.ts` (comment) and vendored `hig | `src/harness/runtime2/restore.ts` | name `Lane` wherever it names `Lane` | | `src/harness/session/types.ts` | add `stats: SessionStats` to `CommitResult` (§2.12) | | `src/harness/session/commit.ts` | `PreparedCommit.result` becomes `Omit` | -| `src/harness/session/storage-state.ts` | `applyValidated` returns the post-apply `SessionStats` | +| `src/harness/session/in-memory-storage-state.ts` | `InMemoryStorageState.applyValidated` returns the post-apply `SessionStats` for Memory and JSONL | | `src/harness/session/memory.ts`, `src/harness/session/jsonl/storage.ts` | compose `{ ...prepared.result, stats }` | | `packages/session-backends/sqlite-node/src/sqlite/storage.ts` | return `{ ...prepared.result, stats: readSessionStats(this.db) }` from **inside** the `db.transaction` closure (§2.12) | | `packages/session-backends/sqlite-node/test/storage-conformance.test.ts`, `packages/session-backends/sqlite-node/test/storage.test.ts` | SQLite conformance plus a post-write-snapshot assertion | @@ -1991,7 +1991,7 @@ Milestone in parentheses where creation order matters. - `packages/agent/src/harness/session/session.ts` - `packages/agent/src/harness/session/types.ts` - `packages/agent/src/harness/session/commit.ts` -- `packages/agent/src/harness/session/storage-state.ts` +- `packages/agent/src/harness/session/in-memory-storage-state.ts` - `packages/agent/src/harness/session/memory.ts` - `packages/agent/src/harness/session/jsonl/storage.ts` - `packages/agent/src/harness/session/testing/conformance/storage.ts` diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 3eed41715..240bcbf36 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -16,7 +16,7 @@ import type { BranchPreparation, BranchSummaryResult } from "./compaction/branch import type { CompactionPreparation, CompactionSettings, CompactResult } from "./compaction/compaction.ts"; import type { Context } from "./context.ts"; import { type Result, TaggedError } from "./result.ts"; -import { createAgentHarness } from "./runtime2/index.ts"; +import { createAgentHarness } from "./runtime2/harness.ts"; import type { BranchSummaryEntry, CompactionEntry, @@ -640,4 +640,4 @@ export interface AgentHarnessConstructor { } /** Runtime constructor for attaching the durable harness to one open session. */ -export const AgentHarness: AgentHarnessConstructor = { create: createAgentHarness }; +export const AgentHarness = { create: createAgentHarness } satisfies AgentHarnessConstructor; diff --git a/packages/agent/src/harness/execution/effect-gate.ts b/packages/agent/src/harness/execution/effect-gate.ts index 789558a0a..e722f4b49 100644 --- a/packages/agent/src/harness/execution/effect-gate.ts +++ b/packages/agent/src/harness/execution/effect-gate.ts @@ -9,53 +9,64 @@ export class AbortRequested extends Error { } } -type EffectGateState = +/** Invocation cancellation or owner abandonment. Never durable cancellation. */ +export class DriveAbandoned extends Error { + constructor() { + super("Drive pass was abandoned by its installing invocation"); + this.name = "DriveAbandoned"; + } +} + +/** Procedure-facing synchronous admission capability for one drive pass. */ +export interface Gate { + readonly signal: AbortSignal; + admit(invoke: () => T): T; +} + +/** Owner-facing lifecycle controls for one drive pass. */ +export interface GateControl { + beginAbort(cancellation: Promise): void; + signalAbort(): void; + close(error: Error): void; +} + +type GateState = | { status: "open" } | { status: "aborting"; cancellation: Promise } | { status: "closed"; error: Error }; -/** Process-local admission gate for one operation's external work. */ -export interface EffectGate { - /** The operation-owned cooperative signal. */ - readonly signal: AbortSignal; - /** Synchronously throws when ordinary work may no longer start. */ - assertOpen(): void; - /** Close ordinary starts before the durable cancellation mutation begins. */ - beginAbort(cancellation: Promise): void; - /** Pull the cooperative signal only after cancellation is durable. */ - signalAbort(): void; - /** Permanently stop starts and signal already-admitted work. */ - close(error: Error): void; -} - -/** Default {@link EffectGate} implementation used by drive passes. */ -export class OperationEffectGate implements EffectGate { - private readonly controller = new AbortController(); - private state: EffectGateState = { status: "open" }; - - get signal(): AbortSignal { - return this.controller.signal; - } - - assertOpen(): void { - if (this.state.status === "aborting") throw new AbortRequested(this.state.cancellation); - if (this.state.status === "closed") throw this.state.error; - } - - beginAbort(cancellation: Promise): void { - if (this.state.status === "open") { - this.state = { status: "aborting", cancellation }; - } - } - - signalAbort(): void { - if (this.state.status !== "aborting" || this.controller.signal.aborted) return; - this.controller.abort(new AbortRequested(this.state.cancellation)); - } - - close(error: Error): void { - if (this.state.status === "closed") return; - this.state = { status: "closed", error }; - if (!this.controller.signal.aborted) this.controller.abort(error); - } +/** Create separate procedure-facing and owner-facing views of one effect gate. */ +export function createGate(): { gate: Gate; control: GateControl } { + let state: GateState = { status: "open" }; + const controller = new AbortController(); + + const check = (): void => { + if (state.status === "aborting") throw new AbortRequested(state.cancellation); + if (state.status === "closed") throw state.error; + }; + + return { + gate: { + admit(invoke: () => T): T { + check(); + return invoke(); + }, + signal: controller.signal, + }, + control: { + beginAbort(cancellation) { + if (state.status !== "open") return; + state = { status: "aborting", cancellation }; + }, + signalAbort() { + if (state.status !== "aborting" || controller.signal.aborted) return; + controller.abort(new AbortRequested(state.cancellation)); + }, + close(error) { + if (state.status === "closed") return; + state = { status: "closed", error }; + if (!controller.signal.aborted) controller.abort(error); + }, + }, + }; } diff --git a/packages/agent/src/harness/execution/tools.ts b/packages/agent/src/harness/execution/tools.ts index e2087f7ff..480490a51 100644 --- a/packages/agent/src/harness/execution/tools.ts +++ b/packages/agent/src/harness/execution/tools.ts @@ -2,7 +2,7 @@ import { type ToolResultMessage, validateToolArguments } from "@earendil-works/p import type { AgentTool, AgentToolCall, AgentToolResult } from "../../types.ts"; import { type Context, withAbortSignal } from "../context.ts"; import type { JsonValue } from "../session/types.ts"; -import type { EffectGate } from "./effect-gate.ts"; +import type { Gate } from "./effect-gate.ts"; /** A tool call whose tool exists and whose prepared arguments passed validation. */ export interface PreparedToolCall { @@ -118,29 +118,30 @@ export function applyBeforeToolDecision( } /** Execute one cleared external tool effect, converting expected tool throws to error output. */ -export async function executeToolCall( +export function executeToolCall( call: ClearedToolCall, - effectGate: EffectGate, + gate: Gate, onUpdate: (partial: AgentToolResult) => void, context: Context, ): Promise { let acceptingUpdates = true; - effectGate.assertOpen(); - const admittedContext = withAbortSignal(effectGate.signal, context); - admittedContext.abortSignal?.throwIfAborted(); - try { - const result = await call.tool.execute(call.toolCall.id, call.args, admittedContext.abortSignal, (partial) => { - if (acceptingUpdates) onUpdate(partial); - }); - return { result, isError: false }; - } catch (error) { - return { - result: createErrorToolResult(error instanceof Error ? error.message : String(error)), - isError: true, - }; - } finally { - acceptingUpdates = false; - } + return gate.admit(async () => { + const admittedContext = withAbortSignal(gate.signal, context); + admittedContext.abortSignal?.throwIfAborted(); + try { + const result = await call.tool.execute(call.toolCall.id, call.args, admittedContext.abortSignal, (partial) => { + if (acceptingUpdates) onUpdate(partial); + }); + return { result, isError: false }; + } catch (error) { + return { + result: createErrorToolResult(error instanceof Error ? error.message : String(error)), + isError: true, + }; + } finally { + acceptingUpdates = false; + } + }); } /** Apply an after-tool patch field by field. */ diff --git a/packages/agent/src/harness/hooks.ts b/packages/agent/src/harness/hooks.ts index a7bed67e2..1a6d18f41 100644 --- a/packages/agent/src/harness/hooks.ts +++ b/packages/agent/src/harness/hooks.ts @@ -1,6 +1,6 @@ import type { HookHandler, HookInvocation, HookMap, HookName, Hooks } from "./agent-harness.ts"; import { type Context, withAbortSignal } from "./context.ts"; -import type { EffectGate } from "./execution/effect-gate.ts"; +import type { Gate } from "./execution/effect-gate.ts"; import { startHarnessSpan } from "./telemetry.ts"; import type { AgentHarnessStreamOptions, AgentHarnessStreamOptionsPatch } from "./types.ts"; @@ -11,13 +11,6 @@ interface HookRegistration { type HookErrorReporter = (error: Error, hook: HookName, lane: string, context: Context) => void | Promise; -function admitEffect(effectGate: EffectGate, context: Context): Context { - effectGate.assertOpen(); - const admittedContext = withAbortSignal(effectGate.signal, context); - admittedContext.abortSignal?.throwIfAborted(); - return admittedContext; -} - /** Ordered harness hook registry and aggregate runner. */ export class HookRegistry implements Hooks { private readonly registrations = new Map(); @@ -51,25 +44,32 @@ export class HookRegistry implements Hooks { runWithGate( name: TName, event: HookInvocation, - effectGate: EffectGate, + gate: Gate, context: Context, ): Promise { - return this.runAdmitted(name, event, admitEffect(effectGate, context)); + return gate.admit(() => { + const admittedContext = withAbortSignal(gate.signal, context); + admittedContext.abortSignal?.throwIfAborted(); + return this.runAdmitted(name, event, admittedContext); + }); } /** Invoke a tool-hook aggregate with one telemetry span per registered handler. */ runToolWithGate( name: TName, event: HookInvocation, - effectGate: EffectGate, + gate: Gate, context: Context, ): Promise { - const admittedContext = admitEffect(effectGate, context); - return ( - name === "before_tool" - ? this.beforeTool(event as HookInvocation<"before_tool">, admittedContext) - : this.afterTool(event as HookInvocation<"after_tool">, admittedContext) - ) as Promise; + return gate.admit(() => { + const admittedContext = withAbortSignal(gate.signal, context); + admittedContext.abortSignal?.throwIfAborted(); + return ( + name === "before_tool" + ? this.beforeTool(event as HookInvocation<"before_tool">, admittedContext) + : this.afterTool(event as HookInvocation<"after_tool">, admittedContext) + ) as Promise; + }); } close(error: Error): void { diff --git a/packages/agent/src/harness/runtime2/harness.ts b/packages/agent/src/harness/runtime2/harness.ts index 3caa50f15..ff5dbb052 100644 --- a/packages/agent/src/harness/runtime2/harness.ts +++ b/packages/agent/src/harness/runtime2/harness.ts @@ -34,11 +34,10 @@ import { type Config, type LaneState, SliceNotImplemented } from "./types.ts"; type GlobalConfigProperty = GlobalConfigEventPayload["property"]; /** Runtime2 implementation of AgentHarness. The harness is the main lane. */ -export class Harness extends Lane implements AgentHarness { - readonly hooks: HookRegistry; +export class Harness extends Lane implements AgentHarness { readonly events: HarnessEventBus; readonly seed: LaneConfiguration; - readonly lanesByName = new Map(); + readonly lanesByName = new Map>(); private readonly configStore: { value: Config }; closePromise: Promise | undefined; faultError: HarnessFault | undefined; @@ -46,6 +45,20 @@ export class Harness extends Lane implement constructor(options: AgentHarnessOptions, seed: LaneConfiguration, restored: Map) { const main = restored.get("main"); if (main === undefined) throw new SessionInvariantError("Session is missing main lane"); + const events = new HarnessEventBus(); + const hooks = new HookRegistry((error, hook, lane, context) => + events.emit( + { + type: "handler_error", + kind: "hook", + hook, + error: error.message, + ...(error.stack === undefined ? {} : { stack: error.stack }), + lane, + }, + context, + ), + ); const configStore: { value: Config } = { value: { tools: options.tools ?? [], @@ -66,6 +79,7 @@ export class Harness extends Lane implement "main", options.session, options.models, + hooks, main, (cause, context) => this.fault(cause, context), (events, context) => this.events.emitBatch(events, context), @@ -74,20 +88,7 @@ export class Harness extends Lane implement ); this.seed = seed; this.configStore = configStore; - this.events = new HarnessEventBus(); - this.hooks = new HookRegistry((error, hook, lane, context) => - this.events.emit( - { - type: "handler_error", - kind: "hook", - hook, - error: error.message, - ...(error.stack === undefined ? {} : { stack: error.stack }), - lane, - }, - context, - ), - ); + this.events = events; this.lanesByName.set("main", this); for (const [name, state] of restored) { if (name !== "main") this.lanesByName.set(name, this.buildLane(name, state)); @@ -215,11 +216,12 @@ export class Harness extends Lane implement throw new SliceNotImplemented("watchSession"); } - private buildLane(name: string, state: LaneState): Lane { - return new Lane( + private buildLane(name: string, state: LaneState): Lane { + return new Lane( name, this.session, this.models, + this.hooks, state, (cause, context) => this.fault(cause, context), (events, context) => this.events.emitBatch(events, context), @@ -265,7 +267,7 @@ export class Harness extends Lane implement /* * Close is a process-local admission boundary, not an abort. Public lane work checks Lane.assertOpen() at * entry. Durable work then enters Session.mutate(), whose mutation line is the final write-admission gate. - * External effects check their EffectGate immediately before invocation. Closing seals all applicable gates + * External effects use Gate.admit() immediately around invocation. Closing seals all applicable gates * before awaiting anything. * * If close wins, new work and queued mutation callbacks that have not started reject. An active mutation may diff --git a/packages/agent/src/harness/runtime2/lane.ts b/packages/agent/src/harness/runtime2/lane.ts index dfda3ff4a..a7bba33dc 100644 --- a/packages/agent/src/harness/runtime2/lane.ts +++ b/packages/agent/src/harness/runtime2/lane.ts @@ -18,6 +18,7 @@ import { type LaneSnapshot, type ModelIdentity, type OperationAdmissionResult, + OperationMismatch, type OperationRequest, type QueuedItem, UnknownSkill, @@ -25,13 +26,13 @@ import { type WatchHandle, } from "../agent-harness.ts"; import type { Context } from "../context.ts"; +import type { HookRegistry } from "../hooks.ts"; import { formatPromptTemplateInvocation } from "../prompt-templates.ts"; import { Result } from "../result.ts"; import { insertEntry } from "../session/commit.ts"; import { SessionInvariantError, SessionPendingAssistantMessageError } from "../session/session.ts"; import type { BranchScan, - CommitResult, Entry, LaneLastResult, Operation, @@ -40,7 +41,6 @@ import type { SessionReader, SessionTree, StructuralDecision, - Write, } from "../session/types.ts"; import { deleteValue, @@ -58,23 +58,19 @@ import { setValue, } from "../session/values.ts"; import { formatSkillInvocation } from "../skills.ts"; -import { type AcceptanceConfig, type LaneState, SliceNotImplemented } from "./types.ts"; +import { + type ActiveDrive, + type Config, + type DriveScope, + type LaneCommand, + type LaneDriveCapabilities, + type LaneState, + SliceNotImplemented, +} from "./types.ts"; type EmitBatch = (events: readonly HarnessEvent[], context: Context) => Promise; type WatchHandler = (snapshot: T, filter: (event: HarnessEvent) => boolean, context: Context) => WatchHandle; type FaultHandler = (cause: unknown, context: Context) => Error; -type Synchronous = TResult extends PromiseLike ? never : TResult; - -type LaneCommand = - | { - kind: "commit"; - writes: Write[]; - next: LaneState; - materialize(commit: CommitResult): Synchronous; - events?(commit: CommitResult): HarnessEvent[]; - } - | { kind: "return"; result: TResult } - | { kind: "reject"; error: Error }; type LaneCommandOutcome = | { kind: "return"; result: TResult; delivery?: Promise } @@ -112,16 +108,18 @@ function capturedModel(operation: Operation): ModelIdentity | undefined { } /** Runtime2 implementation of one configured lane. */ -export class Lane implements AgentLane { +export class Lane implements AgentLane, LaneDriveCapabilities { readonly name: string; readonly sessionTree: SessionTree; + readonly session: Session; + readonly models: Models; + readonly hooks: HookRegistry; + readonly emitBatch: EmitBatch; private readonly sessionView: SessionTree; - protected readonly session: Session; - protected readonly models: Models; private readonly onFault: FaultHandler; - private readonly emitBatch: EmitBatch; private readonly installWatch: WatchHandler; - private readonly getAcceptanceConfig: () => AcceptanceConfig; + private readonly config: () => Config; + private activeDrive: ActiveDrive | undefined; state: LaneState; closedError: Error | undefined; @@ -129,14 +127,16 @@ export class Lane implements AgentLane { name: string, session: Session, models: Models, + hooks: HookRegistry, state: LaneState, onFault: FaultHandler, emitBatch: EmitBatch, installWatch: WatchHandler, - getAcceptanceConfig: () => AcceptanceConfig, + readConfig: () => Config, ) { this.session = session; this.models = models; + this.hooks = hooks; this.name = name; this.sessionView = session.view(name); this.sessionTree = { @@ -154,7 +154,7 @@ export class Lane implements AgentLane { this.onFault = onFault; this.emitBatch = emitBatch; this.installWatch = installWatch; - this.getAcceptanceConfig = getAcceptanceConfig; + this.config = readConfig; } async getLeafId(_context: Context): Promise { @@ -167,6 +167,24 @@ export class Lane implements AgentLane { return this.state.lastResult; } + readConfig(): Config { + return this.config(); + } + + ownsDrive(scope: DriveScope): boolean { + return this.activeDrive === scope.owner; + } + + mismatch(expected: string, currentOperationId: string | null, last: LaneLastResult | undefined): OperationMismatch { + return new OperationMismatch({ + lane: this.name, + expectedOperationId: expected, + ...(currentOperationId === null ? {} : { currentOperationId }), + ...(last === undefined ? {} : { lastOperationId: last.operationId }), + message: `Operation ${expected} does not own lane ${JSON.stringify(this.name)}`, + }); + } + /** * Run one effect-free command on this lane's serialized mutation line. Owned `state` is authoritative; the * planner receives that state plus a read-only reader for bounded payload lookups. Committed values and owned @@ -185,11 +203,7 @@ export class Lane implements AgentLane { * timers, event handlers, or wait for task completion here; perform those after `command()` returns. */ async command( - plan: ( - state: LaneState, - reader: SessionReader, - context: Context, - ) => LaneCommand | Promise>, + plan: (state: LaneState, reader: SessionReader) => LaneCommand | Promise>, context: Context, ): Promise { this.assertOpen(); @@ -200,7 +214,7 @@ export class Lane implements AgentLane { async (mutator) => { this.assertOpen(); try { - const decision = await plan(this.state, mutator, context); + const decision = await plan(this.state, mutator); switch (decision.kind) { case "return": return { kind: "return", result: decision.result }; @@ -245,7 +259,7 @@ export class Lane implements AgentLane { const startedAt = Date.now(); const operationId = request.operationId ?? this.session.idGenerator.next(startedAt); - const acceptanceConfig = this.getAcceptanceConfig(); + const acceptanceConfig = this.readConfig(); let messages: AgentMessage[]; switch (request.kind) { case "prompt": diff --git a/packages/agent/src/harness/runtime2/progress.ts b/packages/agent/src/harness/runtime2/progress.ts new file mode 100644 index 000000000..0c2e9d78e --- /dev/null +++ b/packages/agent/src/harness/runtime2/progress.ts @@ -0,0 +1,119 @@ +import type { AssistantMessageFrame } from "@earendil-works/pi-ai"; +import type { AgentToolResult } from "../../types.ts"; +import { SessionInvariantError } from "../session/session.ts"; +import type { RunState, SessionReader, Write } from "../session/types.ts"; +import { + appendList, + deleteList, + deleteValue, + laneState as laneStateValue, + operationState as operationStateValue, + pendingAssistantFrames, + pendingToolOutput, + setValue, +} from "../session/values.ts"; +import type { DriveScope } from "./types.ts"; + +export interface ProgressChannel { + write(item: T): void; + seal(): void; + drain(): Promise; + clearWrite(): Write; +} + +function openProgress( + scope: DriveScope, + commitWrite: (item: T) => Write, + clear: () => Write, + stillOwns: (reader: SessionReader) => Promise, +): ProgressChannel { + let sealed = false; + let latest: Promise = Promise.resolve(); + return { + write(item) { + if (sealed) return; + const write = scope.lane.command(async (projection, reader) => { + if (!(await stillOwns(reader))) return { kind: "return", result: undefined }; + return { + kind: "commit", + writes: [commitWrite(item)], + next: projection, + materialize: () => undefined, + }; + }, scope.context); + latest = write; + void write.catch(() => {}); + }, + seal() { + sealed = true; + }, + async drain() { + await latest; + }, + clearWrite: clear, + }; +} + +async function readOwnedRun( + scope: DriveScope, + reader: SessionReader, +): Promise { + const durableLane = await reader.getValue(laneStateValue(scope.lane.name), scope.context); + if (durableLane === undefined) { + throw new SessionInvariantError(`Lane ${JSON.stringify(scope.lane.name)} has no lane state`); + } + if (durableLane.value.currentOperationId !== scope.operationId) return undefined; + const durableState = await reader.getValue(operationStateValue(scope.operationId), scope.context); + if (durableState === undefined) { + throw new SessionInvariantError(`Operation ${scope.operationId} has no operation state`); + } + if (durableState.value.kind !== "run") { + throw new SessionInvariantError(`Operation ${scope.operationId} is ${durableState.value.kind}, expected run`); + } + if (!scope.lane.ownsDrive(scope)) return undefined; + return durableState.value; +} + +export function openFrameProgress( + scope: DriveScope, + responseEntryId: string, +): ProgressChannel { + const address = pendingAssistantFrames(scope.operationId, responseEntryId); + return openProgress( + scope, + (frame) => appendList(address, frame), + () => deleteList(address), + async (reader) => { + const state = await readOwnedRun(scope, reader); + if (state === undefined) return false; + if (state.phase.kind === "assistant") { + const generation = state.phase.generation; + return generation.status === "effect_pending" && generation.responseEntryId === responseEntryId; + } + if (state.phase.kind === "deferred") { + const deferred = state.phase.deferred; + return deferred.status === "effect_pending" && deferred.responseEntryId === responseEntryId; + } + return false; + }, + ); +} + +export function openToolProgress( + scope: DriveScope, + invocationId: string, +): ProgressChannel> { + const address = pendingToolOutput(scope.operationId, invocationId); + return openProgress( + scope, + (snapshot) => setValue(address, snapshot), + () => deleteValue(address), + async (reader) => { + const state = await readOwnedRun(scope, reader); + if (state?.phase.kind !== "tools") return false; + return state.phase.batch.calls.some( + (call) => call.resultEntryId === invocationId && call.status === "effect_pending", + ); + }, + ); +} diff --git a/packages/agent/src/harness/runtime2/types.ts b/packages/agent/src/harness/runtime2/types.ts index 9ef572fd4..10c343bf0 100644 --- a/packages/agent/src/harness/runtime2/types.ts +++ b/packages/agent/src/harness/runtime2/types.ts @@ -1,8 +1,26 @@ -import type { RetryPolicy } from "@earendil-works/pi-ai"; +import type { Models, RetryPolicy } from "@earendil-works/pi-ai"; import type { QueueMode } from "../../types.ts"; -import type { AgentHarnessOptions, Resources } from "../agent-harness.ts"; +import type { + AgentHarnessOptions, + DriveOutcome, + HarnessEvent, + OperationMismatch, + Resources, + TerminalOperationOutcome, +} from "../agent-harness.ts"; import type { CompactionSettings } from "../compaction/compaction.ts"; -import type { LaneConfiguration, LaneLastResult, Operation } from "../session/types.ts"; +import type { Context } from "../context.ts"; +import { createGate, type Gate, type GateControl } from "../execution/effect-gate.ts"; +import type { HookRegistry } from "../hooks.ts"; +import type { + CommitResult, + LaneConfiguration, + LaneLastResult, + Operation, + Session, + SessionReader, + Write, +} from "../session/types.ts"; import type { AgentHarnessStreamOptions, AgentHarnessTool } from "../types.ts"; export class SliceNotImplemented extends Error { @@ -45,3 +63,89 @@ export interface LaneState { readonly lastResult?: LaneLastResult; readonly operation: Operation | null; } + +type Synchronous = TResult extends PromiseLike ? never : TResult; + +/** One effect-free decision made on a lane's serialized mutation line. */ +export type LaneCommand = + | { + kind: "commit"; + writes: Write[]; + next: LaneState; + materialize(commit: CommitResult): Synchronous; + events?(commit: CommitResult): HarnessEvent[]; + } + | { kind: "return"; result: TResult } + | { kind: "reject"; error: Error }; + +export type CompletionResult = { kind: "ok"; outcome: DriveOutcome } | { kind: "error"; error: unknown }; + +export interface DriveSupervisor { + settleCompletion(result: CompletionResult): void; + retainOwner(): void; + removeOwner(): void; +} + +/** Process-local ownership record for one drive pass. */ +export interface ActiveDrive { + readonly operationId: string; + readonly completion: Promise; + readonly gate: Gate; + readonly control: GateControl; + finalizedOutcome?: DriveOutcome; + supervisor?: DriveSupervisor; +} + +export interface ActiveDriveHandle { + readonly active: ActiveDrive; + settle(outcome: DriveOutcome): void; + fail(error: unknown): void; +} + +export function createActiveDrive(operationId: string): ActiveDriveHandle { + let settle!: (outcome: DriveOutcome) => void; + let fail!: (error: unknown) => void; + const completion = new Promise((resolve, reject) => { + settle = resolve; + fail = reject; + }); + void completion.catch(() => {}); + const { gate, control } = createGate(); + return { active: { operationId, completion, gate, control }, settle, fail }; +} + +export interface PassPolicy { + waitForRetry: boolean; + deferredPermits: number; +} + +export interface DriveScope { + readonly lane: LaneDriveCapabilities; + readonly operationId: string; + readonly owner: ActiveDrive; + readonly context: Context; + readonly pass: PassPolicy; +} + +export type StepResult = + | { kind: "advance" } + | { kind: "reload" } + | { kind: "waiting"; outcome: DriveOutcome } + | { kind: "settled"; outcome: TerminalOperationOutcome } + | { kind: "lost_ownership" }; + +/** Package-internal capability surface available to direct drive procedures. */ +export interface LaneDriveCapabilities { + readonly name: string; + readonly session: Session; + readonly models: Models; + readonly hooks: HookRegistry; + command( + plan: (projection: LaneState, reader: SessionReader) => LaneCommand | Promise>, + context: Context, + ): Promise; + emitBatch(events: readonly HarnessEvent[], context: Context): Promise; + readConfig(): Config; + ownsDrive(scope: DriveScope): boolean; + mismatch(expected: string, currentOperationId: string | null, last: LaneLastResult | undefined): OperationMismatch; +} diff --git a/packages/agent/src/harness/session/commit.ts b/packages/agent/src/harness/session/commit.ts index 453d4ebb5..f6e94552c 100644 --- a/packages/agent/src/harness/session/commit.ts +++ b/packages/agent/src/harness/session/commit.ts @@ -42,7 +42,7 @@ export type CommittedWrite = export interface PreparedCommit { writes: CommittedWrite[]; - result: CommitResult; + result: Omit; } export interface CommitValidationState { diff --git a/packages/agent/src/harness/session/storage-state.ts b/packages/agent/src/harness/session/in-memory-storage-state.ts similarity index 94% rename from packages/agent/src/harness/session/storage-state.ts rename to packages/agent/src/harness/session/in-memory-storage-state.ts index c5c465a63..ae1b1737e 100644 --- a/packages/agent/src/harness/session/storage-state.ts +++ b/packages/agent/src/harness/session/in-memory-storage-state.ts @@ -53,8 +53,13 @@ function compareKeys(left: string, right: string): number { return leftCodePoints.length - rightCodePoints.length; } -/** Current in-memory projection shared by storage backends. */ -export class StorageState { +/** + * Complete materialized session state for MemoryStorage and JsonlStorage. + * + * This is intentionally unsuitable for database backends and long-running sessions that may not fit in memory. + * Those backends should query indexed durable state and update durable aggregates within each commit transaction. + */ +export class InMemoryStorageState { private readonly entries: Map; private readonly entriesBySeq: Entry[]; private readonly scalarValues: Map>; @@ -86,8 +91,8 @@ export class StorageState { }); } - /** Apply writes already accepted by validateCommitted(). */ - applyValidated(writes: readonly CommittedWrite[]): void { + /** Apply writes already accepted by validateCommitted() and return the post-apply totals. */ + applyValidated(writes: readonly CommittedWrite[]): SessionStats { for (const write of writes) { switch (write.kind) { case "entry": { @@ -137,6 +142,7 @@ export class StorageState { } this.nextSeq = write.seq + 1; } + return this.stats; } advanceNextSeq(nextSeq: number): void { diff --git a/packages/agent/src/harness/session/index.ts b/packages/agent/src/harness/session/index.ts index b6dc491e0..e883222a4 100644 --- a/packages/agent/src/harness/session/index.ts +++ b/packages/agent/src/harness/session/index.ts @@ -28,6 +28,7 @@ export { SessionPendingAssistantMessageError, SessionUnknownTargetError, StorageBackedSession, + type StorageBackedSessionOptions, } from "./session.ts"; export * from "./types.ts"; export * from "./values.ts"; diff --git a/packages/agent/src/harness/session/jsonl/storage.ts b/packages/agent/src/harness/session/jsonl/storage.ts index b7e558845..8fe4201c0 100644 --- a/packages/agent/src/harness/session/jsonl/storage.ts +++ b/packages/agent/src/harness/session/jsonl/storage.ts @@ -11,8 +11,8 @@ import { type CommittedValueDeleteWrite, type CommittedValueSetWrite, type CommittedWrite, - StorageState, -} from "../storage-state.ts"; + InMemoryStorageState, +} from "../in-memory-storage-state.ts"; import type { CommitResult, Entry, @@ -127,7 +127,7 @@ export class JsonlStorage implements Storage { private readonly now: () => number; readonly header: JsonlStorageHeader; private backing: JsonlBacking; - private readonly storageState = new StorageState(); + private readonly storageState = new InMemoryStorageState(); private commitQueue: Promise = Promise.resolve(); private state: "open" | "closing" | "closed" = "open"; private closePromise: Promise | undefined; @@ -256,8 +256,8 @@ export class JsonlStorage implements Storage { `Failed to append JSONL storage ${this.path}`, ); } - this.storageState.applyValidated(prepared.writes); - return prepared.result; + const stats = this.storageState.applyValidated(prepared.writes); + return { ...prepared.result, stats: this.withImportedUsage(stats) }; } /** Atomically upgrade legacy v3 backing and preserve the first caller write as a v4 transaction. */ @@ -345,12 +345,11 @@ export class JsonlStorage implements Storage { getStats(_context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed")); - const stats = this.storageState.getStats(); - if (this.backing.kind === "v4") return Promise.resolve(stats); - return Promise.resolve({ - ...stats, - usage: this.backing.importedUsage, - }); + return Promise.resolve(this.withImportedUsage(this.storageState.getStats())); + } + + private withImportedUsage(stats: SessionStats): SessionStats { + return this.backing.kind === "v4" ? stats : { ...stats, usage: this.backing.importedUsage }; } /** Capture the state needed to fork at one serialized boundary between commits. */ diff --git a/packages/agent/src/harness/session/memory.ts b/packages/agent/src/harness/session/memory.ts index 162698dfe..52fb9f641 100644 --- a/packages/agent/src/harness/session/memory.ts +++ b/packages/agent/src/harness/session/memory.ts @@ -7,8 +7,8 @@ import { type ForkSourceSnapshot, forkSnapshotWrites, } from "./fork.ts"; +import { InMemoryStorageState } from "./in-memory-storage-state.ts"; import { StorageBackedSession } from "./session.ts"; -import { StorageState } from "./storage-state.ts"; import type { BranchScan, CommitResult, @@ -54,7 +54,7 @@ export interface MemorySessionRepoOptions { export class MemoryStorage implements Storage { private readonly now: () => number; - private storageState = new StorageState(); + private storageState = new InMemoryStorageState(); private commitQueue: Promise = Promise.resolve(); private state: "open" | "closing" | "closed" = "open"; private closePromise: Promise | undefined; @@ -67,8 +67,8 @@ export class MemoryStorage implements Storage { if (this.state !== "open") throw new Error("MemoryStorage is closed"); const result = this.commitQueue.then(() => { const prepared = this.storageState.prepareCommit(writes, this.now()); - this.storageState.applyValidated(prepared.writes); - return prepared.result; + const stats = this.storageState.applyValidated(prepared.writes); + return { ...prepared.result, stats }; }); this.commitQueue = result.then( () => undefined, diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index 53c14de9b..7fb2b849a 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -44,8 +44,9 @@ import { type ValueList, } from "./values.ts"; -interface StorageBackedSessionOptions { +export interface StorageBackedSessionOptions { laneMutationLine?: LaneMutationLine; + idGenerator?: IdGenerator; onClose?: () => void; } @@ -182,7 +183,7 @@ class StorageBackedSessionMutator implements SessionMutator { /** Package-internal typed boundary shared by concrete session repositories. */ export class StorageBackedSession implements Session { readonly metadata: TMetadata; - readonly idGenerator: IdGenerator = { next: uuidv7 }; + readonly idGenerator: IdGenerator; private readonly storage: Storage; private readonly laneMutationLine: LaneMutationLine; private readonly onClose: (() => void) | undefined; @@ -192,6 +193,7 @@ export class StorageBackedSession { + deepStrictEqual(result.stats, await storage.getStats(BACKGROUND_CONTEXT)); +} + /** Creates fresh, runner-independent cases for the durable Storage contract. */ export function createStorageConformance(factory: () => Promise): readonly ConformanceCase[] { return [ @@ -139,6 +145,7 @@ export function createStorageConformance(factory: () => Promise) strictEqual(result.seqs.length, 3); strictEqual(result.firstSeq, result.seqs[0]); + await assertCommitStats(storage, result); assertStrictlyIncreasing(result.seqs); ok(Number.isSafeInteger(result.timestamp) && result.timestamp >= 0); deepStrictEqual( @@ -822,17 +829,18 @@ export function createStorageConformance(factory: () => Promise) deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), { messageCount: 0, usage: zeroUsage() }); const firstUsage = usage(2, 3, { cacheWrite1h: 4, reasoning: 1 }); - await storage.commit( + const first = await storage.commit( [ insertEntry(userEntry("message")), insertUsage({ id: "usage-1", usage: firstUsage, adjustment: false }), ], BACKGROUND_CONTEXT, ); - deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), { messageCount: 1, usage: firstUsage }); + await assertCommitStats(storage, first); + deepStrictEqual(first.stats, { messageCount: 1, usage: firstUsage }); const secondUsage = usage(5, 7, { cacheWrite1h: 6, reasoning: 2 }); - await storage.commit( + const second = await storage.commit( [ insertEntry(customEntry("custom", "message")), insertEntry(compactionEntry("compaction", "custom")), @@ -840,7 +848,8 @@ export function createStorageConformance(factory: () => Promise) ], BACKGROUND_CONTEXT, ); - deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), { + await assertCommitStats(storage, second); + deepStrictEqual(second.stats, { messageCount: 1, usage: { input: 7, @@ -872,6 +881,9 @@ export function createStorageConformance(factory: () => Promise) const [firstResult, secondResult] = await Promise.all([first, second]); ok(firstResult.seqs[0]! < secondResult.seqs[0]!); + deepStrictEqual(firstResult.stats, { messageCount: 1, usage: zeroUsage() }); + deepStrictEqual(secondResult.stats, { messageCount: 2, usage: zeroUsage() }); + await assertCommitStats(storage, secondResult); deepStrictEqual(ids(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT)), ["first", "second"]); }, ), diff --git a/packages/agent/src/harness/session/testing/gating-storage.ts b/packages/agent/src/harness/session/testing/gating-storage.ts new file mode 100644 index 000000000..eaad63b6d --- /dev/null +++ b/packages/agent/src/harness/session/testing/gating-storage.ts @@ -0,0 +1,114 @@ +import type { Context } from "../../context.ts"; +import type { CommitResult, Write } from "../types.ts"; +import { StorageDecorator } from "./storage-decorator.ts"; + +/** Thrown for every commit rejected after simulated storage loss. */ +export class CommitDiscarded extends Error { + constructor(message: string) { + super(message); + this.name = "CommitDiscarded"; + } +} + +interface ParkedCommit { + release(): void; + drop(error: Error): void; + readonly landing: Promise; +} + +interface PendingWaiter { + count: number; + resolve(): void; + reject(error: Error): void; +} + +/** Test-only storage decorator that deterministically parks admitted commits. */ +export class GatingStorage extends StorageDecorator { + private armed = false; + private discarded = false; + private readonly queue: ParkedCommit[] = []; + private readonly waiters: PendingWaiter[] = []; + + /** Fixture setup bypasses gating until explicitly armed. */ + arm(): void { + this.armed = true; + } + + pending(): number { + return this.queue.length; + } + + /** Wait until at least `count` commits are parked. */ + waitPending(count = 1): Promise { + if (!Number.isSafeInteger(count) || count < 1) { + return Promise.reject(new RangeError("Pending commit count must be a positive safe integer")); + } + if (this.discarded) return Promise.reject(new CommitDiscarded("storage discarded")); + if (this.queue.length >= count) return Promise.resolve(); + return new Promise((resolve, reject) => { + this.waiters.push({ count, resolve, reject }); + }); + } + + override async commit(writes: Write[], context: Context): Promise { + if (this.discarded) throw new CommitDiscarded("commit rejected: storage discarded"); + if (!this.armed) return super.commit(writes, context); + + let landed!: () => void; + let lost!: (error: Error) => void; + const landing = new Promise((resolve, reject) => { + landed = resolve; + lost = reject; + }); + void landing.catch(() => {}); + + const released = new Promise((resolve, reject) => { + this.queue.push({ release: resolve, drop: reject, landing }); + }); + this.notifyWaiters(); + + try { + await released; + if (this.discarded) throw new CommitDiscarded("commit rejected: storage discarded"); + const result = await super.commit(writes, context); + landed(); + return result; + } catch (error) { + const normalized = error instanceof Error ? error : new Error(String(error)); + lost(normalized); + throw error; + } + } + + /** Release `count` commits in FIFO order and wait until each write lands. */ + async next(count = 1): Promise { + if (!Number.isSafeInteger(count) || count < 1) { + throw new RangeError("Released commit count must be a positive safe integer"); + } + for (let index = 0; index < count; index++) { + await this.waitPending(); + const parked = this.queue.shift(); + if (parked === undefined) throw new Error("No parked commit"); + parked.release(); + await parked.landing; + } + } + + /** Drop parked commits and permanently reject every later commit. */ + discard(): void { + if (this.discarded) return; + this.discarded = true; + const error = new CommitDiscarded("commit discarded"); + for (const parked of this.queue.splice(0)) parked.drop(error); + for (const waiter of this.waiters.splice(0)) waiter.reject(error); + } + + private notifyWaiters(): void { + for (let index = this.waiters.length - 1; index >= 0; index--) { + const waiter = this.waiters[index]!; + if (this.queue.length < waiter.count) continue; + this.waiters.splice(index, 1); + waiter.resolve(); + } + } +} diff --git a/packages/agent/src/harness/session/testing/index.ts b/packages/agent/src/harness/session/testing/index.ts index e0e32dffb..5032b9b74 100644 --- a/packages/agent/src/harness/session/testing/index.ts +++ b/packages/agent/src/harness/session/testing/index.ts @@ -28,6 +28,7 @@ export { createSessionRepoOwnershipConformance, } from "./conformance/session-repo.ts"; export { createStorageConformance } from "./conformance/storage.ts"; +export { CommitDiscarded, GatingStorage } from "./gating-storage.ts"; export { InstrumentedStorage } from "./instrumented-storage.ts"; export { StorageDecorator } from "./storage-decorator.ts"; export type { ConformanceCase, StorageFixture } from "./types.ts"; diff --git a/packages/agent/src/harness/session/types.ts b/packages/agent/src/harness/session/types.ts index a2792cba8..4f135792a 100644 --- a/packages/agent/src/harness/session/types.ts +++ b/packages/agent/src/harness/session/types.ts @@ -392,6 +392,8 @@ export interface CommitResult { firstSeq: number; seqs: number[]; timestamp: number; + /** Session totals immediately after this commit was applied. */ + stats: SessionStats; } export interface EntryStructure { diff --git a/packages/agent/test/harness/execution-primitives.test.ts b/packages/agent/test/harness/execution-primitives.test.ts index 3e5398d54..0f94281c0 100644 --- a/packages/agent/test/harness/execution-primitives.test.ts +++ b/packages/agent/test/harness/execution-primitives.test.ts @@ -8,7 +8,7 @@ import { withTelemetryContext, } from "../../src/harness/context.ts"; import { HarnessEventBus } from "../../src/harness/events.ts"; -import { AbortRequested, OperationEffectGate } from "../../src/harness/execution/effect-gate.ts"; +import { AbortRequested, createGate, DriveAbandoned } from "../../src/harness/execution/effect-gate.ts"; import { HookRegistry } from "../../src/harness/hooks.ts"; import { InMemoryTelemetryContext } from "../../src/index.ts"; @@ -20,15 +20,15 @@ function deferred(): { promise: Promise; resolve(): void } { return { promise, resolve: () => resolvePromise?.() }; } -describe("OperationEffectGate", () => { +describe("createGate", () => { it("closes starts synchronously and signals only after cancellation commits", async () => { - const gate = new OperationEffectGate(); + const { gate, control } = createGate(); const cancellation = deferred(); - gate.beginAbort(cancellation.promise); + control.beginAbort(cancellation.promise); let refusal: unknown; try { - gate.assertOpen(); + gate.admit(() => undefined); } catch (error) { refusal = error; } @@ -37,17 +37,22 @@ describe("OperationEffectGate", () => { expect(gate.signal.aborted).toBe(false); cancellation.resolve(); await cancellation.promise; - gate.signalAbort(); + control.signalAbort(); expect(gate.signal.aborted).toBe(true); }); it("permanently closes and signals admitted work", () => { - const gate = new OperationEffectGate(); + const { gate, control } = createGate(); const error = new Error("closed"); - gate.close(error); - expect(() => gate.assertOpen()).toThrow(error); + control.close(error); + expect(() => gate.admit(() => undefined)).toThrow(error); expect(gate.signal.aborted).toBe(true); }); + + it("keeps invocation abandonment distinct from durable abort control flow", () => { + expect(new DriveAbandoned()).not.toBeInstanceOf(AbortRequested); + expect(new DriveAbandoned().name).toBe("DriveAbandoned"); + }); }); describe("HookRegistry", () => { @@ -71,7 +76,7 @@ describe("HookRegistry", () => { prompt: [{ role: "user", content: "prompt", timestamp: 1 }], resources: {}, }, - new OperationEffectGate(), + createGate().gate, BACKGROUND_CONTEXT, ); @@ -101,12 +106,12 @@ describe("HookRegistry", () => { }); const event = { lane: "main", runId: "run", prompt: [], resources: {} }; const closed = new Error("closed"); - const closedGate = new OperationEffectGate(); - closedGate.close(closed); - expect(() => hooks.runWithGate("before_run", event, closedGate, BACKGROUND_CONTEXT)).toThrow(closed); + const closedGate = createGate(); + closedGate.control.close(closed); + expect(() => hooks.runWithGate("before_run", event, closedGate.gate, BACKGROUND_CONTEXT)).toThrow(closed); expect(calls).toEqual([]); - const running = hooks.runWithGate("before_run", event, new OperationEffectGate(), BACKGROUND_CONTEXT); + const running = hooks.runWithGate("before_run", event, createGate().gate, BACKGROUND_CONTEXT); expect(calls).toEqual(["first:start"]); hooks.on("before_run", () => { calls.push("late"); @@ -123,7 +128,7 @@ describe("HookRegistry", () => { const preAbortedHooks = new HookRegistry(() => {}); const handler = vi.fn(); preAbortedHooks.on("before_drive", handler); - const preAbortedGate = new OperationEffectGate(); + const preAbortedGate = createGate(); const invocationController = new AbortController(); const abortReason = new Error("invocation aborted"); invocationController.abort(abortReason); @@ -132,24 +137,24 @@ describe("HookRegistry", () => { preAbortedHooks.runWithGate( "before_drive", event, - preAbortedGate, + preAbortedGate.gate, withAbortSignal(invocationController.signal, BACKGROUND_CONTEXT), ), ).toThrow(abortReason); expect(handler).not.toHaveBeenCalled(); - expect(() => preAbortedGate.assertOpen()).not.toThrow(); + expect(() => preAbortedGate.gate.admit(() => undefined)).not.toThrow(); const admittedHooks = new HookRegistry(() => {}); - const admittedGate = new OperationEffectGate(); + const admittedGate = createGate(); let admittedSignal: AbortSignal | undefined; admittedHooks.on("before_drive", (_event, context) => { admittedSignal = context.abortSignal; }); - await admittedHooks.runWithGate("before_drive", event, admittedGate, BACKGROUND_CONTEXT); + await admittedHooks.runWithGate("before_drive", event, admittedGate.gate, BACKGROUND_CONTEXT); expect(admittedSignal?.aborted).toBe(false); const gateReason = new Error("gate closed"); - admittedGate.close(gateReason); + admittedGate.control.close(gateReason); expect(admittedSignal?.aborted).toBe(true); expect(admittedSignal?.reason).toBe(gateReason); }); @@ -175,7 +180,7 @@ describe("HookRegistry", () => { toolName: "tool", args: {}, }, - new OperationEffectGate(), + createGate().gate, withContextValue(valueKey, "preserved", withTelemetryContext(invocationSpan, BACKGROUND_CONTEXT)), ), ); @@ -209,7 +214,7 @@ describe("HookRegistry", () => { hooks.runWithGate( "before_drive", { lane: "main", runId: "run", operation: "run" }, - new OperationEffectGate(), + createGate().gate, BACKGROUND_CONTEXT, ), ).rejects.toBe(failure); @@ -243,7 +248,7 @@ describe("HookRegistry", () => { messages: [{ role: "user", content: "original", timestamp: 1 }], systemPrompt: "base", }, - new OperationEffectGate(), + createGate().gate, BACKGROUND_CONTEXT, ); @@ -270,7 +275,7 @@ describe("HookRegistry", () => { attempt: 1, streamOptions: { headers: { a: "1", b: "2" }, metadata: { x: 1 } }, }, - new OperationEffectGate(), + createGate().gate, BACKGROUND_CONTEXT, ); expect(result).toEqual({ @@ -296,7 +301,7 @@ describe("HookRegistry", () => { content: [{ type: "text", text: "raw" }], isError: true, }, - new OperationEffectGate(), + createGate().gate, BACKGROUND_CONTEXT, ); expect(result).toEqual({ content: [{ type: "text", text: "patched" }], isError: false }); @@ -324,7 +329,7 @@ describe("HookRegistry", () => { totalTokens: 0, }, }, - new OperationEffectGate(), + createGate().gate, BACKGROUND_CONTEXT, ); expect(result).toEqual({ decline: false, summary: selected }); @@ -346,19 +351,19 @@ describe("HookRegistry", () => { }); const event = { operation: "run" as const, lane: "main", runId: "run" }; - const abortFirstGate = new OperationEffectGate(); + const abortFirstGate = createGate(); const cancellation = Promise.resolve(); - abortFirstGate.beginAbort(cancellation); - expect(() => hooks.runWithGate("before_drive", event, abortFirstGate, BACKGROUND_CONTEXT)).toThrow( + abortFirstGate.control.beginAbort(cancellation); + expect(() => hooks.runWithGate("before_drive", event, abortFirstGate.gate, BACKGROUND_CONTEXT)).toThrow( AbortRequested, ); expect(calls).toEqual([]); - const startFirstGate = new OperationEffectGate(); - const running = hooks.runWithGate("before_drive", event, startFirstGate, BACKGROUND_CONTEXT); + const startFirstGate = createGate(); + const running = hooks.runWithGate("before_drive", event, startFirstGate.gate, BACKGROUND_CONTEXT); expect(calls).toEqual(["first:start"]); - startFirstGate.beginAbort(cancellation); - startFirstGate.signalAbort(); + startFirstGate.control.beginAbort(cancellation); + startFirstGate.control.signalAbort(); release.resolve(); await running; expect(calls).toEqual(["first:start", "first:end", "second"]); diff --git a/packages/agent/test/harness/execution-tools.test.ts b/packages/agent/test/harness/execution-tools.test.ts index 7bba9c776..10f48f520 100644 --- a/packages/agent/test/harness/execution-tools.test.ts +++ b/packages/agent/test/harness/execution-tools.test.ts @@ -1,7 +1,7 @@ import { Type } from "typebox"; import { describe, expect, it, vi } from "vitest"; import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; -import { AbortRequested, OperationEffectGate } from "../../src/harness/execution/effect-gate.ts"; +import { AbortRequested, createGate, type Gate } from "../../src/harness/execution/effect-gate.ts"; import { applyBeforeToolDecision, createToolResultMessage, @@ -52,8 +52,8 @@ function clearPrepared(callToClear: ReturnType) { return cleared; } -function effectGate(): OperationEffectGate { - return new OperationEffectGate(); +function effectGate(): Gate { + return createGate().gate; } describe("tool execution primitives", () => { @@ -155,10 +155,10 @@ describe("tool execution primitives", () => { it("lets abort-first gate refusal escape without invoking the tool", async () => { const execute = vi.fn(tool().execute); const cleared = clearPrepared(prepareToolCall(call(), [tool({ execute })])); - const gate = effectGate(); - gate.beginAbort(Promise.resolve()); + const { gate, control } = createGate(); + control.beginAbort(Promise.resolve()); - await expect(executeToolCall(cleared, gate, () => {}, BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(AbortRequested); + expect(() => executeToolCall(cleared, gate, () => {}, BACKGROUND_CONTEXT)).toThrow(AbortRequested); expect(execute).not.toHaveBeenCalled(); }); diff --git a/packages/agent/test/harness/gating-storage.test.ts b/packages/agent/test/harness/gating-storage.test.ts new file mode 100644 index 000000000..08f6e4597 --- /dev/null +++ b/packages/agent/test/harness/gating-storage.test.ts @@ -0,0 +1,110 @@ +import { describe, expect, it } from "vitest"; +import { BACKGROUND_CONTEXT, type Context } from "../../src/harness/context.ts"; +import { MemoryStorage } from "../../src/harness/session/memory.ts"; +import { CommitDiscarded, GatingStorage, InstrumentedStorage } from "../../src/harness/session/testing/index.ts"; +import type { CommitResult, Write } from "../../src/harness/session/types.ts"; +import * as storedValues from "../../src/harness/session/values.ts"; +import { deferred } from "./runtime2/test-utils.ts"; + +class ControlledLandingStorage extends MemoryStorage { + readonly started = deferred(); + readonly release = deferred(); + + override async commit(writes: Write[], context: Context): Promise { + this.started.resolve(); + await this.release.promise; + return super.commit(writes, context); + } +} + +describe("GatingStorage", () => { + it("bypasses setup writes until armed and waits for commits parked after waitPending", async () => { + const storage = new GatingStorage(new MemoryStorage({ now: () => 10 })); + await storage.commit([storedValues.setValue(storedValues.sessionName, "setup")], BACKGROUND_CONTEXT); + expect(storage.pending()).toBe(0); + + storage.arm(); + let waitingResolved = false; + const waiting = storage.waitPending().then(() => { + waitingResolved = true; + }); + await Promise.resolve(); + expect(waitingResolved).toBe(false); + const commit = storage.commit([storedValues.setValue(storedValues.sessionName, "parked")], BACKGROUND_CONTEXT); + await waiting; + expect(storage.pending()).toBe(1); + expect((await storage.getValue(storedValues.sessionName, BACKGROUND_CONTEXT))?.value).toBe("setup"); + + await storage.next(); + await commit; + expect(storage.pending()).toBe(0); + expect((await storage.getValue(storedValues.sessionName, BACKGROUND_CONTEXT))?.value).toBe("parked"); + await storage.close(BACKGROUND_CONTEXT); + }); + + it("releases in FIFO order and next resolves only after the backend write lands", async () => { + const delegate = new ControlledLandingStorage({ now: () => 10 }); + const storage = new GatingStorage(delegate); + storage.arm(); + const first = storage.commit([storedValues.setValue(storedValues.sessionName, "first")], BACKGROUND_CONTEXT); + const second = storage.commit([storedValues.setValue(storedValues.sessionName, "second")], BACKGROUND_CONTEXT); + await storage.waitPending(2); + expect(storage.pending()).toBe(2); + + let released = false; + const next = storage.next().then(() => { + released = true; + }); + await delegate.started.promise; + await Promise.resolve(); + expect(released).toBe(false); + expect(storage.pending()).toBe(1); + delegate.release.resolve(); + await next; + await first; + expect(released).toBe(true); + expect((await storage.getValue(storedValues.sessionName, BACKGROUND_CONTEXT))?.value).toBe("first"); + + await storage.next(); + await second; + expect((await storage.getValue(storedValues.sessionName, BACKGROUND_CONTEXT))?.value).toBe("second"); + await storage.close(BACKGROUND_CONTEXT); + }); + + it("permanently rejects parked, waiting, and later commits after discard", async () => { + const storage = new GatingStorage(new MemoryStorage()); + storage.arm(); + const waiting = storage.waitPending(2); + const parked = storage.commit([storedValues.setValue(storedValues.sessionName, "lost")], BACKGROUND_CONTEXT); + await storage.waitPending(); + expect(storage.pending()).toBe(1); + + storage.discard(); + expect(storage.pending()).toBe(0); + await expect(parked).rejects.toBeInstanceOf(CommitDiscarded); + await expect(waiting).rejects.toBeInstanceOf(CommitDiscarded); + await expect(storage.next()).rejects.toBeInstanceOf(CommitDiscarded); + await expect(storage.commit([], BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(CommitDiscarded); + await storage.close(BACKGROUND_CONTEXT); + + const unarmed = new GatingStorage(new MemoryStorage()); + unarmed.discard(); + await expect(unarmed.commit([], BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(CommitDiscarded); + await unarmed.close(BACKGROUND_CONTEXT); + }); + + it("records attempts before gating parks them", async () => { + const gating = new GatingStorage(new MemoryStorage()); + const storage = new InstrumentedStorage(gating); + gating.arm(); + const writes = [storedValues.setValue(storedValues.sessionName, "recorded")]; + + const commit = storage.commit(writes, BACKGROUND_CONTEXT); + expect(storage.getCommitAttempts()).toEqual([writes]); + await gating.waitPending(); + expect(gating.pending()).toBe(1); + await gating.next(); + await commit; + await storage.close(BACKGROUND_CONTEXT); + }); +}); diff --git a/packages/agent/test/harness/instrumented-storage.test.ts b/packages/agent/test/harness/instrumented-storage.test.ts index c9eec4463..6561ce60b 100644 --- a/packages/agent/test/harness/instrumented-storage.test.ts +++ b/packages/agent/test/harness/instrumented-storage.test.ts @@ -3,9 +3,21 @@ import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import * as sessionWrites from "../../src/harness/session/commit.ts"; import { MemoryStorage } from "../../src/harness/session/memory.ts"; import { InstrumentedStorage } from "../../src/harness/session/testing/index.ts"; -import type { CommitResult, Write } from "../../src/harness/session/types.ts"; +import type { CommitResult, SessionStats, Write } from "../../src/harness/session/types.ts"; import * as storedValues from "../../src/harness/session/values.ts"; +const EMPTY_STATS: SessionStats = { + messageCount: 0, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, +}; + class ControlledCommitStorage extends MemoryStorage { private readonly pending: Array<{ resolve: (result: CommitResult) => void; @@ -55,11 +67,11 @@ describe("InstrumentedStorage", () => { expect(secondCommit).toBe(delegate.lastCommit); expect(storage.getCommitAttempts()).toEqual([firstTransaction, secondTransaction]); - const firstResult = { firstSeq: 1, seqs: [1], timestamp: 10 }; + const firstResult = { firstSeq: 1, seqs: [1], timestamp: 10, stats: EMPTY_STATS }; delegate.resolveNextCommit(firstResult); expect(await firstCommit).toBe(firstResult); expect(storage.getCommitAttempts()).toEqual([firstTransaction, secondTransaction]); - const secondResult = { firstSeq: 2, seqs: [2], timestamp: 20 }; + const secondResult = { firstSeq: 2, seqs: [2], timestamp: 20, stats: EMPTY_STATS }; delegate.resolveNextCommit(secondResult); expect(await secondCommit).toBe(secondResult); await storage.close(BACKGROUND_CONTEXT); @@ -71,7 +83,7 @@ describe("InstrumentedStorage", () => { const commit = storage.commit(transaction, BACKGROUND_CONTEXT); expect(storage.getCommitAttempts()[0]).toBe(transaction); - delegate.resolveNextCommit({ firstSeq: 1, seqs: [1], timestamp: 10 }); + delegate.resolveNextCommit({ firstSeq: 1, seqs: [1], timestamp: 10, stats: EMPTY_STATS }); await commit; await storage.close(BACKGROUND_CONTEXT); }); diff --git a/packages/agent/test/harness/jsonl-storage.test.ts b/packages/agent/test/harness/jsonl-storage.test.ts index 09b890c99..e7ae26fe5 100644 --- a/packages/agent/test/harness/jsonl-storage.test.ts +++ b/packages/agent/test/harness/jsonl-storage.test.ts @@ -197,7 +197,7 @@ describe("JsonlStorage persistence", () => { expect( (await reopened.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id, seq }) => ({ id, seq })), ).toEqual([{ id: "usage", seq: committed.seqs[2] }]); - expect(await reopened.getStats(BACKGROUND_CONTEXT)).toEqual({ + const historicalStats = { messageCount: 1, usage: { input: 1, @@ -207,9 +207,11 @@ describe("JsonlStorage persistence", () => { totalTokens: 3, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, - }); + }; + expect(await reopened.getStats(BACKGROUND_CONTEXT)).toEqual(historicalStats); const next = await reopened.commit([], BACKGROUND_CONTEXT); expect(next.firstSeq).toBe(5); + expect(next.stats).toEqual(historicalStats); await reopened.close(BACKGROUND_CONTEXT); }); }); diff --git a/packages/agent/test/harness/memory-conformance.test.ts b/packages/agent/test/harness/memory-conformance.test.ts index 27d9aa6ab..dc850c284 100644 --- a/packages/agent/test/harness/memory-conformance.test.ts +++ b/packages/agent/test/harness/memory-conformance.test.ts @@ -1,5 +1,6 @@ -import { describe, it } from "vitest"; +import { describe, expect, it } from "vitest"; import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; +import * as sessionWrites from "../../src/harness/session/commit.ts"; import { MemorySessionRepo } from "../../src/harness/session/index.ts"; import { MemoryStorage } from "../../src/harness/session/memory.ts"; import { @@ -8,6 +9,7 @@ import { createStorageConformance, type StorageFixture, } from "../../src/harness/session/testing/index.ts"; +import * as storedValues from "../../src/harness/session/values.ts"; const NOW = 1_700_000_000_000; @@ -34,6 +36,53 @@ registerConformance( }), ); +describe("MemoryStorage commit statistics", () => { + it("includes historical totals in the first commit after session reopen", async () => { + const repo = new MemorySessionRepo({ now: () => NOW }); + const session = await repo.create({}, BACKGROUND_CONTEXT); + const usage = { + input: 1, + output: 2, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 3, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + try { + await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + sessionWrites.insertEntry({ + id: "history", + parentId: null, + type: "message", + message: { role: "user", content: "history", timestamp: NOW }, + }), + sessionWrites.insertUsage({ id: "usage", usage, adjustment: false }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, + ); + await session.close(BACKGROUND_CONTEXT); + const reopened = await repo.open(session.metadata, BACKGROUND_CONTEXT); + const result = await reopened.mutate( + "main", + (mutator) => + mutator.commit([storedValues.setValue(storedValues.sessionName, "reopened")], BACKGROUND_CONTEXT), + BACKGROUND_CONTEXT, + ); + expect(result.stats).toEqual({ messageCount: 1, usage }); + expect(result.stats).toEqual(await reopened.getStats(BACKGROUND_CONTEXT)); + await reopened.close(BACKGROUND_CONTEXT); + } finally { + await repo.close(BACKGROUND_CONTEXT); + } + }); +}); + let memoryRepo: MemorySessionRepo; registerConformance( "MemorySessionRepo conformance", diff --git a/packages/agent/test/harness/runtime2/lane.test.ts b/packages/agent/test/harness/runtime2/lane.test.ts index 3f472be50..b358cbad7 100644 --- a/packages/agent/test/harness/runtime2/lane.test.ts +++ b/packages/agent/test/harness/runtime2/lane.test.ts @@ -4,6 +4,8 @@ import { HarnessClosed, type HarnessEvent } from "../../../src/harness/agent-har import { DEFAULT_COMPACTION_SETTINGS } from "../../../src/harness/compaction/compaction.ts"; import { BACKGROUND_CONTEXT, type Context } from "../../../src/harness/context.ts"; import { HarnessEventBus } from "../../../src/harness/events.ts"; +import { HookRegistry } from "../../../src/harness/hooks.ts"; +import { convertToLlm } from "../../../src/harness/messages.ts"; import { Lane } from "../../../src/harness/runtime2/lane.ts"; import { restoreLane } from "../../../src/harness/runtime2/restore.ts"; import { StorageBackedSession } from "../../../src/harness/session/session.ts"; @@ -21,7 +23,7 @@ const configuration: LaneConfiguration = { async function createLane( emitBatch: (events: readonly HarnessEvent[], context: Context) => Promise = () => Promise.resolve(), ): Promise<{ - lane: Lane; + lane: Lane; model: Model; session: Session; storage: ControlledMemoryStorage; @@ -49,20 +51,28 @@ async function createLane( const models = createModels(); models.setProvider(faux.provider); return { - lane: new Lane( + lane: new Lane( "main", session, models, + new HookRegistry(() => {}), await restoreLane(session, "main", BACKGROUND_CONTEXT), (cause) => (cause instanceof Error ? cause : new Error(String(cause))), emitBatch, (snapshot) => ({ snapshot, start: () => {}, unsubscribe: () => {} }), () => ({ + tools: [], resources: {}, + streamOptions: {}, + retryPolicy: { enabled: true, maxRetries: 3, baseDelayMs: 1_000 }, compaction: DEFAULT_COMPACTION_SETTINGS, steeringMode: "all", followUpMode: "all", toolExecution: "parallel", + toolContext: undefined, + systemPrompt: undefined, + toProviderMessages: (messages) => convertToLlm(messages), + entryProjectors: {}, }), ), model: faux.getModel(), @@ -72,7 +82,7 @@ async function createLane( } function setThinkingLevel( - lane: Lane, + lane: Lane, thinkingLevel: LaneConfiguration["thinkingLevel"], observed?: LaneConfiguration["thinkingLevel"][], ): Promise { diff --git a/packages/agent/test/harness/runtime2/progress.test.ts b/packages/agent/test/harness/runtime2/progress.test.ts new file mode 100644 index 000000000..21487e510 --- /dev/null +++ b/packages/agent/test/harness/runtime2/progress.test.ts @@ -0,0 +1,330 @@ +import { type AssistantMessageFrame, createModels } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it } from "vitest"; +import { OperationMismatch } from "../../../src/harness/agent-harness.ts"; +import { DEFAULT_COMPACTION_SETTINGS } from "../../../src/harness/compaction/compaction.ts"; +import { BACKGROUND_CONTEXT, type Context } from "../../../src/harness/context.ts"; +import { HookRegistry } from "../../../src/harness/hooks.ts"; +import { openFrameProgress, openToolProgress } from "../../../src/harness/runtime2/progress.ts"; +import { + type ActiveDrive, + type Config, + createActiveDrive, + type DriveScope, + type LaneCommand, + type LaneDriveCapabilities, + type LaneState as RuntimeLaneState, +} from "../../../src/harness/runtime2/types.ts"; +import { MemoryStorage } from "../../../src/harness/session/memory.ts"; +import { StorageBackedSession } from "../../../src/harness/session/session.ts"; +import type { LaneLastResult, RunState, Session, SessionReader, Write } from "../../../src/harness/session/types.ts"; +import * as storedValues from "../../../src/harness/session/values.ts"; +import { createDriveScope } from "./test-utils.ts"; + +const sessions: Session[] = []; +const operationId = "operation"; +const configuration = { + model: { provider: "test", modelId: "model" }, + thinkingLevel: "off" as const, + activeToolNames: [], +}; + +class FailingStorage extends MemoryStorage { + failure: Error | undefined; + + override commit(writes: Write[], context: Context) { + if (this.failure === undefined) return super.commit(writes, context); + const failure = this.failure; + this.failure = undefined; + return Promise.reject(failure); + } +} + +class TestDriveLane implements LaneDriveCapabilities { + readonly name = "main"; + readonly models = createModels(); + readonly hooks = new HookRegistry(() => {}); + readonly session: Session; + private projection: RuntimeLaneState; + active: ActiveDrive | undefined; + + constructor(session: Session, state: RunState) { + this.session = session; + this.projection = { + leafId: null, + configuration, + pendingNextRun: [], + operation: { + meta: { + operationId, + lane: this.name, + sourceLeafId: null, + startedAt: 1, + intent: { kind: "run", promptEntryIds: [] }, + }, + state, + }, + }; + } + + command( + plan: ( + projection: RuntimeLaneState, + reader: SessionReader, + ) => LaneCommand | Promise>, + context: Context, + ): Promise { + return this.session.mutate( + this.name, + async (mutator) => { + const decision = await plan(this.projection, mutator); + switch (decision.kind) { + case "return": + return decision.result; + case "reject": + throw decision.error; + case "commit": { + const commit = await mutator.commit(decision.writes, context); + this.projection = decision.next; + return decision.materialize(commit); + } + } + }, + context, + ); + } + + emitBatch(): Promise { + return Promise.resolve(); + } + + readConfig(): Config { + return { + tools: [], + resources: {}, + streamOptions: {}, + retryPolicy: { enabled: true, maxRetries: 3, baseDelayMs: 1_000 }, + compaction: DEFAULT_COMPACTION_SETTINGS, + steeringMode: "all", + followUpMode: "all", + toolExecution: "parallel", + toolContext: undefined, + systemPrompt: undefined, + toProviderMessages: () => [], + entryProjectors: {}, + }; + } + + ownsDrive(scope: DriveScope): boolean { + return this.active === scope.owner; + } + + mismatch(expected: string, currentOperationId: string | null, last: LaneLastResult | undefined): OperationMismatch { + return new OperationMismatch({ + lane: this.name, + expectedOperationId: expected, + ...(currentOperationId === null ? {} : { currentOperationId }), + ...(last === undefined ? {} : { lastOperationId: last.operationId }), + message: "mismatch", + }); + } +} + +function runState(phase: RunState["phase"]): RunState { + return { + kind: "run", + control: { status: "running" }, + settings: { + compaction: DEFAULT_COMPACTION_SETTINGS, + steeringMode: "all", + followUpMode: "all", + toolExecution: "parallel", + }, + phase, + inbox: { steer: [], followUp: [], writes: [] }, + latestAssistantEntryId: null, + }; +} + +async function createFixture( + state: RunState, + storage = new FailingStorage(), +): Promise<{ + lane: TestDriveLane; + scope: DriveScope; + storage: FailingStorage; +}> { + const session = new StorageBackedSession( + { id: `progress-${sessions.length}`, createdAt: 1, storageVersion: 1 }, + storage, + ); + sessions.push(session); + const lane = new TestDriveLane(session, state); + await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + storedValues.setValue(storedValues.laneLeaf("main"), null), + storedValues.setValue(storedValues.laneConfig("main"), configuration), + storedValues.setValue(storedValues.laneState("main"), { + currentOperationId: operationId, + pendingNextRun: [], + }), + storedValues.setValue(storedValues.operationMeta(operationId), { + operationId, + lane: "main", + sourceLeafId: null, + startedAt: 1, + intent: { kind: "run", promptEntryIds: [] }, + }), + storedValues.setValue(storedValues.operationState(operationId), state), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, + ); + const { scope, handle } = createDriveScope(lane, { operationId }); + lane.active = handle.active; + return { lane, scope, storage }; +} + +afterEach(async () => { + for (const session of sessions.splice(0)) await session.close(BACKGROUND_CONTEXT); +}); + +describe("runtime2 active drive ownership", () => { + it("settles or rejects the shared completion without exposing owner controls through the gate", async () => { + const settled = createActiveDrive(operationId); + const outcome = { kind: "waiting", operationId, reason: "retry", notBefore: 10 } as const; + settled.settle(outcome); + await expect(settled.active.completion).resolves.toBe(outcome); + + const failed = createActiveDrive(operationId); + const error = new Error("drive failed"); + failed.fail(error); + await expect(failed.active.completion).rejects.toBe(error); + expect("beginAbort" in failed.active.gate).toBe(false); + }); +}); + +describe("runtime2 progress channels", () => { + it("enqueues assistant frames in order, seals admission, and exposes the settlement delete", async () => { + const responseEntryId = "response"; + const { scope } = await createFixture( + runState({ + kind: "assistant", + generation: { + status: "effect_pending", + context: { + stepId: "step", + triggerEntryId: "trigger", + configuration, + streamOptions: {}, + retryPolicy: { maxAttempts: 2, baseDelayMs: 1 }, + overflowRecoveryUsed: false, + }, + attempt: 1, + responseEntryId, + usageId: "usage", + intendedOutputLimit: 100, + contextWindow: 1_000, + }, + }), + ); + const progress = openFrameProgress(scope, responseEntryId); + const frames: AssistantMessageFrame[] = [ + { type: "text_delta", contentIndex: 0, delta: "a" }, + { type: "text_delta", contentIndex: 0, delta: "b" }, + ]; + + progress.write(frames[0]!); + progress.write(frames[1]!); + progress.seal(); + progress.write({ type: "text_delta", contentIndex: 0, delta: "late" }); + await progress.drain(); + + expect( + ( + await scope.lane.session.readList( + storedValues.pendingAssistantFrames(operationId, responseEntryId), + undefined, + BACKGROUND_CONTEXT, + ) + ).map(({ value }) => value), + ).toEqual(frames); + expect(progress.clearWrite()).toEqual( + storedValues.deleteList(storedValues.pendingAssistantFrames(operationId, responseEntryId)), + ); + }); + + it("replaces tool checkpoints and fences writes after owner identity changes", async () => { + const invocationId = "result"; + const { lane, scope } = await createFixture( + runState({ + kind: "tools", + batch: { + assistantEntryId: "assistant", + configuration, + turnId: "turn", + calls: [{ status: "effect_pending", sourceIndex: 0, resultEntryId: invocationId, replay: "safe" }], + }, + }), + ); + const progress = openToolProgress(scope, invocationId); + progress.write({ content: [{ type: "text", text: "first" }], details: {} }); + progress.write({ content: [{ type: "text", text: "second" }], details: {} }); + await progress.drain(); + expect( + ( + await scope.lane.session.getValue( + storedValues.pendingToolOutput(operationId, invocationId), + BACKGROUND_CONTEXT, + ) + )?.value.content, + ).toEqual([{ type: "text", text: "second" }]); + + lane.active = createActiveDrive(operationId).active; + progress.write({ content: [{ type: "text", text: "stale" }], details: {} }); + await progress.drain(); + expect( + ( + await scope.lane.session.getValue( + storedValues.pendingToolOutput(operationId, invocationId), + BACKGROUND_CONTEXT, + ) + )?.value.content, + ).toEqual([{ type: "text", text: "second" }]); + }); + + it("retains the rejecting write promise so drain propagates commit failure", async () => { + const responseEntryId = "response"; + const { scope, storage } = await createFixture( + runState({ + kind: "assistant", + generation: { + status: "effect_pending", + context: { + stepId: "step", + triggerEntryId: "trigger", + configuration, + streamOptions: {}, + retryPolicy: { maxAttempts: 2, baseDelayMs: 1 }, + overflowRecoveryUsed: false, + }, + attempt: 1, + responseEntryId, + usageId: "usage", + intendedOutputLimit: 100, + contextWindow: 1_000, + }, + }), + ); + const failure = new Error("frame commit failed"); + storage.failure = failure; + const progress = openFrameProgress(scope, responseEntryId); + + progress.write({ type: "text_delta", contentIndex: 0, delta: "lost" }); + + await expect(progress.drain()).rejects.toBe(failure); + }); +}); diff --git a/packages/agent/test/harness/runtime2/test-utils.ts b/packages/agent/test/harness/runtime2/test-utils.ts index b1e4883d3..6b55f6b95 100644 --- a/packages/agent/test/harness/runtime2/test-utils.ts +++ b/packages/agent/test/harness/runtime2/test-utils.ts @@ -1,4 +1,11 @@ -import { BACKGROUND_CONTEXT } from "../../../src/harness/context.ts"; +import type { DriveOptions } from "../../../src/harness/agent-harness.ts"; +import { BACKGROUND_CONTEXT, type Context, withoutAbortSignal } from "../../../src/harness/context.ts"; +import { + type ActiveDriveHandle, + createActiveDrive, + type DriveScope, + type LaneDriveCapabilities, +} from "../../../src/harness/runtime2/types.ts"; import { MemoryStorage } from "../../../src/harness/session/memory.ts"; import type { CommitResult, Write } from "../../../src/harness/session/types.ts"; @@ -23,6 +30,35 @@ export class ControlledMemoryStorage extends MemoryStorage { } } +export function createDriveScope( + lane: LaneDriveCapabilities, + options: DriveOptions, + context: Context = BACKGROUND_CONTEXT, +): { scope: DriveScope; handle: ActiveDriveHandle } { + const handle = createActiveDrive(options.operationId); + return { + handle, + scope: { + lane, + operationId: options.operationId, + owner: handle.active, + context: withoutAbortSignal(context), + pass: { + waitForRetry: options.waitForRetry ?? false, + deferredPermits: options.pollDeferred === true ? 1 : 0, + }, + }, + }; +} + +export function createDriveFixture( + lane: LaneDriveCapabilities, + operationId = "operation", + context: Context = BACKGROUND_CONTEXT, +): { scope: DriveScope; handle: ActiveDriveHandle } { + return createDriveScope(lane, { operationId }, context); +} + export function deferred(): { promise: Promise; resolve(): void } { let resolvePromise: (() => void) | undefined; const promise = new Promise((resolve) => { diff --git a/packages/agent/test/harness/storage-backed-session.test.ts b/packages/agent/test/harness/storage-backed-session.test.ts index e2337c6ea..538bcd3e8 100644 --- a/packages/agent/test/harness/storage-backed-session.test.ts +++ b/packages/agent/test/harness/storage-backed-session.test.ts @@ -235,6 +235,17 @@ describe("StorageBackedSession", () => { await session.close(BACKGROUND_CONTEXT); }); + it("accepts an injected id generator for deterministic execution tests", async () => { + let next = 0; + const idGenerator = { next: (timestampMs?: number) => `${timestampMs ?? "now"}:${++next}` }; + const session = new StorageBackedSession(metadata, new MemoryStorage({ now: () => NOW }), { idGenerator }); + + expect(session.idGenerator).toBe(idGenerator); + expect(session.idGenerator.next(7)).toBe("7:1"); + expect(session.idGenerator.next()).toBe("now:2"); + await session.close(BACKGROUND_CONTEXT); + }); + it("exposes metadata directly and the shared UUIDv7 id generator", async () => { const sourceMetadata = { ...metadata }; const session = new StorageBackedSession(sourceMetadata, new MemoryStorage({ now: () => NOW })); diff --git a/packages/agent/test/harness/types.test.ts b/packages/agent/test/harness/types.test.ts index 0ea45b473..a94ab806c 100644 --- a/packages/agent/test/harness/types.test.ts +++ b/packages/agent/test/harness/types.test.ts @@ -2,6 +2,7 @@ import type { AssistantMessage, AssistantMessageFrame, Usage } from "@earendil-w import { expectTypeOf, it } from "vitest"; import * as storedValues from "../../src/harness/session/values.ts"; import type { + AgentHarness, AgentHarnessOptions, AgentHarnessStreamOptions, AgentHarnessTool, @@ -53,6 +54,7 @@ import type { SessionSearchHit, SessionSearchService, SessionSnapshot, + SessionStats, SessionTree, SettledAssistantMessage, Storage, @@ -372,7 +374,7 @@ it("covers storage, session, repository, search, and identity signatures", () => ( transactionToCommit: Write[], context: Context, - ) => Promise<{ firstSeq: number; seqs: number[]; timestamp: number }> + ) => Promise<{ firstSeq: number; seqs: number[]; timestamp: number; stats: SessionStats }> >(); expectTypeOf().toEqualTypeOf< ( @@ -385,7 +387,7 @@ it("covers storage, session, repository, search, and identity signatures", () => ( transactionToCommit: Write[], context: Context, - ) => Promise<{ firstSeq: number; seqs: number[]; timestamp: number }> + ) => Promise<{ firstSeq: number; seqs: number[]; timestamp: number; stats: SessionStats }> >(); expectTypeOf().toEqualTypeOf< (query: StorageBranchScan, context: Context) => Promise @@ -472,6 +474,12 @@ it("covers Part 5 results, events, hooks, snapshots, tools, and stream options", expectTypeOf().toEqualTypeOf<"operationId" | "waitForRetry" | "pollDeferred">(); expectTypeOf().toEqualTypeOf<"settled" | "waiting">(); expectTypeOf().returns.toEqualTypeOf>(); + expectTypeOf< + Extract + >().toEqualTypeOf(); + expectTypeOf< + Extract + >().toEqualTypeOf(); expectTypeOf().toEqualTypeOf< "prompt" | "skill" | "prompt_template" | "compaction" | "navigation" >(); diff --git a/packages/session-backends/sqlite-node/CHANGELOG.md b/packages/session-backends/sqlite-node/CHANGELOG.md index 30064cb53..be219b1a9 100644 --- a/packages/session-backends/sqlite-node/CHANGELOG.md +++ b/packages/session-backends/sqlite-node/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- Returned exact post-commit session totals from SQLite storage commits. + ## [0.84.2] - 2026-08-14 ## [0.84.1] - 2026-08-07 diff --git a/packages/session-backends/sqlite-node/src/sqlite/storage.ts b/packages/session-backends/sqlite-node/src/sqlite/storage.ts index 1f0df84c8..fc8323d8d 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/storage.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/storage.ts @@ -197,7 +197,7 @@ export class SqliteStorage implements Storage { } } advanceNextSeq(this.db, firstSeq + prepared.writes.length); - return prepared.result; + return { ...prepared.result, stats: readSessionStats(this.db) }; }); } diff --git a/packages/session-backends/sqlite-node/test/storage.test.ts b/packages/session-backends/sqlite-node/test/storage.test.ts index 7e2928b79..b417f87b4 100644 --- a/packages/session-backends/sqlite-node/test/storage.test.ts +++ b/packages/session-backends/sqlite-node/test/storage.test.ts @@ -7,6 +7,15 @@ import { applyInitialSchema } from "../src/sqlite/migrations.ts"; import { advanceNextSeq, readNextSeq } from "../src/sqlite/session/session-sequences.ts"; import { listValueReadQuery } from "../src/sqlite/session/values.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 }, +}; + async function withStorage(run: (storage: SqliteStorage, db: SqliteDatabase) => Promise): Promise { const db = await createNodeSqliteFactory().open(":memory:"); try { @@ -41,14 +50,7 @@ function insertCommitSessionRow(db: SqliteDatabase, nextSeq = 1): void { ${1}, ${null}, ${0}, - ${JSON.stringify({ - input: 0, - output: 0, - cacheRead: 0, - cacheWrite: 0, - totalTokens: 0, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - })}, + ${JSON.stringify(ZERO_USAGE)}, ${nextSeq} )`.run(db); } @@ -70,7 +72,12 @@ describe("SqliteStorage", () => { ], BACKGROUND_CONTEXT, ), - ).toEqual({ firstSeq: 1, seqs: [1], timestamp: 1_700_000_000_000 }); + ).toEqual({ + firstSeq: 1, + seqs: [1], + timestamp: 1_700_000_000_000, + stats: { messageCount: 1, usage: ZERO_USAGE }, + }); expect( await storage.commit( [ @@ -83,7 +90,12 @@ describe("SqliteStorage", () => { ], BACKGROUND_CONTEXT, ), - ).toEqual({ firstSeq: 2, seqs: [2], timestamp: 1_700_000_000_000 }); + ).toEqual({ + firstSeq: 2, + seqs: [2], + timestamp: 1_700_000_000_000, + stats: { messageCount: 2, usage: ZERO_USAGE }, + }); expect(sql`SELECT branch_id, tip_entry_id, tip_seq FROM branch_meta`.all(db)).toEqual([ { branch_id: "root", tip_entry_id: "child", tip_seq: 2 }, @@ -521,6 +533,47 @@ describe("SqliteStorage", () => { VALUES (${1}, ${null}, ${1}, ${null}, ${2}, ${JSON.stringify(usage)}, ${3})`.run(db); expect(await storage.getStats(BACKGROUND_CONTEXT)).toEqual({ messageCount: 2, usage }); + const next = await storage.commit( + [storedValues.setValue(storedValues.sessionName, "after-history")], + BACKGROUND_CONTEXT, + ); + expect(next.stats).toEqual({ messageCount: 2, usage }); + }); + }); + + it("includes historical totals in the first commit after storage reopen", async () => { + await withStorage(async (storage, db) => { + insertCommitSessionRow(db); + const usage = { + input: 1, + output: 2, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 3, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + await storage.commit( + [ + sessionWrites.insertEntry({ + id: "history", + parentId: null, + type: "message", + message: { role: "user", content: "history", timestamp: 1 }, + }), + sessionWrites.insertUsage({ id: "usage", usage, adjustment: false }), + ], + BACKGROUND_CONTEXT, + ); + await storage.close(BACKGROUND_CONTEXT); + + const reopened = new SqliteStorage(db, { now: () => 1_700_000_000_000 }); + const result = await reopened.commit( + [storedValues.setValue(storedValues.sessionName, "reopened")], + BACKGROUND_CONTEXT, + ); + expect(result.stats).toEqual({ messageCount: 1, usage }); + expect(result.stats).toEqual(await reopened.getStats(BACKGROUND_CONTEXT)); + await reopened.close(BACKGROUND_CONTEXT); }); });