From 3fbfcf48ff7d97d2acb6685902f8c07f5a3a50c1 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Mon, 17 Aug 2026 19:54:16 +0200 Subject: [PATCH] feat(agent): thread invocation context through harness APIs --- packages/agent/src/harness/agent-harness.ts | 126 ++-- .../compaction/branch-summarization.ts | 8 +- packages/agent/src/harness/events.ts | 71 +- packages/agent/src/harness/hooks.ts | 188 +++-- .../agent/src/harness/runtime2/harness.ts | 173 ++--- packages/agent/src/harness/runtime2/lane.ts | 179 ++--- .../agent/src/harness/runtime2/restore.ts | 25 +- packages/agent/src/harness/session/context.ts | 7 +- .../agent/src/harness/session/jsonl/repo.ts | 100 +-- .../src/harness/session/jsonl/storage.ts | 40 +- packages/agent/src/harness/session/memory.ts | 213 +++--- packages/agent/src/harness/session/session.ts | 493 ++++++++------ .../session/testing/benchmark/session-repo.ts | 33 +- .../session/testing/benchmark/storage.ts | 17 +- .../testing/conformance/session-repo.ts | 476 +++++++------ .../session/testing/conformance/storage.ts | 641 +++++++++++------- .../session/testing/instrumented-storage.ts | 5 +- .../session/testing/storage-decorator.ts | 49 +- packages/agent/src/harness/session/types.ts | 115 ++-- .../test/harness/branch-summarization.test.ts | 13 +- .../agent/test/harness/compaction.test.ts | 7 +- .../test/harness/execution-primitives.test.ts | 138 +++- .../test/harness/instrumented-storage.test.ts | 107 +-- .../jsonl-session-repo-conformance.test.ts | 29 +- .../test/harness/jsonl-session-repo.test.ts | 71 +- .../harness/jsonl-storage-conformance.test.ts | 4 +- .../agent/test/harness/jsonl-storage.test.ts | 204 +++--- .../test/harness/memory-conformance.test.ts | 5 +- .../test/harness/memory-session-repo.test.ts | 55 +- .../agent/test/harness/memory-storage.test.ts | 31 +- .../test/harness/runtime2/harness.test.ts | 559 ++++++++------- .../agent/test/harness/runtime2/lane.test.ts | 77 ++- .../test/harness/runtime2/restore.test.ts | 143 ++-- .../agent/test/harness/runtime2/test-utils.ts | 5 +- .../test/harness/session-context.test.ts | 52 +- .../test/harness/session-create-lane.test.ts | 114 ++-- .../agent/test/harness/session-tree.test.ts | 281 ++++---- .../harness/storage-backed-session.test.ts | 123 ++-- packages/agent/test/harness/types.test.ts | 25 +- packages/agent/test/harness/values.test.ts | 13 +- .../coding-agent/src/experimental/server.ts | 12 +- .../src/experimental/session-worker.ts | 55 +- .../test/experimental-session-support.ts | 26 +- .../test/fixtures/faux-session-worker.ts | 19 +- packages/server/src/testing/host.ts | 20 +- .../sqlite-node/src/sqlite/repo.ts | 15 +- .../sqlite-node/src/sqlite/session.ts | 151 +++-- .../sqlite-node/src/sqlite/storage.ts | 30 +- .../sqlite-node/test/repo.test.ts | 296 ++++---- .../test/storage-conformance.test.ts | 3 +- .../sqlite-node/test/storage.test.ts | 249 ++++--- 51 files changed, 3469 insertions(+), 2422 deletions(-) diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index d4b810333..c7fb7826a 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -15,6 +15,7 @@ import type { TelemetryContext } from "@earendil-works/pi-telemetry"; import type { AgentMessage, AgentToolResult, QueueMode, ThinkingLevel } from "../types.ts"; import type { BranchPreparation, BranchSummaryResult } from "./compaction/branch-summarization.ts"; 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 type { @@ -499,7 +500,10 @@ export type HarnessEvent = | (HandlerErrorPayload & ({ lane: string; recovery?: true } | { lane?: never; recovery?: never })); export type HarnessEventType = HarnessEvent["type"]; -export type EventListener = (event: TEvent) => void | Promise; +export type EventListener = ( + event: TEvent, + context: Context, +) => void | Promise; export interface Events { on( @@ -591,6 +595,7 @@ export type HookInvocation = HookMap[TName]["event"] & { }; export type HookHandler = ( event: HookInvocation, + context: Context, ) => Promise | HookMap[TName]["result"]; export interface Hooks { @@ -606,8 +611,8 @@ export interface AgentHarnessOptions[]; - toolContext?: TContext | (() => TContext | Promise); - systemPrompt?: string | ((context: TContext) => string | Promise); + toolContext?: TContext | ((context: Context) => TContext | Promise); + systemPrompt?: string | ((toolContext: TContext, context: Context) => string | Promise); resources?: Resources; streamOptions?: AgentHarnessStreamOptions; retry?: RetryPolicy; @@ -616,74 +621,87 @@ export interface AgentHarnessOptions Message[] | Promise; + toProviderMessages?: (messages: AgentMessage[], context: Context) => Message[] | Promise; entryProjectors?: Record; telemetryContext?: TelemetryContext; } export interface AgentLane { readonly name: string; - getLeafId(): Promise; - getLastResult(): Promise; - accept(request: OperationRequest): Promise; - drive(options: DriveOptions): Promise; - requestAbort(operationId: string): Promise; - inspectExecution(): Promise; - prompt(text: string, images?: ImageContent[]): Promise; - prompt(message: AgentMessage | AgentMessage[]): Promise; - skill(name: string, additionalInstructions?: string): Promise; - promptFromTemplate(name: string, args?: string[]): Promise; - compact(options?: { customInstructions?: string }): Promise; - navigateTree(targetId: string | null, options?: NavigateOptions): Promise; - resume(): Promise; - abort(): Promise; - steer(message: string | AgentMessage, images?: ImageContent[]): Promise; - followUp(message: string | AgentMessage, images?: ImageContent[]): Promise; - nextRun(message: string | AgentMessage, images?: ImageContent[]): Promise; - cancelQueued(entryId: string): Promise; - recordUsage(usage: Usage, options?: { entryId?: string; details?: JsonValue }): Promise; - waitForIdle(): Promise; - runWhenIdle(callback: () => void | Promise): Promise; - peekAction(): Promise; - executeAction(): Promise; - runToCompletion(): Promise; - getModel(): Promise | undefined>; - setModel(model: Model): Promise; - getThinkingLevel(): Promise; - setThinkingLevel(level: ThinkingLevel): Promise; - getActiveTools(): Promise; - setActiveTools(names: string[]): Promise; + getLeafId(context: Context): Promise; + getLastResult(context: Context): Promise; + accept(request: OperationRequest, context: Context): Promise; + drive(options: DriveOptions, context: Context): Promise; + requestAbort(operationId: string, context: Context): Promise; + inspectExecution(context: Context): Promise; + prompt(text: string, images: ImageContent[] | undefined, context: Context): Promise; + prompt(message: AgentMessage | AgentMessage[], context: Context): Promise; + skill(name: string, additionalInstructions: string | undefined, context: Context): Promise; + promptFromTemplate(name: string, args: string[] | undefined, context: Context): Promise; + compact(options: { customInstructions?: string } | undefined, context: Context): Promise; + navigateTree( + targetId: string | null, + options: NavigateOptions | undefined, + context: Context, + ): Promise; + resume(context: Context): Promise; + abort(context: Context): Promise; + steer(message: string | AgentMessage, images: ImageContent[] | undefined, context: Context): Promise; + followUp(message: string | AgentMessage, images: ImageContent[] | undefined, context: Context): Promise; + nextRun( + message: string | AgentMessage, + images: ImageContent[] | undefined, + context: Context, + ): Promise; + cancelQueued(entryId: string, context: Context): Promise; + recordUsage( + usage: Usage, + options: { entryId?: string; details?: JsonValue } | undefined, + context: Context, + ): Promise; + waitForIdle(context: Context): Promise; + runWhenIdle(callback: (context: Context) => void | Promise, context: Context): Promise; + peekAction(context: Context): Promise; + executeAction(context: Context): Promise; + runToCompletion(context: Context): Promise; + getModel(context: Context): Promise | undefined>; + setModel(model: Model, context: Context): Promise; + getThinkingLevel(context: Context): Promise; + setThinkingLevel(level: ThinkingLevel, context: Context): Promise; + getActiveTools(context: Context): Promise; + setActiveTools(names: string[], context: Context): Promise; readonly sessionTree: SessionTree; - watch(): Promise>; + watch(context: Context): Promise>; } export interface AgentHarness extends AgentLane { - lane(name: string): Promise; - createLane(name: string, at: string | null): Promise; - lanes(): Promise; - getTools(): Promise[]>; - setTools(tools: AgentHarnessTool[]): Promise; - getResources(): Promise; - setResources(resources: Resources): Promise; - getStreamOptions(): Promise; - setStreamOptions(options: AgentHarnessStreamOptions): Promise; - getRetryPolicy(): Promise; - setRetryPolicy(policy: RetryPolicy): Promise; - getCompactionSettings(): Promise; - setCompactionSettings(settings: CompactionSettings): Promise; - getSteeringMode(): Promise; - setSteeringMode(mode: QueueMode): Promise; - getFollowUpMode(): Promise; - setFollowUpMode(mode: QueueMode): Promise; - watchSession(): Promise>; + lane(name: string, context: Context): Promise; + createLane(name: string, at: string | null, context: Context): Promise; + lanes(context: Context): Promise; + getTools(context: Context): Promise[]>; + setTools(tools: AgentHarnessTool[], context: Context): Promise; + getResources(context: Context): Promise; + setResources(resources: Resources, context: Context): Promise; + getStreamOptions(context: Context): Promise; + setStreamOptions(options: AgentHarnessStreamOptions, context: Context): Promise; + getRetryPolicy(context: Context): Promise; + setRetryPolicy(policy: RetryPolicy, context: Context): Promise; + getCompactionSettings(context: Context): Promise; + setCompactionSettings(settings: CompactionSettings, context: Context): Promise; + getSteeringMode(context: Context): Promise; + setSteeringMode(mode: QueueMode, context: Context): Promise; + getFollowUpMode(context: Context): Promise; + setFollowUpMode(mode: QueueMode, context: Context): Promise; + watchSession(context: Context): Promise>; readonly hooks: Hooks; readonly events: Events; - close(): Promise; + close(context: Context): Promise; } export interface AgentHarnessConstructor { create( options: AgentHarnessOptions, + context: Context, ): Promise<{ harness: AgentHarness; suspended: SuspendedOperation[] }>; } diff --git a/packages/agent/src/harness/compaction/branch-summarization.ts b/packages/agent/src/harness/compaction/branch-summarization.ts index fc4db8df6..c3560e31f 100644 --- a/packages/agent/src/harness/compaction/branch-summarization.ts +++ b/packages/agent/src/harness/compaction/branch-summarization.ts @@ -9,6 +9,7 @@ import { } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; +import type { Context } from "../context.ts"; import { convertToLlm, createBranchSummaryMessage, createCompactionSummaryMessage } from "../messages.ts"; import type { Entry, SessionTree } from "../session/index.ts"; import { BranchSummaryError, err, ok, type Result } from "../types.ts"; @@ -83,12 +84,13 @@ export async function collectEntriesForBranchSummary( session: Pick, oldLeafId: string | null, targetId: string, + context: Context, ): Promise { if (!oldLeafId) { return { entries: [], commonAncestorId: null }; } - const oldPath = new Set((await session.findEntriesOnBranch({ start: oldLeafId })).map((entry) => entry.id)); - const targetPath = await session.findEntriesOnBranch({ start: targetId }); + const oldPath = new Set((await session.findEntriesOnBranch({ start: oldLeafId }, context)).map((entry) => entry.id)); + const targetPath = await session.findEntriesOnBranch({ start: targetId }, context); let commonAncestorId: string | null = null; for (const entry of targetPath) { if (oldPath.has(entry.id)) { @@ -100,7 +102,7 @@ export async function collectEntriesForBranchSummary( let current: string | null = oldLeafId; while (current && current !== commonAncestorId) { - const entry = await session.getEntry(current); + const entry = await session.getEntry(current, context); if (!entry) throw new Error(`Corrupt session: entry ${current} not found`); entries.push(entry); current = entry.parentId; diff --git a/packages/agent/src/harness/events.ts b/packages/agent/src/harness/events.ts index 3368f70f3..45a151bb8 100644 --- a/packages/agent/src/harness/events.ts +++ b/packages/agent/src/harness/events.ts @@ -1,6 +1,7 @@ import type { EventListener, Events, HarnessEvent, HarnessEventType, WatchHandle } from "./agent-harness.ts"; +import type { Context } from "./context.ts"; -type UntypedEventListener = (event: HarnessEvent) => void | Promise; +type UntypedEventListener = (event: HarnessEvent, context: Context) => void | Promise; /** Passive harness event bus with isolated handler failures. */ export class HarnessEventBus implements Events { @@ -14,7 +15,8 @@ export class HarnessEventBus implements Events { listener: EventListener>, ): () => void { if (this.closedError !== undefined) throw this.closedError; - const wrapped: UntypedEventListener = (event) => listener(event as Extract); + const wrapped: UntypedEventListener = (event, context) => + listener(event as Extract, context); let listeners = this.listeners.get(type); if (listeners === undefined) { listeners = new Set(); @@ -24,26 +26,27 @@ export class HarnessEventBus implements Events { return () => listeners?.delete(wrapped); } - emit(event: HarnessEvent): Promise { + emit(event: HarnessEvent, context: Context): Promise { if (this.closedError !== undefined) return Promise.resolve(); - const delivery = this.deliveryTail.then(() => this.deliver(event, true)); + const delivery = this.deliveryTail.then(() => this.deliver(event, true, context)); this.deliveryTail = delivery.catch(() => {}); return delivery; } - watch(snapshot: T, filter: (event: HarnessEvent) => boolean): WatchHandle { + watch(snapshot: T, filter: (event: HarnessEvent) => boolean, _context: Context): WatchHandle { if (this.closedError !== undefined) throw this.closedError; return this.installWatcher(snapshot, filter); } async watchFromSnapshot( - capture: () => Promise, + capture: (context: Context) => Promise, filter: (event: HarnessEvent) => boolean, + context: Context, ): Promise> { if (this.closedError !== undefined) throw this.closedError; const watcher = this.installWatcher(undefined, filter); try { - watcher.setSnapshot(await capture()); + watcher.setSnapshot(await capture(context)); return watcher; } catch (error) { watcher.unsubscribe(); @@ -63,32 +66,35 @@ export class HarnessEventBus implements Events { snapshot: T | undefined, filter: (event: HarnessEvent) => boolean, ): BufferedEventWatcher { - const watcher = new BufferedEventWatcher(snapshot, async (error, event) => { + const watcher = new BufferedEventWatcher(snapshot, async (error, event, context) => { if (event.type === "handler_error") return; const normalized = error instanceof Error ? error : new Error(String(error)); const lane = "lane" in event && typeof event.lane === "string" ? event.lane : undefined; - await this.emit({ - type: "handler_error", - kind: "event", - event: event.type, - error: normalized.message, - ...(normalized.stack === undefined ? {} : { stack: normalized.stack }), - ...(lane === undefined ? {} : { lane }), - }); + await this.emit( + { + type: "handler_error", + kind: "event", + event: event.type, + error: normalized.message, + ...(normalized.stack === undefined ? {} : { stack: normalized.stack }), + ...(lane === undefined ? {} : { lane }), + }, + context, + ); }); - const watchListener: UntypedEventListener = (event) => { - if (filter(event)) watcher.push(event); + const watchListener: UntypedEventListener = (event, context) => { + if (filter(event)) watcher.push(event, context); }; this.watchListeners.add(watchListener); watcher.setUnsubscribe(() => this.watchListeners.delete(watchListener)); return watcher; } - private async deliver(event: HarnessEvent, reportErrors: boolean): Promise { + private async deliver(event: HarnessEvent, reportErrors: boolean, context: Context): Promise { const listeners = [...(this.listeners.get(event.type) ?? []), ...this.watchListeners]; for (const listener of listeners) { try { - await listener(structuredClone(event)); + await listener(structuredClone(event), context); } catch (error) { if (!reportErrors || event.type === "handler_error") continue; const normalized = error instanceof Error ? error : new Error(String(error)); @@ -101,7 +107,7 @@ export class HarnessEventBus implements Events { ...(normalized.stack === undefined ? {} : { stack: normalized.stack }), ...(lane === undefined ? {} : { lane }), }; - await this.deliver(handlerError, false); + await this.deliver(handlerError, false, context); } } } @@ -109,14 +115,17 @@ export class HarnessEventBus implements Events { class BufferedEventWatcher implements WatchHandle { snapshot: T; - private readonly onError: (error: unknown, event: HarnessEvent) => void | Promise; - private buffer: HarnessEvent[] = []; + private readonly onError: (error: unknown, event: HarnessEvent, context: Context) => void | Promise; + private buffer: Array<{ event: HarnessEvent; context: Context }> = []; private listener: EventListener | undefined; private unsubscribeCallback: (() => void) | undefined; private deliveryTail: Promise = Promise.resolve(); private state: "buffering" | "started" | "unsubscribed" = "buffering"; - constructor(snapshot: T | undefined, onError: (error: unknown, event: HarnessEvent) => void | Promise) { + constructor( + snapshot: T | undefined, + onError: (error: unknown, event: HarnessEvent, context: Context) => void | Promise, + ) { this.snapshot = snapshot as T; this.onError = onError; } @@ -131,7 +140,7 @@ class BufferedEventWatcher implements WatchHandle { this.listener = listener; const buffered = this.buffer; this.buffer = []; - for (const event of buffered) this.enqueue(event); + for (const bufferedEvent of buffered) this.enqueue(bufferedEvent.event, bufferedEvent.context); } unsubscribe(): void { @@ -143,29 +152,29 @@ class BufferedEventWatcher implements WatchHandle { this.unsubscribeCallback = undefined; } - push(event: HarnessEvent): void { + push(event: HarnessEvent, context: Context): void { if (this.state === "unsubscribed") return; if (this.state === "buffering") { - this.buffer.push(event); + this.buffer.push({ event, context }); return; } - this.enqueue(event); + this.enqueue(event, context); } setUnsubscribe(callback: () => void): void { this.unsubscribeCallback = callback; } - private enqueue(event: HarnessEvent): void { + private enqueue(event: HarnessEvent, context: Context): void { const listener = this.listener; if (listener === undefined) return; this.deliveryTail = this.deliveryTail .then(async () => { - if (this.state === "started") await listener(event); + if (this.state === "started") await listener(event, context); }) .catch(async (error) => { try { - await this.onError(error, event); + await this.onError(error, event, context); } catch {} }); } diff --git a/packages/agent/src/harness/hooks.ts b/packages/agent/src/harness/hooks.ts index fdfe5638c..0c2336fe2 100644 --- a/packages/agent/src/harness/hooks.ts +++ b/packages/agent/src/harness/hooks.ts @@ -1,15 +1,22 @@ -import type { TelemetryContext } from "@earendil-works/pi-telemetry"; import type { HookHandler, HookInvocation, HookMap, HookName, Hooks } from "./agent-harness.ts"; +import { type Context, withAbortSignal, withTelemetryContext } from "./context.ts"; import type { EffectGate } from "./execution/effect-gate.ts"; import { startHarnessSpan } from "./telemetry.ts"; import type { AgentHarnessStreamOptions, AgentHarnessStreamOptionsPatch } from "./types.ts"; interface HookRegistration { id?: string; - handler: (event: unknown) => unknown | Promise; + handler: (event: unknown, context: Context) => unknown | Promise; } -type HookErrorReporter = (error: Error, hook: HookName, lane: string) => void | Promise; +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 { @@ -26,7 +33,7 @@ export class HookRegistry implements Hooks { const registrations = this.registrations.get(name) ?? []; const registration: HookRegistration = { ...(options.id === undefined ? {} : { id: options.id }), - handler: (event) => handler(event as HookInvocation), + handler: (event, context) => handler(event as HookInvocation, context), }; registrations.push(registration); this.registrations.set(name, registrations); @@ -45,9 +52,9 @@ export class HookRegistry implements Hooks { name: TName, event: HookInvocation, effectGate: EffectGate, + context: Context, ): Promise { - effectGate.assertOpen(); - return this.runAdmitted(name, event); + return this.runAdmitted(name, event, admitEffect(effectGate, context)); } /** Invoke a tool-hook aggregate with one telemetry span per registered handler. */ @@ -55,13 +62,13 @@ export class HookRegistry implements Hooks { name: TName, event: HookInvocation, effectGate: EffectGate, - telemetryContext: TelemetryContext, + context: Context, ): Promise { - effectGate.assertOpen(); + const admittedContext = admitEffect(effectGate, context); return ( name === "before_tool" - ? this.beforeTool(event as HookInvocation<"before_tool">, telemetryContext) - : this.afterTool(event as HookInvocation<"after_tool">, telemetryContext) + ? this.beforeTool(event as HookInvocation<"before_tool">, admittedContext) + : this.afterTool(event as HookInvocation<"after_tool">, admittedContext) ) as Promise; } @@ -72,58 +79,75 @@ export class HookRegistry implements Hooks { private async runAdmitted( name: TName, event: HookInvocation, + context: Context, ): Promise { if (this.closedError !== undefined) throw this.closedError; - const result = await this.aggregate(name, event); + const result = await this.aggregate(name, event, context); return result as HookMap[TName]["result"]; } - private async aggregate(name: HookName, event: HookInvocation): Promise { + private async aggregate(name: HookName, event: HookInvocation, context: Context): Promise { switch (name) { case "before_run": - return this.beforeRun(event as HookInvocation<"before_run">); + return this.beforeRun(event as HookInvocation<"before_run">, context); case "before_drive": - await this.invokeAllFailClosed(name, event as HookInvocation<"before_drive">); + await this.invokeAllFailClosed(name, event as HookInvocation<"before_drive">, context); return undefined; case "before_run_end": { let followUp: string | undefined; - await this.invokeAll(name, event, (value) => { - const result = value as HookMap["before_run_end"]["result"]; - if (result?.followUp !== undefined) followUp = result.followUp; - }); + await this.invokeAll( + name, + event, + (value) => { + const result = value as HookMap["before_run_end"]["result"]; + if (result?.followUp !== undefined) followUp = result.followUp; + }, + context, + ); return followUp === undefined ? undefined : { followUp }; } case "transform_context": - return this.transformContext(event as HookInvocation<"transform_context">); + return this.transformContext(event as HookInvocation<"transform_context">, context); case "before_request": - return this.beforeRequest(event as HookInvocation<"before_request">); + return this.beforeRequest(event as HookInvocation<"before_request">, context); case "before_payload": - return this.beforePayload(event as HookInvocation<"before_payload">); + return this.beforePayload(event as HookInvocation<"before_payload">, context); case "after_response": - return this.afterResponse(event as HookInvocation<"after_response">); + return this.afterResponse(event as HookInvocation<"after_response">, context); case "before_tool": - return this.beforeTool(event as HookInvocation<"before_tool">); + return this.beforeTool(event as HookInvocation<"before_tool">, context); case "after_tool": - return this.afterTool(event as HookInvocation<"after_tool">); + return this.afterTool(event as HookInvocation<"after_tool">, context); case "before_compaction": - return this.firstStructural(name, event as HookInvocation<"before_compaction">, "compaction"); + return this.firstStructural(name, event as HookInvocation<"before_compaction">, "compaction", context); case "before_navigation": - return this.firstStructural(name, event as HookInvocation<"before_navigation">, "summary"); + return this.firstStructural(name, event as HookInvocation<"before_navigation">, "summary", context); } } - private async beforeRun(event: HookInvocation<"before_run">): Promise { + private async beforeRun( + event: HookInvocation<"before_run">, + context: Context, + ): Promise { let prompt = event.prompt; let injected: HookMap["before_run"]["event"]["prompt"] = []; for (const registration of this.registrationsFor("before_run")) { try { - const result = (await registration.handler({ ...event, prompt })) as HookMap["before_run"]["result"]; + const result = (await registration.handler( + { ...event, prompt }, + context, + )) as HookMap["before_run"]["result"]; if (result?.messages !== undefined) { injected = [...injected, ...result.messages]; prompt = [...prompt, ...result.messages]; } } catch (error) { - await this.reportError(error instanceof Error ? error : new Error(String(error)), "before_run", event.lane); + await this.reportError( + error instanceof Error ? error : new Error(String(error)), + "before_run", + event.lane, + context, + ); } } return injected.length === 0 ? undefined : { messages: injected }; @@ -131,7 +155,7 @@ export class HookRegistry implements Hooks { private async beforeTool( event: HookInvocation<"before_tool">, - telemetryContext?: TelemetryContext, + context: Context, ): Promise { let args = event.args; let block: { reason: string; terminate?: boolean } | undefined; @@ -141,7 +165,7 @@ export class HookRegistry implements Hooks { "before_tool", registration, { ...event, args }, - telemetryContext, + context, )) as HookMap["before_tool"]["result"]; if (result?.args !== undefined) args = result.args; if (result?.block !== undefined) { @@ -150,7 +174,7 @@ export class HookRegistry implements Hooks { } } catch (error) { const normalized = error instanceof Error ? error : new Error(String(error)); - await this.reportError(normalized, "before_tool", event.lane); + await this.reportError(normalized, "before_tool", event.lane, context); block = { reason: normalized.message }; break; } @@ -163,16 +187,20 @@ export class HookRegistry implements Hooks { private async transformContext( event: HookInvocation<"transform_context">, + context: Context, ): Promise { let messages = event.messages; let systemPrompt = event.systemPrompt; for (const registration of this.registrationsFor("transform_context")) { try { - const result = (await registration.handler({ - ...event, - messages, - systemPrompt, - })) as HookMap["transform_context"]["result"]; + const result = (await registration.handler( + { + ...event, + messages, + systemPrompt, + }, + context, + )) as HookMap["transform_context"]["result"]; if (result?.messages !== undefined) messages = result.messages; if (result?.systemPrompt !== undefined) systemPrompt = result.systemPrompt; } catch (error) { @@ -180,21 +208,25 @@ export class HookRegistry implements Hooks { error instanceof Error ? error : new Error(String(error)), "transform_context", event.lane, + context, ); } } return { messages, systemPrompt }; } - private async beforeRequest(event: HookInvocation<"before_request">): Promise { + private async beforeRequest( + event: HookInvocation<"before_request">, + context: Context, + ): Promise { let streamOptions = event.streamOptions; let changed = false; for (const registration of this.registrationsFor("before_request")) { try { - const result = (await registration.handler({ - ...event, - streamOptions, - })) as HookMap["before_request"]["result"]; + const result = (await registration.handler( + { ...event, streamOptions }, + context, + )) as HookMap["before_request"]["result"]; if (result?.streamOptions !== undefined) { streamOptions = applyStreamOptionsPatch(streamOptions, result.streamOptions); changed = true; @@ -204,40 +236,55 @@ export class HookRegistry implements Hooks { error instanceof Error ? error : new Error(String(error)), "before_request", event.lane, + context, ); } } return changed ? { streamOptions: createStreamOptionsPatch(event.streamOptions, streamOptions) } : undefined; } - private async beforePayload(event: HookInvocation<"before_payload">): Promise { + private async beforePayload( + event: HookInvocation<"before_payload">, + context: Context, + ): Promise { let payload = event.payload; for (const registration of this.registrationsFor("before_payload")) { try { - const result = (await registration.handler({ ...event, payload })) as HookMap["before_payload"]["result"]; + const result = (await registration.handler( + { ...event, payload }, + context, + )) as HookMap["before_payload"]["result"]; if (result?.payload !== undefined) payload = result.payload; } catch (error) { await this.reportError( error instanceof Error ? error : new Error(String(error)), "before_payload", event.lane, + context, ); } } return { payload }; } - private async afterResponse(event: HookInvocation<"after_response">): Promise { + private async afterResponse( + event: HookInvocation<"after_response">, + context: Context, + ): Promise { let message = event.message; for (const registration of this.registrationsFor("after_response")) { try { - const result = (await registration.handler({ ...event, message })) as HookMap["after_response"]["result"]; + const result = (await registration.handler( + { ...event, message }, + context, + )) as HookMap["after_response"]["result"]; if (result?.message !== undefined) message = result.message; } catch (error) { await this.reportError( error instanceof Error ? error : new Error(String(error)), "after_response", event.lane, + context, ); } } @@ -246,7 +293,7 @@ export class HookRegistry implements Hooks { private async afterTool( event: HookInvocation<"after_tool">, - telemetryContext?: TelemetryContext, + context: Context, ): Promise { let current = { content: event.content, @@ -261,7 +308,7 @@ export class HookRegistry implements Hooks { "after_tool", registration, { ...event, ...current }, - telemetryContext, + context, )) as HookMap["after_tool"]["result"]; if (result === undefined) continue; if (result.content !== undefined) aggregate.content = result.content; @@ -276,7 +323,12 @@ export class HookRegistry implements Hooks { usage: result.usage === undefined ? current.usage : result.usage, }; } catch (error) { - await this.reportError(error instanceof Error ? error : new Error(String(error)), "after_tool", event.lane); + await this.reportError( + error instanceof Error ? error : new Error(String(error)), + "after_tool", + event.lane, + context, + ); } } return Object.keys(aggregate).length === 0 ? undefined : aggregate; @@ -286,10 +338,11 @@ export class HookRegistry implements Hooks { name: "before_compaction" | "before_navigation", event: HookInvocation<"before_compaction"> | HookInvocation<"before_navigation">, resultField: "compaction" | "summary", + context: Context, ): Promise { for (const registration of this.registrationsFor(name)) { try { - const value = await registration.handler(event); + const value = await registration.handler(event, context); if (value === undefined || value === null || typeof value !== "object") continue; const result = value as Record; if (result.decline === true && result[resultField] !== undefined) { @@ -297,12 +350,18 @@ export class HookRegistry implements Hooks { new Error(`${name} hook cannot return both decline and ${resultField}`), name, event.lane, + context, ); continue; } if (result.decline === true || result[resultField] !== undefined) return value; } catch (error) { - await this.reportError(error instanceof Error ? error : new Error(String(error)), name, event.lane); + await this.reportError( + error instanceof Error ? error : new Error(String(error)), + name, + event.lane, + context, + ); } } return undefined; @@ -312,11 +371,10 @@ export class HookRegistry implements Hooks { name: "before_tool" | "after_tool", registration: HookRegistration, event: HookInvocation<"before_tool"> | HookInvocation<"after_tool">, - telemetryContext: TelemetryContext | undefined, + context: Context, ): Promise { - if (telemetryContext === undefined) return Promise.resolve(registration.handler(event)); return startHarnessSpan( - telemetryContext, + context.telemetryContext, "pi.harness.hook", { "pi.lane.name": event.lane, @@ -326,7 +384,7 @@ export class HookRegistry implements Hooks { }, async (span) => { try { - const result = await registration.handler(event); + const result = await registration.handler(event, withTelemetryContext(span, context)); const blocked = name === "before_tool" && result !== null && @@ -348,13 +406,17 @@ export class HookRegistry implements Hooks { return [...(this.registrations.get(name) ?? [])]; } - private async invokeAllFailClosed(name: "before_drive", event: HookInvocation<"before_drive">): Promise { + private async invokeAllFailClosed( + name: "before_drive", + event: HookInvocation<"before_drive">, + context: Context, + ): Promise { for (const registration of this.registrationsFor(name)) { try { - await registration.handler(event); + await registration.handler(event, context); } catch (error) { const normalized = error instanceof Error ? error : new Error(String(error)); - await this.reportError(normalized, name, event.lane); + await this.reportError(normalized, name, event.lane, context); throw normalized; } } @@ -364,12 +426,18 @@ export class HookRegistry implements Hooks { name: HookName, event: HookInvocation, apply: (value: unknown) => void, + context: Context, ): Promise { for (const registration of this.registrationsFor(name)) { try { - apply(await registration.handler(event)); + apply(await registration.handler(event, context)); } catch (error) { - await this.reportError(error instanceof Error ? error : new Error(String(error)), name, event.lane); + await this.reportError( + error instanceof Error ? error : new Error(String(error)), + name, + event.lane, + context, + ); } } } diff --git a/packages/agent/src/harness/runtime2/harness.ts b/packages/agent/src/harness/runtime2/harness.ts index 197e2cfb3..c092dca3c 100644 --- a/packages/agent/src/harness/runtime2/harness.ts +++ b/packages/agent/src/harness/runtime2/harness.ts @@ -22,6 +22,7 @@ import { validateRetryPolicy, validateToolNames, } from "../config.ts"; +import type { Context } from "../context.ts"; import { HarnessEventBus } from "../events.ts"; import { HookRegistry } from "../hooks.ts"; import { convertToLlm } from "../messages.ts"; @@ -64,8 +65,8 @@ export class Harness extends Lane implement options.session, options.models, main, - (cause) => this.fault(cause), - (event) => this.events.emit(event), + (cause, context) => this.fault(cause, context), + (event, context) => this.events.emit(event, context), suspensionsByLane.get("main"), ); this.seed = seed; @@ -86,15 +87,18 @@ export class Harness extends Lane implement telemetryContext: options.telemetryContext ?? NOOP_TELEMETRY_CONTEXT, }; this.events = new HarnessEventBus(); - this.hooks = new HookRegistry((error, hook, lane) => - this.events.emit({ - type: "handler_error", - kind: "hook", - hook, - error: error.message, - ...(error.stack === undefined ? {} : { stack: error.stack }), - lane, - }), + 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.lanesByName.set("main", this); for (const [name, state] of restored) { @@ -102,14 +106,16 @@ export class Harness extends Lane implement } } - async lane(name: string): Promise { + async lane(name: string, _context: Context): Promise { this.assertOpen(); return this.lanesByName.get(name); } - async lanes(): Promise { + async lanes(context: Context): Promise { this.assertOpen(); - const executions = await Promise.all([...this.lanesByName.values()].map((lane) => lane.inspectExecution())); + const executions = await Promise.all( + [...this.lanesByName.values()].map((lane) => lane.inspectExecution(context)), + ); return executions.map((execution) => ({ name: execution.lane, leafId: execution.leafId, @@ -124,7 +130,7 @@ export class Harness extends Lane implement })); } - async createLane(name: string, at: string | null): Promise { + async createLane(name: string, at: string | null, context: Context): Promise { // Public Result errors are expected caller/lifecycle outcomes. Harness faults reject every API call because the // owned state can no longer advance safely; they are intentionally not members of CreateLaneResult. if (this.faultError !== undefined) throw this.faultError; @@ -136,11 +142,11 @@ export class Harness extends Lane implement operation: null, }; try { - await this.session.createLane(name, at, this.seed); + await this.session.createLane(name, at, this.seed, context); const lane = this.buildLane(name, state); if (this.closedError !== undefined) lane.seal(this.closedError); this.lanesByName.set(name, lane); - await this.events.emit({ type: "lane_created", lane: name, at }); + await this.events.emit({ type: "lane_created", lane: name, at }, context); return Result.ok(lane); } catch (error) { // Session has no Result layer, so its expected validation failures are thrown. Translate only those failures @@ -154,70 +160,70 @@ export class Harness extends Lane implement if (error instanceof SessionUnknownTargetError) { return Result.err(new UnknownTarget({ targetId: error.targetId, message: error.message })); } - throw this.fault(error); + throw this.fault(error, context); } } - getTools(): Promise[]> { - return this.getConfig("tools"); + getTools(context: Context): Promise[]> { + return this.getConfig("tools", context); } - setTools(tools: AgentHarnessTool[]): Promise { + setTools(tools: AgentHarnessTool[], context: Context): Promise { validateToolNames(tools); - return this.setConfig("tools", tools, "tools"); + return this.setConfig("tools", tools, "tools", context); } - getResources(): Promise { - return this.getConfig("resources"); + getResources(context: Context): Promise { + return this.getConfig("resources", context); } - setResources(resources: Resources): Promise { - return this.setConfig("resources", resources, "resources"); + setResources(resources: Resources, context: Context): Promise { + return this.setConfig("resources", resources, "resources", context); } - getStreamOptions(): Promise { - return this.getConfig("streamOptions"); + getStreamOptions(context: Context): Promise { + return this.getConfig("streamOptions", context); } - setStreamOptions(options: AgentHarnessStreamOptions): Promise { - return this.setConfig("streamOptions", options, "streamOptions"); + setStreamOptions(options: AgentHarnessStreamOptions, context: Context): Promise { + return this.setConfig("streamOptions", options, "streamOptions", context); } - getRetryPolicy(): Promise { - return this.getConfig("retryPolicy"); + getRetryPolicy(context: Context): Promise { + return this.getConfig("retryPolicy", context); } - setRetryPolicy(policy: RetryPolicy): Promise { + setRetryPolicy(policy: RetryPolicy, context: Context): Promise { validateRetryPolicy(policy); - return this.setConfig("retryPolicy", policy, "retryPolicy"); + return this.setConfig("retryPolicy", policy, "retryPolicy", context); } - getCompactionSettings(): Promise { - return this.getConfig("compaction"); + getCompactionSettings(context: Context): Promise { + return this.getConfig("compaction", context); } - setCompactionSettings(compaction: CompactionSettings): Promise { + setCompactionSettings(compaction: CompactionSettings, context: Context): Promise { validateCompactionSettings(compaction); - return this.setConfig("compaction", compaction, "compactionSettings"); + return this.setConfig("compaction", compaction, "compactionSettings", context); } - getSteeringMode(): Promise { - return this.getConfig("steeringMode"); + getSteeringMode(context: Context): Promise { + return this.getConfig("steeringMode", context); } - setSteeringMode(steeringMode: QueueMode): Promise { - return this.setConfig("steeringMode", steeringMode, "steeringMode"); + setSteeringMode(steeringMode: QueueMode, context: Context): Promise { + return this.setConfig("steeringMode", steeringMode, "steeringMode", context); } - getFollowUpMode(): Promise { - return this.getConfig("followUpMode"); + getFollowUpMode(context: Context): Promise { + return this.getConfig("followUpMode", context); } - setFollowUpMode(followUpMode: QueueMode): Promise { - return this.setConfig("followUpMode", followUpMode, "followUpMode"); + setFollowUpMode(followUpMode: QueueMode, context: Context): Promise { + return this.setConfig("followUpMode", followUpMode, "followUpMode", context); } - async watchSession(): Promise { + async watchSession(_context: Context): Promise { throw new SliceNotImplemented("watchSession"); } @@ -227,13 +233,16 @@ export class Harness extends Lane implement this.session, this.models, state, - (cause) => this.fault(cause), - (event) => this.events.emit(event), + (cause, context) => this.fault(cause, context), + (event, context) => this.events.emit(event, context), suspension, ); } - private async getConfig>(key: TKey): Promise[TKey]> { + private async getConfig>( + key: TKey, + _context: Context, + ): Promise[TKey]> { this.assertOpen(); return this.config[key]; } @@ -242,13 +251,14 @@ export class Harness extends Lane implement key: TKey, value: Config[TKey], property: GlobalConfigProperty, + context: Context, ): Promise { this.assertOpen(); this.config = { ...this.config, [key]: value }; - await this.events.emit({ type: "config_update", property }); + await this.events.emit({ type: "config_update", property }, context); } - fault(cause: unknown): Error { + fault(cause: unknown, context: Context): Error { if (this.faultError !== undefined) return this.faultError; if (this.closedError !== undefined) return this.closedError; const normalized = cause instanceof Error ? cause : new Error(String(cause)); @@ -256,12 +266,12 @@ export class Harness extends Lane implement this.faultError = fault; for (const lane of this.lanesByName.values()) lane.seal(fault); this.hooks.close(fault); - void this.events.emit({ type: "fault", code: "harness_fault", message: fault.message }); + void this.events.emit({ type: "fault", code: "harness_fault", message: fault.message }, context); this.events.close(fault); return fault; } - close(): Promise { + close(context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; /* * Close is a process-local admission boundary, not an abort. Public lane work checks Lane.assertOpen() at @@ -286,7 +296,7 @@ export class Harness extends Lane implement for (const lane of this.lanesByName.values()) lane.seal(error); this.hooks.close(error); this.events.close(error); - this.closePromise = this.session.close(); + this.closePromise = this.session.close(context); return this.closePromise; } } @@ -294,6 +304,7 @@ export class Harness extends Lane implement /** Attach runtime2 without starting provider, tool, hook, or timer effects. */ export async function createAgentHarness( options: AgentHarnessOptions, + context: Context, ): Promise<{ harness: AgentHarness; suspended: SuspendedOperation[] }> { const tools = options.tools ?? []; validateToolNames(tools); @@ -305,12 +316,12 @@ export async function createAgentHarness tool.name), }; try { - await seedMain(options.session, seed); - const restored = await restoreSession(options.session); + await seedMain(options.session, seed, context); + const restored = await restoreSession(options.session, context); const suspended: SuspendedOperation[] = []; for (const [lane, state] of restored) { if (state.operation !== null) { - suspended.push(await describeSuspension(options.session, options.models, tools, lane, state)); + suspended.push(await describeSuspension(options.session, options.models, tools, lane, state, context)); } } const suspensionsByLane = new Map(suspended.map((descriptor) => [descriptor.lane, descriptor])); @@ -326,6 +337,7 @@ async function describeSuspension( tools: readonly { name: string }[], lane: string, state: LaneState, + context: Context, ): Promise { const operation = state.operation; if (operation === null) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is not suspended`); @@ -334,9 +346,10 @@ async function describeSuspension( operation.state.kind === "run" && operation.state.phase.kind === "deferred" ? operation.state.phase.deferred.sourceEntryId : undefined; - const entries = await session.getEntries([ - ...new Set([...promptIds, ...(deferredSourceId === undefined ? [] : [deferredSourceId])]), - ]); + const entries = await session.getEntries( + [...new Set([...promptIds, ...(deferredSourceId === undefined ? [] : [deferredSourceId])])], + context, + ); const prompt = promptIds.map((id) => { const entry = entries.get(id); if (entry?.type !== "message") throw new SessionInvariantError(`Prompt entry ${id} is missing`); @@ -375,21 +388,25 @@ async function describeSuspension( return { ...base, reason: "crash" }; } -async function seedMain(session: Session, seed: LaneConfiguration): Promise { - await session.mutate("main", async (mutator) => { - const [leaf, state, configuration, lastResult] = await Promise.all([ - mutator.getValue(laneLeaf("main")), - mutator.getValue(laneState("main")), - mutator.getValue(laneConfig("main")), - mutator.getValue(laneLastResult("main")), - ]); - if (leaf === undefined || state === undefined) { - throw new SessionInvariantError("Session main lane has incomplete durable state"); - } - if (configuration !== undefined) return; - if (state.value.currentOperationId !== null || lastResult !== undefined) { - throw new SessionInvariantError("Configured or active main lane is missing lane.config"); - } - await mutator.commit([setValue(laneConfig("main"), seed)]); - }); +async function seedMain(session: Session, seed: LaneConfiguration, context: Context): Promise { + await session.mutate( + "main", + async (mutator) => { + const [leaf, state, configuration, lastResult] = await Promise.all([ + mutator.getValue(laneLeaf("main"), context), + mutator.getValue(laneState("main"), context), + mutator.getValue(laneConfig("main"), context), + mutator.getValue(laneLastResult("main"), context), + ]); + if (leaf === undefined || state === undefined) { + throw new SessionInvariantError("Session main lane has incomplete durable state"); + } + if (configuration !== undefined) return; + if (state.value.currentOperationId !== null || lastResult !== undefined) { + throw new SessionInvariantError("Configured or active main lane is missing lane.config"); + } + await mutator.commit([setValue(laneConfig("main"), seed)], context); + }, + context, + ); } diff --git a/packages/agent/src/harness/runtime2/lane.ts b/packages/agent/src/harness/runtime2/lane.ts index 9e8339545..4084caa36 100644 --- a/packages/agent/src/harness/runtime2/lane.ts +++ b/packages/agent/src/harness/runtime2/lane.ts @@ -1,5 +1,5 @@ -import type { Api, Model, Models } from "@earendil-works/pi-ai"; -import type { ThinkingLevel } from "../../types.ts"; +import type { Api, ImageContent, Model, Models } from "@earendil-works/pi-ai"; +import type { AgentMessage, ThinkingLevel } from "../../types.ts"; import type { AgentLane, HarnessEvent, @@ -7,6 +7,7 @@ import type { LaneExecutionInfo, SuspendedOperation, } from "../agent-harness.ts"; +import type { Context } from "../context.ts"; import { insertEntry } from "../session/commit.ts"; import { SessionPendingAssistantMessageError } from "../session/session.ts"; import type { @@ -29,8 +30,8 @@ import { } from "../session/values.ts"; import { type LaneState, SliceNotImplemented } from "./types.ts"; -type EventHandler = (event: HarnessEvent) => Promise; -type FaultHandler = (cause: unknown) => Error; +type EventHandler = (event: HarnessEvent, context: Context) => Promise; +type FaultHandler = (cause: unknown, context: Context) => Error; type Synchronous = TResult extends PromiseLike ? never : TResult; type LaneCommand = @@ -78,14 +79,14 @@ export class Lane implements AgentLane { this.sessionView = session.view(name); this.sessionTree = { ...this.sessionView, - getLeafId: () => this.getLeafId(), - setName: (name) => this.setName(name), - setLabel: (targetId, label) => this.setLabel(targetId, label), - findEntriesOnBranch: (query) => this.findEntriesOnBranch(query), - findEntryOnBranch: (query) => this.findEntryOnBranch(query), - appendMessage: (message) => this.append({ type: "message", payload: message }), - appendCustomEntry: (customType, data) => - this.append({ type: "custom", customType, ...(data === undefined ? {} : { payload: data }) }), + getLeafId: (context) => this.getLeafId(context), + setName: (name, context) => this.setName(name, context), + setLabel: (targetId, label, context) => this.setLabel(targetId, label, context), + findEntriesOnBranch: (query, context) => this.findEntriesOnBranch(query, context), + findEntryOnBranch: (query, context) => this.findEntryOnBranch(query, context), + appendMessage: (message, context) => this.append({ type: "message", payload: message }, context), + appendCustomEntry: (customType, data, context) => + this.append({ type: "custom", customType, ...(data === undefined ? {} : { payload: data }) }, context), }; this.state = state; this.onFault = onFault; @@ -93,12 +94,12 @@ export class Lane implements AgentLane { this.suspension = suspension; } - async getLeafId(): Promise { + async getLeafId(_context: Context): Promise { this.assertOpen(); return this.state.leafId; } - async getLastResult(): Promise { + async getLastResult(_context: Context): Promise { this.assertOpen(); return this.state.lastResult; } @@ -121,36 +122,47 @@ 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) => LaneCommand | Promise>, + plan: ( + state: LaneState, + reader: SessionReader, + context: Context, + ) => LaneCommand | Promise>, + context: Context, ): Promise { this.assertOpen(); let outcome: LaneCommandOutcome; try { - outcome = await this.session.mutate(this.name, async (mutator) => { - this.assertOpen(); - try { - const decision = await plan(this.state, mutator); - switch (decision.kind) { - case "return": - return { kind: "return", result: decision.result }; - case "reject": - return { kind: "reject", error: decision.error }; - case "commit": { - const commit = await mutator.commit(decision.writes); - this.state = decision.next; - if (this.suspension?.operationId !== decision.next.operation?.meta.operationId) { - this.suspension = undefined; + outcome = await this.session.mutate( + this.name, + async (mutator) => { + this.assertOpen(); + try { + const decision = await plan(this.state, mutator, context); + switch (decision.kind) { + case "return": + return { kind: "return", result: decision.result }; + case "reject": + return { kind: "reject", error: decision.error }; + case "commit": { + const commit = await mutator.commit(decision.writes, context); + this.state = decision.next; + if (this.suspension?.operationId !== decision.next.operation?.meta.operationId) { + this.suspension = undefined; + } + const result = decision.materialize(commit); + if (isPromiseLike(result)) { + throw new TypeError("Lane command materialize() must be synchronous"); + } + return { kind: "return", result }; } - const result = decision.materialize(commit); - if (isPromiseLike(result)) throw new TypeError("Lane command materialize() must be synchronous"); - return { kind: "return", result }; } + } catch (error) { + if (this.closedError !== undefined) throw this.closedError; + throw this.onFault(error, context); } - } catch (error) { - if (this.closedError !== undefined) throw this.closedError; - throw this.onFault(error); - } - }); + }, + context, + ); } catch (error) { if (this.closedError !== undefined) throw this.closedError; throw error; @@ -159,19 +171,19 @@ export class Lane implements AgentLane { return outcome.result; } - async accept(): Promise { + async accept(..._args: Parameters): Promise { throw new SliceNotImplemented("accept"); } - async drive(): Promise { + async drive(..._args: Parameters): Promise { throw new SliceNotImplemented("drive"); } - async requestAbort(): Promise { + async requestAbort(..._args: Parameters): Promise { throw new SliceNotImplemented("requestAbort"); } - async inspectExecution(): Promise { + async inspectExecution(_context: Context): Promise { this.assertOpen(); const operation = this.state.operation; if (operation === null) { @@ -201,80 +213,84 @@ export class Lane implements AgentLane { }; } - async prompt(): Promise { + async prompt( + ..._args: + | [text: string, images: ImageContent[] | undefined, context: Context] + | [message: AgentMessage | AgentMessage[], context: Context] + ): Promise { throw new SliceNotImplemented("prompt"); } - async skill(): Promise { + async skill(..._args: Parameters): Promise { throw new SliceNotImplemented("skill"); } - async promptFromTemplate(): Promise { + async promptFromTemplate(..._args: Parameters): Promise { throw new SliceNotImplemented("promptFromTemplate"); } - async compact(): Promise { + async compact(..._args: Parameters): Promise { throw new SliceNotImplemented("compact"); } - async navigateTree(): Promise { + async navigateTree(..._args: Parameters): Promise { throw new SliceNotImplemented("navigateTree"); } - async resume(): Promise { + async resume(..._args: Parameters): Promise { throw new SliceNotImplemented("resume"); } - async abort(): Promise { + async abort(..._args: Parameters): Promise { throw new SliceNotImplemented("abort"); } - async steer(): Promise { + async steer(..._args: Parameters): Promise { throw new SliceNotImplemented("steer"); } - async followUp(): Promise { + async followUp(..._args: Parameters): Promise { throw new SliceNotImplemented("followUp"); } - async nextRun(): Promise { + async nextRun(..._args: Parameters): Promise { throw new SliceNotImplemented("nextRun"); } - async cancelQueued(): Promise { + async cancelQueued(..._args: Parameters): Promise { throw new SliceNotImplemented("cancelQueued"); } - async recordUsage(): Promise { + async recordUsage(..._args: Parameters): Promise { throw new SliceNotImplemented("recordUsage"); } - async waitForIdle(): Promise { + async waitForIdle(..._args: Parameters): Promise { throw new SliceNotImplemented("waitForIdle"); } - async runWhenIdle(): Promise { + async runWhenIdle(..._args: Parameters): Promise { throw new SliceNotImplemented("runWhenIdle"); } - async peekAction(): Promise { + async peekAction(..._args: Parameters): Promise { throw new SliceNotImplemented("peekAction"); } - async executeAction(): Promise { + async executeAction(..._args: Parameters): Promise { throw new SliceNotImplemented("executeAction"); } - async runToCompletion(): Promise { + async runToCompletion(..._args: Parameters): Promise { throw new SliceNotImplemented("runToCompletion"); } - async getModel(): Promise | undefined> { + async getModel(_context: Context): Promise | undefined> { this.assertOpen(); return this.models.getModel(this.state.configuration.model.provider, this.state.configuration.model.modelId); } - setModel(model: Model): Promise { + setModel(model: Model, context: Context): Promise { return this.setConfiguration( (configuration) => ({ ...configuration, @@ -286,15 +302,16 @@ export class Lane implements AgentLane { previous: previous.model, value: value.model, }), + context, ); } - async getThinkingLevel(): Promise { + async getThinkingLevel(_context: Context): Promise { this.assertOpen(); return this.state.configuration.thinkingLevel; } - setThinkingLevel(thinkingLevel: ThinkingLevel): Promise { + setThinkingLevel(thinkingLevel: ThinkingLevel, context: Context): Promise { return this.setConfiguration( (configuration) => ({ ...configuration, thinkingLevel }), (previous, value) => ({ @@ -303,15 +320,16 @@ export class Lane implements AgentLane { previous: previous.thinkingLevel, value: value.thinkingLevel, }), + context, ); } - async getActiveTools(): Promise { + async getActiveTools(_context: Context): Promise { this.assertOpen(); return this.state.configuration.activeToolNames; } - setActiveTools(activeToolNames: string[]): Promise { + setActiveTools(activeToolNames: string[], context: Context): Promise { return this.setConfiguration( (configuration) => ({ ...configuration, activeToolNames }), (previous, value) => ({ @@ -320,16 +338,18 @@ export class Lane implements AgentLane { previous: previous.activeToolNames, value: value.activeToolNames, }), + context, ); } - async watch(): Promise { + async watch(..._args: Parameters): Promise { throw new SliceNotImplemented("watch"); } private async setConfiguration( update: (configuration: LaneState["configuration"]) => LaneState["configuration"], event: (previous: LaneState["configuration"], value: LaneState["configuration"]) => LaneConfigEventPayload, + context: Context, ): Promise { const payload = await this.command((state) => { const configuration = update(state.configuration); @@ -339,31 +359,32 @@ export class Lane implements AgentLane { next: { ...state, configuration }, materialize: () => event(state.configuration, configuration), }; - }); - await this.onEvent({ ...payload, lane: this.name }); + }, context); + await this.onEvent({ ...payload, lane: this.name }, context); } - private async setName(name: string | undefined): Promise { - await this.sessionView.setName(name); - await this.onEvent({ type: "value_update", value: "session_name", name }); + private async setName(name: string | undefined, context: Context): Promise { + await this.sessionView.setName(name, context); + await this.onEvent({ type: "value_update", value: "session_name", name }, context); } - private async setLabel(targetId: string, label: string | undefined): Promise { - await this.sessionView.setLabel(targetId, label); - await this.onEvent({ type: "value_update", value: "entry_label", targetId, label }); + private async setLabel(targetId: string, label: string | undefined, context: Context): Promise { + await this.sessionView.setLabel(targetId, label, context); + await this.onEvent({ type: "value_update", value: "entry_label", targetId, label }, context); } - private async findEntriesOnBranch(query: BranchScan = {}): Promise { + private async findEntriesOnBranch(query: BranchScan | undefined, context: Context): Promise { + query ??= {}; this.assertOpen(); const start = query.start ?? this.state.leafId; - return start === null ? [] : this.sessionView.findEntriesOnBranch({ ...query, start }); + return start === null ? [] : this.sessionView.findEntriesOnBranch({ ...query, start }, context); } - private async findEntryOnBranch(query: BranchScan = {}): Promise { - return (await this.findEntriesOnBranch({ ...query, limit: 1 }))[0]; + private async findEntryOnBranch(query: BranchScan | undefined, context: Context): Promise { + return (await this.findEntriesOnBranch({ ...query, limit: 1 }, context))[0]; } - private append(pending: PendingEntry): Promise { + private append(pending: PendingEntry, context: Context): Promise { this.assertOpen(); if ( pending.type === "message" && @@ -420,7 +441,7 @@ export class Lane implements AgentLane { next: { ...state, operation: { meta: state.operation.meta, state: operationState } }, materialize: () => id, }; - }); + }, context); } seal(error: Error): void { diff --git a/packages/agent/src/harness/runtime2/restore.ts b/packages/agent/src/harness/runtime2/restore.ts index cb9698df7..d223d5ee5 100644 --- a/packages/agent/src/harness/runtime2/restore.ts +++ b/packages/agent/src/harness/runtime2/restore.ts @@ -1,3 +1,4 @@ +import type { Context } from "../context.ts"; import { SessionInvariantError } from "../session/session.ts"; import type { Session, SessionReader } from "../session/types.ts"; import { @@ -12,28 +13,28 @@ import { import type { LaneState } from "./types.ts"; /** Restore every configured lane in one session without starting work. */ -export async function restoreSession(session: Session): Promise> { - const lanes = await session.scanValues(laneLeafInventoryPrefix()); +export async function restoreSession(session: Session, context: Context): Promise> { + const lanes = await session.scanValues(laneLeafInventoryPrefix(), context); if (!lanes.some((stored) => stored.address.key === "main")) { throw new SessionInvariantError("Session is missing main lane"); } const restored = await Promise.all( - lanes.map(async ({ address }) => ({ key: address.key, state: await restoreLane(session, address.key) })), + lanes.map(async ({ address }) => ({ key: address.key, state: await restoreLane(session, address.key, context) })), ); return new Map(restored.map(({ key, state }) => [key, state])); } /** Restore one configured lane without starting work or interpreting its state. */ -export function restoreLane(session: Session, lane: string): Promise { - return session.mutate(lane, (reader) => restoreLaneState(reader, lane)); +export function restoreLane(session: Session, lane: string, context: Context): Promise { + return session.mutate(lane, (reader) => restoreLaneState(reader, lane, context), context); } -async function restoreLaneState(reader: SessionReader, lane: string): Promise { +async function restoreLaneState(reader: SessionReader, lane: string, context: Context): Promise { const [leaf, configuration, laneState, lastResult] = await Promise.all([ - reader.getValue(laneLeaf(lane)), - reader.getValue(laneConfig(lane)), - reader.getValue(laneStateValue(lane)), - reader.getValue(laneLastResult(lane)), + reader.getValue(laneLeaf(lane), context), + reader.getValue(laneConfig(lane), context), + reader.getValue(laneStateValue(lane), context), + reader.getValue(laneLastResult(lane), context), ]); if (leaf === undefined) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.leaf`); if (configuration === undefined) @@ -44,8 +45,8 @@ async function restoreLaneState(reader: SessionReader, lane: string): Promise { + options ??= {}; const entries = buildContextEntries(pathEntries); const messages: AgentMessage[] = []; for (const entry of entries) { @@ -55,7 +58,7 @@ export async function buildSessionContext( continue; } const projector = options.entryProjectors?.[entry.customType]; - if (projector !== undefined) messages.push(...((await projector(entry)) ?? [])); + if (projector !== undefined) messages.push(...((await projector(entry, context)) ?? [])); } return messages; } diff --git a/packages/agent/src/harness/session/jsonl/repo.ts b/packages/agent/src/harness/session/jsonl/repo.ts index 3b6f254b1..70f07dc2a 100644 --- a/packages/agent/src/harness/session/jsonl/repo.ts +++ b/packages/agent/src/harness/session/jsonl/repo.ts @@ -1,4 +1,5 @@ import { uuidv7 } from "@earendil-works/pi-ai"; +import type { Context } from "../../context.ts"; import type { FileError, FileSystem, Result } from "../../types.ts"; import { createForkSnapshot, type ForkSourceSnapshot } from "../fork.ts"; import { StorageBackedSession } from "../session.ts"; @@ -64,7 +65,6 @@ export class JsonlSessionRepo private readonly now: () => number; private readonly openSessions = new Map(); private readonly pendingCreates = new Set(); - private rootPromise: Promise | undefined; private closed = false; private closePromise: Promise | undefined; @@ -74,17 +74,17 @@ export class JsonlSessionRepo this.now = options.now ?? Date.now; } - async create(options: JsonlSessionCreateOptions): Promise> { + async create(options: JsonlSessionCreateOptions, context: Context): Promise> { this.assertOpen(); const createdAt = this.now(); - const { cwd, id } = await this.resolveCreateDestination(options.cwd, options.id, createdAt); + const { cwd, id } = await this.resolveCreateDestination(options.cwd, options.id, createdAt, context); const key = this.sessionKey(cwd, id); if (this.openSessions.has(key) || this.pendingCreates.has(key)) throw new Error(`Session already exists: ${id}`); this.pendingCreates.add(key); let path: string | undefined; let storage: JsonlStorage | undefined; try { - path = await this.resolveNewSessionPath(cwd, createdAt, id); + path = await this.resolveNewSessionPath(cwd, createdAt, id, context); const header: JsonlStorageHeader = { v: JSONL_FORMAT_VERSION, kind: "header", @@ -95,15 +95,18 @@ export class JsonlSessionRepo ...(options.parentSessionId === undefined ? {} : { parentSessionId: options.parentSessionId }), }; // TODO: do we want to make session creation atomic? - storage = await JsonlStorage.create({ fileSystem: this.fileSystem, path, now: this.now }, header); - await storage.commit([ - setValue(laneLeaf("main"), null), - setValue(laneState("main"), { currentOperationId: null, pendingNextRun: [] }), - ]); + storage = await JsonlStorage.create({ fileSystem: this.fileSystem, path, now: this.now }, header, context); + await storage.commit( + [ + setValue(laneLeaf("main"), null), + setValue(laneState("main"), { currentOperationId: null, pendingNextRun: [] }), + ], + context, + ); const info = fileValue(await this.fileSystem.fileInfo(path), `Failed to read session ${path}`); return this.publishOpenSession(metadataFromHeader(header, path, info.mtimeMs), storage, key); } catch (error) { - await storage?.close().catch(() => undefined); + await storage?.close(context).catch(() => undefined); if (path !== undefined) await this.fileSystem.remove(path, { force: true }); throw error; } finally { @@ -111,21 +114,22 @@ export class JsonlSessionRepo } } - async open(metadata: JsonlSessionMetadata): Promise> { + async open(metadata: JsonlSessionMetadata, context: Context): Promise> { this.assertOpen(); const key = this.sessionKey(metadata.cwd, metadata.id); if (this.openSessions.has(key)) throw new Error(`Session is already open: ${metadata.id}`); let storage: JsonlStorage | undefined; try { - storage = await this.loadStorage(metadata); + storage = await this.loadStorage(metadata, context); return this.publishOpenSession(metadata, storage, key); } catch (error) { - await storage?.close(); + await storage?.close(context); throw error; } } - async list(options: JsonlSessionListOptions = {}): Promise { + async list(options: JsonlSessionListOptions | undefined, context: Context): Promise { + options ??= {}; this.assertOpen(); const cwd = options.cwd === undefined @@ -134,9 +138,10 @@ export class JsonlSessionRepo await this.fileSystem.absolutePath(options.cwd), `Failed to resolve session cwd ${options.cwd}`, ); - const root = await this.root(); + const root = await this.root(context); if (!fileValue(await this.fileSystem.exists(root), `Failed to check sessions root ${root}`)) return []; - const directories = cwd === undefined ? await this.sessionDirectories(root) : [await this.sessionDirectory(cwd)]; + const directories = + cwd === undefined ? await this.sessionDirectories(root) : [await this.sessionDirectory(cwd, context)]; const metadata: JsonlSessionMetadata[] = []; for (const directory of directories) metadata.push(...(await this.listDirectory(directory, cwd))); return metadata.sort( @@ -145,7 +150,7 @@ export class JsonlSessionRepo ); } - async delete(metadata: JsonlSessionMetadata): Promise { + async delete(metadata: JsonlSessionMetadata, _context: Context): Promise { this.assertOpen(); const key = this.sessionKey(metadata.cwd, metadata.id); if (this.openSessions.has(key)) throw new Error(`Session is open: ${metadata.id}`); @@ -155,14 +160,18 @@ export class JsonlSessionRepo fileValue(await this.fileSystem.remove(metadata.path), `Failed to delete session ${metadata.path}`); } - async fork(source: JsonlSessionMetadata, options: ForkOptions): Promise> { + async fork( + source: JsonlSessionMetadata, + options: ForkOptions, + context: Context, + ): Promise> { this.assertOpen(); const createdAt = this.now(); const sourceStorage = this.openSessions.get(this.sessionKey(source.cwd, source.id)); const sourceSnapshot = await (sourceStorage === undefined - ? this.loadClosedForkSourceSnapshot(source) - : sourceStorage.snapshot()); - const { cwd, id } = await this.resolveCreateDestination(source.cwd, options.id, createdAt); + ? this.loadClosedForkSourceSnapshot(source, context) + : sourceStorage.snapshot(context)); + const { cwd, id } = await this.resolveCreateDestination(source.cwd, options.id, createdAt, context); const destinationKey = this.sessionKey(cwd, id); if (this.openSessions.has(destinationKey) || this.pendingCreates.has(destinationKey)) { throw new Error(`Session already exists: ${id}`); @@ -172,7 +181,7 @@ export class JsonlSessionRepo let path: string | undefined; let storage: JsonlStorage | undefined; try { - path = await this.resolveNewSessionPath(cwd, createdAt, id); + path = await this.resolveNewSessionPath(cwd, createdAt, id, context); const snapshot = createForkSnapshot(sourceSnapshot, options); const header: JsonlStorageHeader = { v: JSONL_FORMAT_VERSION, @@ -187,11 +196,12 @@ export class JsonlSessionRepo { fileSystem: this.fileSystem, path, now: this.now }, header, snapshot, + context, ); const info = fileValue(await this.fileSystem.fileInfo(path), `Failed to read session ${path}`); return this.publishOpenSession(metadataFromHeader(header, path, info.mtimeMs), storage, destinationKey); } catch (error) { - await storage?.close().catch(() => undefined); + await storage?.close(context).catch(() => undefined); if (path !== undefined) await this.fileSystem.remove(path, { force: true }); throw error; } finally { @@ -199,7 +209,7 @@ export class JsonlSessionRepo } } - close(): Promise { + close(_context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; this.closed = true; // TODO: Define ownership semantics before deciding whether repository close should close session handles. @@ -211,6 +221,7 @@ export class JsonlSessionRepo cwdInput: string, id: string | undefined, createdAt: number, + _context: Context, ): Promise<{ cwd: string; id: string }> { const destinationId = id ?? uuidv7(createdAt); const cwd = fileValue(await this.fileSystem.absolutePath(cwdInput), `Failed to resolve session cwd ${cwdInput}`); @@ -247,15 +258,15 @@ export class JsonlSessionRepo .map((entry) => entry.path); } - private async sessionDirectory(cwd: string): Promise { + private async sessionDirectory(cwd: string, context: Context): Promise { return fileValue( - await this.fileSystem.joinPath([await this.root(), sessionDirectoryName(cwd)]), + await this.fileSystem.joinPath([await this.root(context), sessionDirectoryName(cwd)]), `Failed to resolve sessions directory for ${cwd}`, ); } - private async resolveNewSessionPath(cwd: string, createdAt: number, id: string): Promise { - const directory = await this.sessionDirectory(cwd); + private async resolveNewSessionPath(cwd: string, createdAt: number, id: string, context: Context): Promise { + const directory = await this.sessionDirectory(cwd, context); await this.assertSessionIdAvailable(directory, id); fileValue(await this.fileSystem.createDir(directory), `Failed to create sessions directory ${directory}`); return fileValue( @@ -276,12 +287,15 @@ export class JsonlSessionRepo if (idExists) throw new Error(`Session already exists: ${id}`); } - private async loadClosedForkSourceSnapshot(source: JsonlSessionMetadata): Promise { - const storage = await this.loadStorage(source); + private async loadClosedForkSourceSnapshot( + source: JsonlSessionMetadata, + context: Context, + ): Promise { + const storage = await this.loadStorage(source, context); try { - return await storage.snapshot(); + return await storage.snapshot(context); } finally { - await storage.close(); + await storage.close(context); } } @@ -304,16 +318,19 @@ export class JsonlSessionRepo return `${cwd}\0${id}`; } - private async loadStorage(metadata: JsonlSessionMetadata): Promise { + private async loadStorage(metadata: JsonlSessionMetadata, context: Context): Promise { // TODO: normalize discovered legacy v3 sessions on open instead of rejecting them here. if (!fileValue(await this.fileSystem.exists(metadata.path), `Failed to check session ${metadata.path}`)) { throw new Error(`Session file does not exist: ${metadata.path}`); } - const storage = await JsonlStorage.open({ - fileSystem: this.fileSystem, - path: metadata.path, - now: this.now, - }); + const storage = await JsonlStorage.open( + { + fileSystem: this.fileSystem, + path: metadata.path, + now: this.now, + }, + context, + ); try { if (storage.header.id !== metadata.id || storage.header.cwd !== metadata.cwd) { throw new Error(`Session identity does not match header: ${metadata.id}`); @@ -323,16 +340,15 @@ export class JsonlSessionRepo } return storage; } catch (error) { - await storage.close(); + await storage.close(context); throw error; } } - private root(): Promise { - this.rootPromise ??= this.fileSystem + private root(_context: Context): Promise { + return this.fileSystem .absolutePath(this.sessionsRootInput) .then((result) => fileValue(result, `Failed to resolve sessions root ${this.sessionsRootInput}`)); - return this.rootPromise; } private assertOpen(): void { diff --git a/packages/agent/src/harness/session/jsonl/storage.ts b/packages/agent/src/harness/session/jsonl/storage.ts index f84226761..fbc933b84 100644 --- a/packages/agent/src/harness/session/jsonl/storage.ts +++ b/packages/agent/src/harness/session/jsonl/storage.ts @@ -1,3 +1,4 @@ +import type { Context } from "../../context.ts"; import type { FileError, FileSystem, Result } from "../../types.ts"; import { type CommittedEntryWrite, @@ -140,7 +141,11 @@ export class JsonlStorage implements Storage { this.header = header; } - static async create(options: JsonlStorageOptions, header: JsonlStorageHeader): Promise { + static async create( + options: JsonlStorageOptions, + header: JsonlStorageHeader, + _context: Context, + ): Promise { fileValue( await options.fileSystem.writeFile(options.path, `${JSON.stringify(header)}\n`), `Failed to create JSONL storage ${options.path}`, @@ -153,15 +158,16 @@ export class JsonlStorage implements Storage { options: JsonlStorageOptions, header: JsonlStorageHeader, snapshot: JsonlSnapshotContents, + context: Context, ): Promise { const writes = snapshotWrites(snapshot); const snapshotHeader = { ...header, nextSeq: snapshot.nextSeq }; const content = `${[JSON.stringify(snapshotHeader), ...writes.map((write) => JSON.stringify(write))].join("\n")}\n`; await publishFileAtomically(options.fileSystem, options.path, content); - return JsonlStorage.open(options); + return JsonlStorage.open(options, context); } - static async open(options: JsonlStorageOptions): Promise { + static async open(options: JsonlStorageOptions, _context: Context): Promise { const content = fileValue( await options.fileSystem.readTextFile(options.path), `Failed to read JSONL storage ${options.path}`, @@ -187,7 +193,7 @@ export class JsonlStorage implements Storage { return storage; } - async commit(writes: Write[]): Promise { + async commit(writes: Write[], _context: Context): Promise { if (this.state !== "open") throw new Error("JsonlStorage is closed"); const result = this.commitQueue.then(() => this.applyCommit(writes)); this.commitQueue = result.then( @@ -212,53 +218,57 @@ export class JsonlStorage implements Storage { return prepared.result; } - getEntries(ids: string[]): Promise> { + getEntries(ids: string[], _context: Context): Promise> { if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed")); return Promise.resolve(this.storageState.getEntries(ids)); } - getValue(address: Value): Promise | undefined> { + getValue(address: Value, _context: Context): Promise | undefined> { if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed")); return Promise.resolve(this.storageState.getValue(address)); } - scanValues(prefix: Value): Promise[]> { + scanValues(prefix: Value, _context: Context): Promise[]> { if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed")); return Promise.resolve(this.storageState.scanValues(prefix)); } - async readList(address: ValueList, options?: ListReadOptions): Promise[]> { + async readList( + address: ValueList, + options: ListReadOptions | undefined, + _context: Context, + ): Promise[]> { if (this.state !== "open") throw new Error("JsonlStorage is closed"); return this.storageState.readList(address, options); } - async scanBranch(query: StorageBranchScan): Promise { + async scanBranch(query: StorageBranchScan, _context: Context): Promise { if (this.state !== "open") throw new Error("JsonlStorage is closed"); return this.storageState.scanBranch(query); } - async scanBranchStructure(query: StorageBranchScan): Promise { + async scanBranchStructure(query: StorageBranchScan, _context: Context): Promise { if (this.state !== "open") throw new Error("JsonlStorage is closed"); return this.storageState.scanBranchStructure(query); } - scanEntries(query: EntryScan): Promise { + scanEntries(query: EntryScan, _context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed")); return Promise.resolve(this.storageState.scanEntries(query)); } - scanUsage(query: UsageScan): Promise { + scanUsage(query: UsageScan, _context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed")); return Promise.resolve(this.storageState.scanUsage(query)); } - getStats(): Promise { + getStats(_context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed")); return Promise.resolve(this.storageState.getStats()); } /** Capture the current entries and values at one serialized boundary between commits. */ - snapshot(): Promise<{ + snapshot(_context: Context): Promise<{ entries: Entry[]; scalarValues: StoredValue[]; listValues: StorageStateSnapshot["listValues"]; @@ -279,7 +289,7 @@ export class JsonlStorage implements Storage { return result; } - close(): Promise { + close(_context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; this.state = "closing"; this.closePromise = this.commitQueue.then(() => { diff --git a/packages/agent/src/harness/session/memory.ts b/packages/agent/src/harness/session/memory.ts index fb730f570..933f52a7a 100644 --- a/packages/agent/src/harness/session/memory.ts +++ b/packages/agent/src/harness/session/memory.ts @@ -1,5 +1,6 @@ import { uuidv7 } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; +import type { Context } from "../context.ts"; import { createForkSnapshot } from "./fork.ts"; import { StorageBackedSession } from "./session.ts"; import { StorageState, type StorageStateSnapshot } from "./storage-state.ts"; @@ -57,7 +58,7 @@ export class MemoryStorage implements Storage { this.now = options.now ?? Date.now; } - async commit(writes: Write[]): Promise { + async commit(writes: Write[], _context: Context): Promise { if (this.state !== "open") throw new Error("MemoryStorage is closed"); const result = this.commitQueue.then(() => { const prepared = this.storageState.prepareCommit(writes, this.now()); @@ -71,53 +72,57 @@ export class MemoryStorage implements Storage { return result; } - getEntries(ids: string[]): Promise> { + getEntries(ids: string[], _context: Context): Promise> { if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed")); return Promise.resolve(this.storageState.getEntries(ids)); } - getValue(address: Value): Promise | undefined> { + getValue(address: Value, _context: Context): Promise | undefined> { if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed")); return Promise.resolve(this.storageState.getValue(address)); } - scanValues(prefix: Value): Promise[]> { + scanValues(prefix: Value, _context: Context): Promise[]> { if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed")); return Promise.resolve(this.storageState.scanValues(prefix)); } - async readList(address: ValueList, options?: ListReadOptions): Promise[]> { + async readList( + address: ValueList, + options: ListReadOptions | undefined, + _context: Context, + ): Promise[]> { if (this.state !== "open") throw new Error("MemoryStorage is closed"); return this.storageState.readList(address, options); } - async scanBranch(query: StorageBranchScan): Promise { + async scanBranch(query: StorageBranchScan, _context: Context): Promise { if (this.state !== "open") throw new Error("MemoryStorage is closed"); return this.storageState.scanBranch(query); } - async scanBranchStructure(query: StorageBranchScan): Promise { + async scanBranchStructure(query: StorageBranchScan, _context: Context): Promise { if (this.state !== "open") throw new Error("MemoryStorage is closed"); return this.storageState.scanBranchStructure(query); } - scanEntries(query: EntryScan): Promise { + scanEntries(query: EntryScan, _context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed")); return Promise.resolve(this.storageState.scanEntries(query)); } - scanUsage(query: UsageScan): Promise { + scanUsage(query: UsageScan, _context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed")); return Promise.resolve(this.storageState.scanUsage(query)); } - getStats(): Promise { + getStats(_context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed")); return Promise.resolve(this.storageState.getStats()); } /** Capture the current stores at one serialized boundary between commits. */ - snapshot(): Promise<{ + snapshot(_context: Context): Promise<{ entries: Entry[]; scalarValues: StoredValue[]; listValues: StorageStateSnapshot["listValues"]; @@ -138,7 +143,7 @@ export class MemoryStorage implements Storage { return result; } - close(): Promise { + close(_context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; this.state = "closing"; this.closePromise = this.commitQueue.then(() => { @@ -180,128 +185,142 @@ class MemorySessionFacade implements Session { this.onClose = onClose; } - async mutate(lane: string, mutation: (mutator: SessionMutator) => T | Promise): Promise { - return this.admit(() => this.session.mutate(lane, mutation)); + async mutate( + lane: string, + mutation: (mutator: SessionMutator, context: Context) => T | Promise, + context: Context, + ): Promise { + return this.admit(() => this.session.mutate(lane, mutation, context)); } - getEntries(ids: string[]): Promise> { - return this.admit(() => this.session.getEntries(ids)); + getEntries(ids: string[], context: Context): Promise> { + return this.admit(() => this.session.getEntries(ids, context)); } - getValue(address: Value): Promise | undefined> { - return this.admit(() => this.session.getValue(address)); + getValue(address: Value, context: Context): Promise | undefined> { + return this.admit(() => this.session.getValue(address, context)); } - scanValues(prefix: Value): Promise[]> { - return this.admit(() => this.session.scanValues(prefix)); + scanValues(prefix: Value, context: Context): Promise[]> { + return this.admit(() => this.session.scanValues(prefix, context)); } - readList(address: ValueList, options?: ListReadOptions): Promise[]> { - return this.admit(() => this.session.readList(address, options)); + readList( + address: ValueList, + options: ListReadOptions | undefined, + context: Context, + ): Promise[]> { + return this.admit(() => this.session.readList(address, options, context)); } view(lane: string): SessionTree { const view = this.session.view(lane); return { - getLeafId: () => this.admit(() => view.getLeafId()), - getEntry: (id) => this.admit(() => view.getEntry(id)), - getStats: () => this.admit(() => view.getStats()), - getValue: (address) => this.admit(() => view.getValue(address)), - scanValues: (prefix) => this.admit(() => view.scanValues(prefix)), - readList: (address, options) => this.admit(() => view.readList(address, options)), - setValue: (address, next) => this.admit(() => view.setValue(address, next)), - deleteValue: (address) => this.admit(() => view.deleteValue(address)), - appendList: (address, element) => this.admit(() => view.appendList(address, element)), - deleteList: (address) => this.admit(() => view.deleteList(address)), - getName: () => this.admit(() => view.getName()), - setName: (name) => this.admit(() => view.setName(name)), - getLabel: (targetId) => this.admit(() => view.getLabel(targetId)), - setLabel: (targetId, label) => this.admit(() => view.setLabel(targetId, label)), - findEntries: (query) => this.admit(() => view.findEntries(query)), - findEntry: (query) => this.admit(() => view.findEntry(query)), - findEntriesOnBranch: (query) => this.admit(() => view.findEntriesOnBranch(query)), - findEntryOnBranch: (query) => this.admit(() => view.findEntryOnBranch(query)), - appendMessage: (message) => this.admit(() => view.appendMessage(message)), - appendCustomEntry: (customType, data) => this.admit(() => view.appendCustomEntry(customType, data)), + getLeafId: (context) => this.admit(() => view.getLeafId(context)), + getEntry: (id, context) => this.admit(() => view.getEntry(id, context)), + getStats: (context) => this.admit(() => view.getStats(context)), + getValue: (address, context) => this.admit(() => view.getValue(address, context)), + scanValues: (prefix, context) => this.admit(() => view.scanValues(prefix, context)), + readList: (address, options, context) => this.admit(() => view.readList(address, options, context)), + setValue: (address, next, context) => this.admit(() => view.setValue(address, next, context)), + deleteValue: (address, context) => this.admit(() => view.deleteValue(address, context)), + appendList: (address, element, context) => this.admit(() => view.appendList(address, element, context)), + deleteList: (address, context) => this.admit(() => view.deleteList(address, context)), + getName: (context) => this.admit(() => view.getName(context)), + setName: (name, context) => this.admit(() => view.setName(name, context)), + getLabel: (targetId, context) => this.admit(() => view.getLabel(targetId, context)), + setLabel: (targetId, label, context) => this.admit(() => view.setLabel(targetId, label, context)), + findEntries: (query, context) => this.admit(() => view.findEntries(query, context)), + findEntry: (query, context) => this.admit(() => view.findEntry(query, context)), + findEntriesOnBranch: (query, context) => this.admit(() => view.findEntriesOnBranch(query, context)), + findEntryOnBranch: (query, context) => this.admit(() => view.findEntryOnBranch(query, context)), + appendMessage: (message, context) => this.admit(() => view.appendMessage(message, context)), + appendCustomEntry: (customType, data, context) => + this.admit(() => view.appendCustomEntry(customType, data, context)), }; } - createLane(name: string, at: string | null, configuration: LaneConfiguration): Promise { + createLane( + name: string, + at: string | null, + configuration: LaneConfiguration, + context: Context, + ): Promise { return this.admit(async () => { - await this.session.createLane(name, at, configuration); + await this.session.createLane(name, at, configuration, context); return this.view(name); }); } - getLeafId(): Promise { - return this.admit(() => this.session.getLeafId()); + getLeafId(context: Context): Promise { + return this.admit(() => this.session.getLeafId(context)); } - getEntry(id: string): Promise { - return this.admit(() => this.session.getEntry(id)); + getEntry(id: string, context: Context): Promise { + return this.admit(() => this.session.getEntry(id, context)); } - getStats(): Promise { - return this.admit(() => this.session.getStats()); + getStats(context: Context): Promise { + return this.admit(() => this.session.getStats(context)); } - setValue(address: Value, next: NoInfer): Promise { - return this.admit(() => this.session.setValue(address, next)); + setValue(address: Value, next: NoInfer, context: Context): Promise { + return this.admit(() => this.session.setValue(address, next, context)); } - deleteValue(address: Value): Promise { - return this.admit(() => this.session.deleteValue(address)); + deleteValue(address: Value, context: Context): Promise { + return this.admit(() => this.session.deleteValue(address, context)); } - appendList(address: ValueList, element: NoInfer): Promise { - return this.admit(() => this.session.appendList(address, element)); + appendList(address: ValueList, element: NoInfer, context: Context): Promise { + return this.admit(() => this.session.appendList(address, element, context)); } - deleteList(address: ValueList): Promise { - return this.admit(() => this.session.deleteList(address)); + deleteList(address: ValueList, context: Context): Promise { + return this.admit(() => this.session.deleteList(address, context)); } - getName(): Promise { - return this.admit(() => this.session.getName()); + getName(context: Context): Promise { + return this.admit(() => this.session.getName(context)); } - setName(name: string | undefined): Promise { - return this.admit(() => this.session.setName(name)); + setName(name: string | undefined, context: Context): Promise { + return this.admit(() => this.session.setName(name, context)); } - getLabel(targetId: string): Promise { - return this.admit(() => this.session.getLabel(targetId)); + getLabel(targetId: string, context: Context): Promise { + return this.admit(() => this.session.getLabel(targetId, context)); } - setLabel(targetId: string, label: string | undefined): Promise { - return this.admit(() => this.session.setLabel(targetId, label)); + setLabel(targetId: string, label: string | undefined, context: Context): Promise { + return this.admit(() => this.session.setLabel(targetId, label, context)); } - findEntries(query?: EntryQuery): Promise { - return this.admit(() => this.session.findEntries(query)); + findEntries(query: EntryQuery | undefined, context: Context): Promise { + return this.admit(() => this.session.findEntries(query, context)); } - findEntry(query?: EntryQuery): Promise { - return this.admit(() => this.session.findEntry(query)); + findEntry(query: EntryQuery | undefined, context: Context): Promise { + return this.admit(() => this.session.findEntry(query, context)); } - findEntriesOnBranch(query?: BranchScan): Promise { - return this.admit(() => this.session.findEntriesOnBranch(query)); + findEntriesOnBranch(query: BranchScan | undefined, context: Context): Promise { + return this.admit(() => this.session.findEntriesOnBranch(query, context)); } - findEntryOnBranch(query?: BranchScan): Promise { - return this.admit(() => this.session.findEntryOnBranch(query)); + findEntryOnBranch(query: BranchScan | undefined, context: Context): Promise { + return this.admit(() => this.session.findEntryOnBranch(query, context)); } - appendMessage(message: AgentMessage): Promise { - return this.admit(() => this.session.appendMessage(message)); + appendMessage(message: AgentMessage, context: Context): Promise { + return this.admit(() => this.session.appendMessage(message, context)); } - appendCustomEntry(customType: string, data?: JsonValue): Promise { - return this.admit(() => this.session.appendCustomEntry(customType, data)); + appendCustomEntry(customType: string, data: JsonValue | undefined, context: Context): Promise { + return this.admit(() => this.session.appendCustomEntry(customType, data, context)); } - close(): Promise { + close(_context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; this.state = "closing"; this.closePromise = Promise.allSettled([...this.admitted]).then(() => { @@ -339,7 +358,7 @@ export class MemorySessionRepo implements SessionRepo { this.now = options.now ?? Date.now; } - async create(options: SessionCreateOptions): Promise { + async create(options: SessionCreateOptions, context: Context): Promise { this.assertOpen(); const createdAt = this.now(); const id = options.id ?? uuidv7(createdAt); @@ -353,24 +372,30 @@ export class MemorySessionRepo implements SessionRepo { const storage = new MemoryStorage({ now: this.now }); const session = new StorageBackedSession(metadata, storage); try { - await session.mutate("main", (mutator) => - mutator.commit([ - setValueWrite(laneLeaf("main"), null), - setValueWrite(laneState("main"), { currentOperationId: null, pendingNextRun: [] }), - ]), + await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + setValueWrite(laneLeaf("main"), null), + setValueWrite(laneState("main"), { currentOperationId: null, pendingNextRun: [] }), + ], + context, + ), + context, ); const record: MemorySessionRecord = { metadata, storage, session, open: true }; this.sessions.set(id, record); return this.openRecord(record); } catch (error) { - await session.close(); + await session.close(context); throw error; } finally { this.pendingIds.delete(id); } } - open(metadata: SessionMetadata): Promise { + open(metadata: SessionMetadata, _context: Context): Promise { // Memory sessions are always created at the current storage version, so // persistent-backend version gating does not apply here. this.assertOpen(); @@ -381,21 +406,21 @@ export class MemorySessionRepo implements SessionRepo { return Promise.resolve(this.openRecord(record)); } - list(): Promise { + list(_options: undefined, _context: Context): Promise { this.assertOpen(); return Promise.resolve([...this.sessions.values()].map(({ metadata }) => metadata)); } - async delete(metadata: SessionMetadata): Promise { + async delete(metadata: SessionMetadata, context: Context): Promise { this.assertOpen(); const record = this.sessions.get(metadata.id); if (record === undefined) throw new Error(`Unknown session: ${metadata.id}`); if (record.open) throw new Error(`Session is open: ${metadata.id}`); - await record.session.close(); + await record.session.close(context); this.sessions.delete(metadata.id); } - async fork(source: SessionMetadata, options: ForkOptions): Promise { + async fork(source: SessionMetadata, options: ForkOptions, context: Context): Promise { this.assertOpen(); const sourceRecord = this.sessions.get(source.id); if (sourceRecord === undefined) throw new Error(`Unknown session: ${source.id}`); @@ -404,7 +429,7 @@ export class MemorySessionRepo implements SessionRepo { this.reserveId(id); try { - const snapshot = createForkSnapshot(await sourceRecord.storage.snapshot(), options); + const snapshot = createForkSnapshot(await sourceRecord.storage.snapshot(context), options); const storage = MemoryStorage.fromSnapshot({ now: this.now }, snapshot); const metadata: SessionMetadata = { id, @@ -421,10 +446,10 @@ export class MemorySessionRepo implements SessionRepo { } } - close(): Promise { + close(context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; this.closed = true; - this.closePromise = Promise.all([...this.sessions.values()].map(({ session }) => session.close())).then( + this.closePromise = Promise.all([...this.sessions.values()].map(({ session }) => session.close(context))).then( () => undefined, ); return this.closePromise; diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index 871e37739..68d67e895 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -1,5 +1,6 @@ import { uuidv7 } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; +import type { Context } from "../context.ts"; import { insertEntry } from "./commit.ts"; import { LaneMutationLine } from "./lane-mutations.ts"; import type { @@ -109,7 +110,7 @@ class StorageBackedSessionMutator implements SessionMutator { this.storage = storage; } - commit(writes: Write[]): Promise { + commit(writes: Write[], context: Context): Promise { this.assertActive(); if (this.commitResult !== undefined) return Promise.reject(new Error("SessionMutator commit already attempted")); try { @@ -123,31 +124,35 @@ class StorageBackedSessionMutator implements SessionMutator { throw new SessionPendingAssistantMessageError(); } } - this.commitResult = this.storage.commit(writes); + this.commitResult = this.storage.commit(writes, context); } catch (error) { this.commitResult = Promise.reject(error); } return this.commitResult; } - getEntries(ids: string[]): Promise> { + getEntries(ids: string[], context: Context): Promise> { this.assertActive(); - return this.storage.getEntries(ids); + return this.storage.getEntries(ids, context); } - getValue(address: Value): Promise | undefined> { + getValue(address: Value, context: Context): Promise | undefined> { this.assertActive(); - return this.storage.getValue(address); + return this.storage.getValue(address, context); } - scanValues(prefix: Value): Promise[]> { + scanValues(prefix: Value, context: Context): Promise[]> { this.assertActive(); - return this.storage.scanValues(prefix); + return this.storage.scanValues(prefix, context); } - readList(address: ValueList, options?: ListReadOptions): Promise[]> { + readList( + address: ValueList, + options: ListReadOptions | undefined, + context: Context, + ): Promise[]> { this.assertActive(); - return this.storage.readList(address, options); + return this.storage.readList(address, options, context); } settle(): Promise { @@ -186,13 +191,17 @@ export class StorageBackedSession(lane: string, mutation: (mutator: SessionMutator) => T | Promise): Promise { + async mutate( + lane: string, + mutation: (mutator: SessionMutator, context: Context) => T | Promise, + context: Context, + ): Promise { this.assertOpen(); return this.laneMutationLine.run(lane, async () => { const mutator = new StorageBackedSessionMutator(lane, this.storage); try { try { - return await mutation(mutator); + return await mutation(mutator, context); } finally { await mutator.settle(); } @@ -202,188 +211,219 @@ export class StorageBackedSession> { + async getEntries(ids: string[], context: Context): Promise> { this.assertOpen(); - return this.storage.getEntries(ids); + return this.storage.getEntries(ids, context); } - async getValue(address: Value): Promise | undefined> { + async getValue(address: Value, context: Context): Promise | undefined> { this.assertOpen(); - return this.storage.getValue(address); + return this.storage.getValue(address, context); } - async scanValues(prefix: Value): Promise[]> { + async scanValues(prefix: Value, context: Context): Promise[]> { this.assertOpen(); - return this.storage.scanValues(prefix); + return this.storage.scanValues(prefix, context); } - async readList(address: ValueList, options?: ListReadOptions): Promise[]> { + async readList( + address: ValueList, + options: ListReadOptions | undefined, + context: Context, + ): Promise[]> { this.assertOpen(); - return this.storage.readList(address, options); + return this.storage.readList(address, options, context); } view(lane: string): SessionTree { return { - getLeafId: () => this.getLeafIdForLane(lane), - getEntry: (id) => this.getEntry(id), - getStats: () => this.getStats(), - getValue: (address) => this.getValue(address), - scanValues: (prefix) => this.scanValues(prefix), - readList: (address, options) => this.readList(address, options), - setValue: (address, next) => this.setValueForLane(lane, address, next), - deleteValue: (address) => this.deleteValueForLane(lane, address), - appendList: (address, element) => this.appendListForLane(lane, address, element), - deleteList: (address) => this.deleteListForLane(lane, address), - getName: () => this.getName(), - setName: (name) => this.setNameForLane(lane, name), - getLabel: (targetId) => this.getLabel(targetId), - setLabel: (targetId, label) => this.setLabelForLane(lane, targetId, label), - findEntries: (query) => this.findEntries(query), - findEntry: (query) => this.findEntry(query), - findEntriesOnBranch: (query) => this.findEntriesOnBranchForLane(lane, query), - findEntryOnBranch: (query) => this.findEntryOnBranchForLane(lane, query), - appendMessage: (message) => this.appendMessageForLane(lane, message), - appendCustomEntry: (customType, data) => this.appendCustomEntryForLane(lane, customType, data), + getLeafId: (context) => this.getLeafIdForLane(lane, context), + getEntry: (id, context) => this.getEntry(id, context), + getStats: (context) => this.getStats(context), + getValue: (address, context) => this.getValue(address, context), + scanValues: (prefix, context) => this.scanValues(prefix, context), + readList: (address, options, context) => this.readList(address, options, context), + setValue: (address, next, context) => this.setValueForLane(lane, address, next, context), + deleteValue: (address, context) => this.deleteValueForLane(lane, address, context), + appendList: (address, element, context) => this.appendListForLane(lane, address, element, context), + deleteList: (address, context) => this.deleteListForLane(lane, address, context), + getName: (context) => this.getName(context), + setName: (name, context) => this.setNameForLane(lane, name, context), + getLabel: (targetId, context) => this.getLabel(targetId, context), + setLabel: (targetId, label, context) => this.setLabelForLane(lane, targetId, label, context), + findEntries: (query, context) => this.findEntries(query, context), + findEntry: (query, context) => this.findEntry(query, context), + findEntriesOnBranch: (query, context) => this.findEntriesOnBranchForLane(lane, query, context), + findEntryOnBranch: (query, context) => this.findEntryOnBranchForLane(lane, query, context), + appendMessage: (message, context) => this.appendMessageForLane(lane, message, context), + appendCustomEntry: (customType, data, context) => + this.appendCustomEntryForLane(lane, customType, data, context), }; } - async createLane(name: string, at: string | null, configuration: LaneConfiguration): Promise { + async createLane( + name: string, + at: string | null, + configuration: LaneConfiguration, + context: Context, + ): Promise { this.assertOpen(); if (name.length === 0) throw new SessionInvalidLaneError(name, "lane name must not be empty"); - return this.mutate(name, async (mutator) => { - // R1 owns complete idle-lane and current-state validation. Slice 2 only - // distinguishes valid existing lane shapes from partial durable lane state. - const [leaf, storedConfiguration, storedLaneState, lastResult] = await Promise.all([ - mutator.getValue(laneLeaf(name)), - mutator.getValue(laneConfig(name)), - mutator.getValue(laneState(name)), - mutator.getValue(laneLastResult(name)), - ]); - const presentCount = [leaf, storedConfiguration, storedLaneState, lastResult].filter( - (stored) => stored !== undefined, - ).length; - if ( - leaf !== undefined && - storedLaneState !== undefined && - (storedConfiguration !== undefined || (name === "main" && lastResult === undefined)) - ) { - throw new SessionLaneExistsError(name); - } - if (presentCount !== 0) { - throw new SessionInvariantError(`Lane ${JSON.stringify(name)} has incomplete durable state`); - } - if (at !== null && !(await mutator.getEntries([at])).has(at)) throw new SessionUnknownTargetError(at); + return this.mutate( + name, + async (mutator) => { + // R1 owns complete idle-lane and current-state validation. Slice 2 only + // distinguishes valid existing lane shapes from partial durable lane state. + const [leaf, storedConfiguration, storedLaneState, lastResult] = await Promise.all([ + mutator.getValue(laneLeaf(name), context), + mutator.getValue(laneConfig(name), context), + mutator.getValue(laneState(name), context), + mutator.getValue(laneLastResult(name), context), + ]); + const presentCount = [leaf, storedConfiguration, storedLaneState, lastResult].filter( + (stored) => stored !== undefined, + ).length; + if ( + leaf !== undefined && + storedLaneState !== undefined && + (storedConfiguration !== undefined || (name === "main" && lastResult === undefined)) + ) { + throw new SessionLaneExistsError(name); + } + if (presentCount !== 0) { + throw new SessionInvariantError(`Lane ${JSON.stringify(name)} has incomplete durable state`); + } + if (at !== null && !(await mutator.getEntries([at], context)).has(at)) { + throw new SessionUnknownTargetError(at); + } - // R6 adds the harness-wide admission barrier. Until then, close may reject - // this lane job before Storage.commit admits it; admitted commits still drain. - await mutator.commit([ - setValueWrite(laneConfig(name), configuration), - setValueWrite(laneLeaf(name), at), - setValueWrite(laneState(name), { currentOperationId: null, pendingNextRun: [] }), - ]); - return this.view(name); - }); + // R6 adds the harness-wide admission barrier. Until then, close may reject + // this lane job before Storage.commit admits it; admitted commits still drain. + await mutator.commit( + [ + setValueWrite(laneConfig(name), configuration), + setValueWrite(laneLeaf(name), at), + setValueWrite(laneState(name), { currentOperationId: null, pendingNextRun: [] }), + ], + context, + ); + return this.view(name); + }, + context, + ); } - getLeafId(): Promise { - return this.getLeafIdForLane("main"); + getLeafId(context: Context): Promise { + return this.getLeafIdForLane("main", context); } - async getEntry(id: string): Promise { - return (await this.getEntries([id])).get(id); + async getEntry(id: string, context: Context): Promise { + return (await this.getEntries([id], context)).get(id); } - async getStats(): Promise { + async getStats(context: Context): Promise { this.assertOpen(); - return this.storage.getStats(); + return this.storage.getStats(context); } - setValue(address: Value, next: NoInfer): Promise { - return this.setValueForLane("main", address, next); + setValue(address: Value, next: NoInfer, context: Context): Promise { + return this.setValueForLane("main", address, next, context); } - deleteValue(address: Value): Promise { - return this.deleteValueForLane("main", address); + deleteValue(address: Value, context: Context): Promise { + return this.deleteValueForLane("main", address, context); } - appendList(address: ValueList, element: NoInfer): Promise { - return this.appendListForLane("main", address, element); + appendList(address: ValueList, element: NoInfer, context: Context): Promise { + return this.appendListForLane("main", address, element, context); } - deleteList(address: ValueList): Promise { - return this.deleteListForLane("main", address); + deleteList(address: ValueList, context: Context): Promise { + return this.deleteListForLane("main", address, context); } - async getName(): Promise { - return (await this.getValue(sessionName))?.value; + async getName(context: Context): Promise { + return (await this.getValue(sessionName, context))?.value; } - setName(name: string | undefined): Promise { - return this.setNameForLane("main", name); + setName(name: string | undefined, context: Context): Promise { + return this.setNameForLane("main", name, context); } - async getLabel(targetId: string): Promise { - return (await this.getValue(entryLabel(targetId)))?.value; + async getLabel(targetId: string, context: Context): Promise { + return (await this.getValue(entryLabel(targetId), context))?.value; } - setLabel(targetId: string, label: string | undefined): Promise { - return this.setLabelForLane("main", targetId, label); + setLabel(targetId: string, label: string | undefined, context: Context): Promise { + return this.setLabelForLane("main", targetId, label, context); } - async findEntries(query: EntryQuery = {}): Promise { + async findEntries(query: EntryQuery | undefined, context: Context): Promise { + query ??= {}; this.assertOpen(); const order = query.order ?? "desc"; if (query.cursor !== undefined) { if (order === "asc" && query.cursor.seq === Number.MAX_SAFE_INTEGER) return []; if (order === "desc" && query.cursor.seq <= 1) return []; } - return this.storage.scanEntries({ - type: query.type, - customType: query.customType, - order, - limit: query.limit, - ...(query.cursor === undefined - ? {} - : order === "asc" - ? { fromSeq: query.cursor.seq + 1 } - : { toSeq: query.cursor.seq - 1 }), - }); + return this.storage.scanEntries( + { + type: query.type, + customType: query.customType, + order, + limit: query.limit, + ...(query.cursor === undefined + ? {} + : order === "asc" + ? { fromSeq: query.cursor.seq + 1 } + : { toSeq: query.cursor.seq - 1 }), + }, + context, + ); } - async findEntry(query: EntryQuery = {}): Promise { - const entries = await this.findEntries({ - ...query, - limit: query.limit === undefined ? 1 : Math.min(query.limit, 1), - }); + async findEntry(query: EntryQuery | undefined, context: Context): Promise { + query ??= {}; + const entries = await this.findEntries( + { + ...query, + limit: query.limit === undefined ? 1 : Math.min(query.limit, 1), + }, + context, + ); return entries[0]; } - findEntriesOnBranch(query: BranchScan = {}): Promise { - return this.findEntriesOnBranchForLane("main", query); + findEntriesOnBranch(query: BranchScan | undefined, context: Context): Promise { + return this.findEntriesOnBranchForLane("main", query, context); } - findEntryOnBranch(query: BranchScan = {}): Promise { - return this.findEntryOnBranchForLane("main", query); + findEntryOnBranch(query: BranchScan | undefined, context: Context): Promise { + return this.findEntryOnBranchForLane("main", query, context); } - appendMessage(message: AgentMessage): Promise { - return this.captureAppend("main", { type: "message", payload: message }); + appendMessage(message: AgentMessage, context: Context): Promise { + return this.captureAppend("main", { type: "message", payload: message }, context); } - appendCustomEntry(customType: string, data?: JsonValue): Promise { - return this.captureAppend("main", { - type: "custom", - customType, - ...(data === undefined ? {} : { payload: data }), - }); + appendCustomEntry(customType: string, data: JsonValue | undefined, context: Context): Promise { + return this.captureAppend( + "main", + { + type: "custom", + customType, + ...(data === undefined ? {} : { payload: data }), + }, + context, + ); } - close(): Promise { + close(context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; this.state = "closing"; this.closePromise = this.laneMutationLine .seal(this.closedError) - .then(() => this.storage.close()) + .then(() => this.storage.close(context)) .finally(() => { this.state = "closed"; this.onClose?.(); @@ -391,75 +431,121 @@ export class StorageBackedSession { - const leaf = await this.getValue(laneLeaf(lane)); + private async getLeafIdForLane(lane: string, context: Context): Promise { + const leaf = await this.getValue(laneLeaf(lane), context); if (leaf === undefined) throw new Error(`Unknown lane: ${lane}`); return leaf.value; } - private async findEntriesOnBranchForLane(lane: string, query: BranchScan = {}): Promise { + private async findEntriesOnBranchForLane( + lane: string, + query: BranchScan | undefined, + context: Context, + ): Promise { + query ??= {}; this.assertOpen(); - const start = query.start ?? (await this.getLeafIdForLane(lane)); + const start = query.start ?? (await this.getLeafIdForLane(lane, context)); if (start === null) return []; - return this.storage.scanBranch({ ...query, start, order: query.order ?? "newestFirst" }); + return this.storage.scanBranch({ ...query, start, order: query.order ?? "newestFirst" }, context); } - private async findEntryOnBranchForLane(lane: string, query: BranchScan = {}): Promise { - const entries = await this.findEntriesOnBranchForLane(lane, { - ...query, - limit: query.limit === undefined ? 1 : Math.min(query.limit, 1), - }); + private async findEntryOnBranchForLane( + lane: string, + query: BranchScan | undefined, + context: Context, + ): Promise { + query ??= {}; + const entries = await this.findEntriesOnBranchForLane( + lane, + { + ...query, + limit: query.limit === undefined ? 1 : Math.min(query.limit, 1), + }, + context, + ); return entries[0]; } - private setValueForLane(lane: string, address: Value, next: NoInfer): Promise { - return this.mutate(lane, async (mutator) => { - await mutator.commit([setValueWrite(address, next)]); - }); + private setValueForLane(lane: string, address: Value, next: NoInfer, context: Context): Promise { + return this.mutate( + lane, + async (mutator) => { + await mutator.commit([setValueWrite(address, next)], context); + }, + context, + ); } - private deleteValueForLane(lane: string, address: Value): Promise { - return this.mutate(lane, async (mutator) => { - await mutator.commit([deleteValueWrite(address)]); - }); + private deleteValueForLane(lane: string, address: Value, context: Context): Promise { + return this.mutate( + lane, + async (mutator) => { + await mutator.commit([deleteValueWrite(address)], context); + }, + context, + ); } - private appendListForLane(lane: string, address: ValueList, element: NoInfer): Promise { - return this.mutate(lane, async (mutator) => { - await mutator.commit([appendListWrite(address, element)]); - }); + private appendListForLane( + lane: string, + address: ValueList, + element: NoInfer, + context: Context, + ): Promise { + return this.mutate( + lane, + async (mutator) => { + await mutator.commit([appendListWrite(address, element)], context); + }, + context, + ); } - private deleteListForLane(lane: string, address: ValueList): Promise { - return this.mutate(lane, async (mutator) => { - await mutator.commit([deleteListWrite(address)]); - }); + private deleteListForLane(lane: string, address: ValueList, context: Context): Promise { + return this.mutate( + lane, + async (mutator) => { + await mutator.commit([deleteListWrite(address)], context); + }, + context, + ); } - private setNameForLane(lane: string, name: string | undefined): Promise { + private setNameForLane(lane: string, name: string | undefined, context: Context): Promise { return name === undefined - ? this.deleteValueForLane(lane, sessionName) - : this.setValueForLane(lane, sessionName, name); + ? this.deleteValueForLane(lane, sessionName, context) + : this.setValueForLane(lane, sessionName, name, context); } - private setLabelForLane(lane: string, targetId: string, label: string | undefined): Promise { + private setLabelForLane(lane: string, targetId: string, label: string | undefined, context: Context): Promise { const address = entryLabel(targetId); - return label === undefined ? this.deleteValueForLane(lane, address) : this.setValueForLane(lane, address, label); + return label === undefined + ? this.deleteValueForLane(lane, address, context) + : this.setValueForLane(lane, address, label, context); } - private appendMessageForLane(lane: string, message: AgentMessage): Promise { - return this.captureAppend(lane, { type: "message", payload: message }); + private appendMessageForLane(lane: string, message: AgentMessage, context: Context): Promise { + return this.captureAppend(lane, { type: "message", payload: message }, context); } - private appendCustomEntryForLane(lane: string, customType: string, data?: JsonValue): Promise { - return this.captureAppend(lane, { - type: "custom", - customType, - ...(data === undefined ? {} : { payload: data }), - }); + private appendCustomEntryForLane( + lane: string, + customType: string, + data: JsonValue | undefined, + context: Context, + ): Promise { + return this.captureAppend( + lane, + { + type: "custom", + customType, + ...(data === undefined ? {} : { payload: data }), + }, + context, + ); } - private async captureAppend(lane: string, pending: PendingEntry): Promise { + private async captureAppend(lane: string, pending: PendingEntry, context: Context): Promise { this.assertOpen(); if ( pending.type === "message" && @@ -468,45 +554,53 @@ export class StorageBackedSession { - await this.mutate(lane, (mutator) => this.appendCapturedIfReady(mutator, id, pending)); + private async appendCaptured(lane: string, id: string, pending: PendingEntry, context: Context): Promise { + await this.mutate(lane, (mutator) => this.appendCapturedIfReady(mutator, id, pending, context), context); return id; } - private async appendCapturedIfReady(mutator: SessionMutator, id: string, pending: PendingEntry): Promise { + private async appendCapturedIfReady( + mutator: SessionMutator, + id: string, + pending: PendingEntry, + context: Context, + ): Promise { const { lane } = mutator; const [leaf, storedLaneState] = await Promise.all([ - mutator.getValue(laneLeaf(lane)), - mutator.getValue(laneState(lane)), + mutator.getValue(laneLeaf(lane), context), + mutator.getValue(laneState(lane), context), ]); if (leaf === undefined) throw new SessionInvariantError(`Unknown lane: ${lane}`); if (storedLaneState === undefined) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.state`); const operationId = storedLaneState.value.currentOperationId; if (operationId === null) { - await mutator.commit([ - insertEntry( - pending.type === "message" - ? { id, parentId: leaf.value, type: "message", message: pending.payload } - : { - id, - parentId: leaf.value, - type: "custom", - customType: pending.customType, - ...(pending.payload === undefined ? {} : { data: pending.payload }), - }, - ), - setValueWrite(laneLeaf(lane), id), - ]); + await mutator.commit( + [ + insertEntry( + pending.type === "message" + ? { id, parentId: leaf.value, type: "message", message: pending.payload } + : { + id, + parentId: leaf.value, + type: "custom", + customType: pending.customType, + ...(pending.payload === undefined ? {} : { data: pending.payload }), + }, + ), + setValueWrite(laneLeaf(lane), id), + ], + context, + ); return; } const [operation, storedOperationState] = await Promise.all([ - mutator.getValue(operationMeta(operationId)), - mutator.getValue(operationState(operationId)), + mutator.getValue(operationMeta(operationId), context), + mutator.getValue(operationState(operationId), context), ]); if (operation === undefined) { throw new SessionInvariantError(`Active operation ${operationId} is missing op.meta`); @@ -521,16 +615,19 @@ export class StorageBackedSession { const metadata: SessionMetadata[] = []; for (let index = 0; index < dataset.sessionCount; index++) { - const session = await repo.create({ id: sessionRepoBenchmarkSessionId(index) }); + const session = await repo.create({ id: sessionRepoBenchmarkSessionId(index) }, BACKGROUND_CONTEXT); metadata.push(session.metadata); - await session.close(); + await session.close(BACKGROUND_CONTEXT); } return metadata; } @@ -74,7 +75,7 @@ export const SESSION_REPO_CATALOG_READ_BENCHMARK_SCENARIOS: readonly SessionRepo return dataset.sessionCount; }, async run(repo) { - return (await repo.list()).length; + return (await repo.list(undefined, BACKGROUND_CONTEXT)).length; }, }, ]; @@ -87,7 +88,7 @@ export const SESSION_REPO_CATALOG_WRITE_BENCHMARK_SCENARIOS: readonly SessionRep prepare(repo) { return Promise.resolve({ async run() { - const session = await repo.create({ id: BENCHMARK_SESSION_ID }); + const session = await repo.create({ id: BENCHMARK_SESSION_ID }, BACKGROUND_CONTEXT); return session.metadata.id === BENCHMARK_SESSION_ID ? 1 : 0; }, }); @@ -97,12 +98,12 @@ export const SESSION_REPO_CATALOG_WRITE_BENCHMARK_SCENARIOS: readonly SessionRep name: "open closed empty session", expectedResult: 1, async prepare(repo) { - const session = await repo.create({ id: BENCHMARK_SESSION_ID }); + const session = await repo.create({ id: BENCHMARK_SESSION_ID }, BACKGROUND_CONTEXT); const { metadata } = session; - await session.close(); + await session.close(BACKGROUND_CONTEXT); return { async run() { - const reopened = await repo.open(metadata); + const reopened = await repo.open(metadata, BACKGROUND_CONTEXT); return reopened.metadata.id === metadata.id ? 1 : 0; }, }; @@ -112,12 +113,12 @@ export const SESSION_REPO_CATALOG_WRITE_BENCHMARK_SCENARIOS: readonly SessionRep name: "delete closed empty session", expectedResult: 1, async prepare(repo) { - const session = await repo.create({ id: BENCHMARK_SESSION_ID }); + const session = await repo.create({ id: BENCHMARK_SESSION_ID }, BACKGROUND_CONTEXT); const { metadata } = session; - await session.close(); + await session.close(BACKGROUND_CONTEXT); return { async run() { - await repo.delete(metadata); + await repo.delete(metadata, BACKGROUND_CONTEXT); return 1; }, }; @@ -136,11 +137,15 @@ export async function seedSessionRepoForkBenchmark( repo: SessionRepo, dataset: StorageBenchmarkDataset, ): Promise { - const session = await repo.create({ id: BENCHMARK_SESSION_ID }); + const session = await repo.create({ id: BENCHMARK_SESSION_ID }, BACKGROUND_CONTEXT); for (const transaction of generateStorageBenchmarkSeedTransactions(dataset)) { - await session.mutate("main", (mutator) => mutator.commit(transaction)); + await session.mutate("main", (mutator) => mutator.commit(transaction, BACKGROUND_CONTEXT), BACKGROUND_CONTEXT); } - await session.mutate("main", (mutator) => mutator.commit([setValue(laneLeaf("main"), dataset.leafId)])); + await session.mutate( + "main", + (mutator) => mutator.commit([setValue(laneLeaf("main"), dataset.leafId)], BACKGROUND_CONTEXT), + BACKGROUND_CONTEXT, + ); return session.metadata; } @@ -152,7 +157,7 @@ export const SESSION_REPO_FORK_WRITE_BENCHMARK_SCENARIOS: readonly SessionRepoFo return dataset.entryCount; }, async run(repo, source, dataset) { - const fork = await repo.fork(source, { id: FORK_DESTINATION_SESSION_ID }); + const fork = await repo.fork(source, { id: FORK_DESTINATION_SESSION_ID }, BACKGROUND_CONTEXT); return fork.metadata.id === FORK_DESTINATION_SESSION_ID && fork.metadata.parentSessionId === source.id ? dataset.entryCount : 0; diff --git a/packages/agent/src/harness/session/testing/benchmark/storage.ts b/packages/agent/src/harness/session/testing/benchmark/storage.ts index e21888c0b..5c66decbd 100644 --- a/packages/agent/src/harness/session/testing/benchmark/storage.ts +++ b/packages/agent/src/harness/session/testing/benchmark/storage.ts @@ -1,3 +1,4 @@ +import { BACKGROUND_CONTEXT } from "../../../context.ts"; import { insertEntry, insertUsage } from "../../commit.ts"; import type { MessageEntry, NewEntry, Storage, Write } from "../../types.ts"; import { laneLeaf, setValue } from "../../values.ts"; @@ -41,7 +42,8 @@ export function* generateStorageBenchmarkSeedTransactions(dataset: StorageBenchm /** Seeds one deterministic synthetic linear branch. */ export async function seedStorageBenchmark(storage: Storage, dataset: StorageBenchmarkDataset): Promise { - for (const transaction of generateStorageBenchmarkSeedTransactions(dataset)) await storage.commit(transaction); + for (const transaction of generateStorageBenchmarkSeedTransactions(dataset)) + await storage.commit(transaction, BACKGROUND_CONTEXT); } /** A steady-state read operation run against a pre-seeded fixture. */ @@ -88,7 +90,7 @@ export const STORAGE_READ_BENCHMARK_SCENARIOS: readonly StorageReadBenchmarkScen return dataset.lookupIds.length; }, async run(storage, dataset) { - return (await storage.getEntries([...dataset.lookupIds])).size; + return (await storage.getEntries([...dataset.lookupIds], BACKGROUND_CONTEXT)).size; }, }, { @@ -97,7 +99,7 @@ export const STORAGE_READ_BENCHMARK_SCENARIOS: readonly StorageReadBenchmarkScen return Math.min(50, dataset.entryCount); }, async run(storage) { - return (await storage.scanEntries({ order: "desc", limit: 50 })).length; + return (await storage.scanEntries({ order: "desc", limit: 50 }, BACKGROUND_CONTEXT)).length; }, }, { @@ -106,7 +108,8 @@ export const STORAGE_READ_BENCHMARK_SCENARIOS: readonly StorageReadBenchmarkScen return dataset.entryCount; }, async run(storage, dataset) { - return (await storage.scanBranchStructure({ start: dataset.leafId, order: "newestFirst" })).length; + return (await storage.scanBranchStructure({ start: dataset.leafId, order: "newestFirst" }, BACKGROUND_CONTEXT)) + .length; }, }, ]; @@ -117,14 +120,14 @@ export const STORAGE_WRITE_BENCHMARK_SCENARIOS: readonly StorageWriteBenchmarkSc name: "commit one message entry", writeCount: 1, async run(storage) { - return (await storage.commit(singleEntryTransaction)).seqs.length; + return (await storage.commit(singleEntryTransaction, BACKGROUND_CONTEXT)).seqs.length; }, }, { name: "commit 100 message entries", writeCount: 100, async run(storage) { - return (await storage.commit(hundredEntryTransaction)).seqs.length; + return (await storage.commit(hundredEntryTransaction, BACKGROUND_CONTEXT)).seqs.length; }, }, { @@ -134,7 +137,7 @@ export const STORAGE_WRITE_BENCHMARK_SCENARIOS: readonly StorageWriteBenchmarkSc return seedStorageBenchmark(storage, WRITE_BASELINE_DATASET); }, async run(storage) { - return (await storage.commit(mixedAppendTransaction)).seqs.length; + return (await storage.commit(mixedAppendTransaction, BACKGROUND_CONTEXT)).seqs.length; }, }, ]; diff --git a/packages/agent/src/harness/session/testing/conformance/session-repo.ts b/packages/agent/src/harness/session/testing/conformance/session-repo.ts index b274cd340..f39d87162 100644 --- a/packages/agent/src/harness/session/testing/conformance/session-repo.ts +++ b/packages/agent/src/harness/session/testing/conformance/session-repo.ts @@ -1,5 +1,6 @@ import { deepStrictEqual, rejects, strictEqual } from "node:assert/strict"; import type { AssistantMessage, StopReason } from "@earendil-works/pi-ai"; +import { BACKGROUND_CONTEXT } from "../../../context.ts"; import { insertEntry, insertUsage } from "../../commit.ts"; import type { JsonValue, LaneConfiguration, LaneState, SessionMetadata, SessionRepo, UsageRow } from "../../types.ts"; import { @@ -122,19 +123,19 @@ export function createSessionRepoLifecycleConformance { - const session = await repo.create({ id: "session" }); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); strictEqual(session.metadata.id, "session"); strictEqual(Number.isSafeInteger(session.metadata.createdAt), true); strictEqual(session.metadata.storageVersion, 1); - strictEqual(await session.getLeafId(), null); - deepStrictEqual((await session.getValue(laneState("main")))?.value, { + strictEqual(await session.getLeafId(BACKGROUND_CONTEXT), null); + deepStrictEqual((await session.getValue(laneState("main"), BACKGROUND_CONTEXT))?.value, { currentOperationId: null, pendingNextRun: [], }); - strictEqual(await session.getValue(laneConfig("main")), undefined); - await rejects(repo.create({ id: "session" })); - await session.close(); + strictEqual(await session.getValue(laneConfig("main"), BACKGROUND_CONTEXT), undefined); + await rejects(repo.create({ id: "session" }, BACKGROUND_CONTEXT)); + await session.close(BACKGROUND_CONTEXT); }, ), createCase( @@ -142,12 +143,12 @@ export function createSessionRepoLifecycleConformance { - const first = await repo.create({ id: "first" }); - await first.setName("preserved"); - const second = await repo.create({ id: "second", parentSessionId: "parent" }); - await second.close(); + const first = await repo.create({ id: "first" }, BACKGROUND_CONTEXT); + await first.setName("preserved", BACKGROUND_CONTEXT); + const second = await repo.create({ id: "second", parentSessionId: "parent" }, BACKGROUND_CONTEXT); + await second.close(BACKGROUND_CONTEXT); - const listed = await repo.list(); + const listed = await repo.list(undefined, BACKGROUND_CONTEXT); deepStrictEqual( listed .map(({ id, parentSessionId }) => ({ id, parentSessionId })) @@ -157,26 +158,26 @@ export function createSessionRepoLifecycleConformance { - const removed = await repo.create({ id: "removed" }); - const retained = await repo.create({ id: "retained" }); - await Promise.all([removed.close(), retained.close()]); + const removed = await repo.create({ id: "removed" }, BACKGROUND_CONTEXT); + const retained = await repo.create({ id: "retained" }, BACKGROUND_CONTEXT); + await Promise.all([removed.close(BACKGROUND_CONTEXT), retained.close(BACKGROUND_CONTEXT)]); - await repo.delete(removed.metadata); + await repo.delete(removed.metadata, BACKGROUND_CONTEXT); deepStrictEqual( - (await repo.list()).map(({ id }) => id), + (await repo.list(undefined, BACKGROUND_CONTEXT)).map(({ id }) => id), ["retained"], ); - await rejects(repo.open(removed.metadata)); - await rejects(repo.delete(removed.metadata)); + await rejects(repo.open(removed.metadata, BACKGROUND_CONTEXT)); + await rejects(repo.delete(removed.metadata, BACKGROUND_CONTEXT)); }), ]; } @@ -189,13 +190,13 @@ export function createSessionRepoOwnershipConformance { - const session = await repo.create({ id: "session" }); - await rejects(repo.open(session.metadata)); - await session.close(); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); + await rejects(repo.open(session.metadata, BACKGROUND_CONTEXT)); + await session.close(BACKGROUND_CONTEXT); - const reopened = await repo.open(session.metadata); - await rejects(repo.open(session.metadata)); - await reopened.close(); + const reopened = await repo.open(session.metadata, BACKGROUND_CONTEXT); + await rejects(repo.open(session.metadata, BACKGROUND_CONTEXT)); + await reopened.close(BACKGROUND_CONTEXT); }), ]; } @@ -212,17 +213,17 @@ export function createSessionRepoMessageConformance { - const session = await repo.create({ id: "session" }); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); - await rejects(session.appendMessage(assistantMessage("pending"))); + await rejects(session.appendMessage(assistantMessage("pending"), BACKGROUND_CONTEXT)); - strictEqual(await session.getLeafId(), null); - deepStrictEqual(await session.findEntries(), []); - await session.close(); + strictEqual(await session.getLeafId(BACKGROUND_CONTEXT), null); + deepStrictEqual(await session.findEntries(undefined, BACKGROUND_CONTEXT), []); + await session.close(BACKGROUND_CONTEXT); }, ), createCase(factory, "messages", "preserves every settled assistant stop reason", async ({ repo }) => { - const session = await repo.create({ id: "session" }); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); const messagesByStopReason = { stop: assistantMessage("stop"), length: assistantMessage("length"), @@ -234,9 +235,9 @@ export function createSessionRepoMessageConformance entry.id), ids, @@ -245,8 +246,8 @@ export function createSessionRepoMessageConformance { - const source = await repo.create({ id: "source" }); - const fork = await repo.fork(source.metadata, { id: "fork" }); + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); + const fork = await repo.fork(source.metadata, { id: "fork" }, BACKGROUND_CONTEXT); strictEqual(fork.metadata.id, "fork"); strictEqual(fork.metadata.parentSessionId, "source"); - strictEqual(await fork.getLeafId(), null); - strictEqual(await fork.getValue(laneConfig("main")), undefined); - deepStrictEqual((await fork.getValue(laneState("main")))?.value, { + strictEqual(await fork.getLeafId(BACKGROUND_CONTEXT), null); + strictEqual(await fork.getValue(laneConfig("main"), BACKGROUND_CONTEXT), undefined); + deepStrictEqual((await fork.getValue(laneState("main"), BACKGROUND_CONTEXT))?.value, { currentOperationId: null, pendingNextRun: [], }); - deepStrictEqual(await fork.findEntries(), []); - deepStrictEqual(await fork.getStats(), { + deepStrictEqual(await fork.findEntries(undefined, BACKGROUND_CONTEXT), []); + deepStrictEqual(await fork.getStats(BACKGROUND_CONTEXT), { messageCount: 0, usage: { input: 0, @@ -282,96 +283,112 @@ export function createSessionRepoForkBehaviorConformance { - const source = await repo.create({ id: "source" }); - await source.mutate("main", (mutator) => - mutator.commit([ - insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }), - insertEntry({ - id: CHILD_ID, - parentId: ROOT_ID, - type: "message", - message: { role: "user", content: "child", timestamp: 1 }, - }), - insertEntry({ - id: SIBLING_ID, - parentId: ROOT_ID, - type: "custom", - customType: "sibling", - }), - setValue(laneLeaf("main"), CHILD_ID), - setValue(laneConfig("main"), configuration), - setValue(laneState("main"), { - currentOperationId: OPERATION_ID, - pendingNextRun: [PENDING_ID], - }), - setValue(laneLastResult("main"), { - operationId: "previous", - kind: "navigation", - leafId: CHILD_ID, - oldLeafId: ROOT_ID, - outcome: "completed", - }), - setValue(sessionName, "source name"), - setValue(applicationValue, { copied: false }), - appendList(applicationList, { copied: false }), - setValue(entryLabel(ROOT_ID), "root label"), - setValue(entryLabel(SIBLING_ID), "sibling label"), - setValue(pendingEntry(PENDING_ID), { type: "custom", customType: "pending" }), - setValue(operationMeta(OPERATION_ID), { - operationId: OPERATION_ID, - lane: "main", - sourceLeafId: CHILD_ID, - startedAt: 1, - intent: { kind: "compaction" }, - }), - setValue(operationState(OPERATION_ID), { - kind: "compaction", - control: { status: "running" }, - structural: { taskId: OPERATION_ID, status: "deciding" }, - }), - setValue(operationToolArgs(OPERATION_ID, ROOT_ID, 0), { argument: true }), - setValue(operationPreparation(OPERATION_ID, OPERATION_ID), { - kind: "compaction", - messagesToSummarize: [], - turnPrefixMessages: [], - retainedTail: [], - isSplitTurn: false, - tokensBefore: 0, - fileOps: { read: [], written: [], edited: [] }, - settings: { enabled: true, reserveTokens: 1, keepRecentTokens: 1 }, - }), - insertUsage(usageRow()), - ]), + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); + await source.mutate( + "main", + (mutator) => + mutator.commit( + [ + insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }), + insertEntry({ + id: CHILD_ID, + parentId: ROOT_ID, + type: "message", + message: { role: "user", content: "child", timestamp: 1 }, + }), + insertEntry({ + id: SIBLING_ID, + parentId: ROOT_ID, + type: "custom", + customType: "sibling", + }), + setValue(laneLeaf("main"), CHILD_ID), + setValue(laneConfig("main"), configuration), + setValue(laneState("main"), { + currentOperationId: OPERATION_ID, + pendingNextRun: [PENDING_ID], + }), + setValue(laneLastResult("main"), { + operationId: "previous", + kind: "navigation", + leafId: CHILD_ID, + oldLeafId: ROOT_ID, + outcome: "completed", + }), + setValue(sessionName, "source name"), + setValue(applicationValue, { copied: false }), + appendList(applicationList, { copied: false }), + setValue(entryLabel(ROOT_ID), "root label"), + setValue(entryLabel(SIBLING_ID), "sibling label"), + setValue(pendingEntry(PENDING_ID), { type: "custom", customType: "pending" }), + setValue(operationMeta(OPERATION_ID), { + operationId: OPERATION_ID, + lane: "main", + sourceLeafId: CHILD_ID, + startedAt: 1, + intent: { kind: "compaction" }, + }), + setValue(operationState(OPERATION_ID), { + kind: "compaction", + control: { status: "running" }, + structural: { taskId: OPERATION_ID, status: "deciding" }, + }), + setValue(operationToolArgs(OPERATION_ID, ROOT_ID, 0), { argument: true }), + setValue(operationPreparation(OPERATION_ID, OPERATION_ID), { + kind: "compaction", + messagesToSummarize: [], + turnPrefixMessages: [], + retainedTail: [], + isSplitTurn: false, + tokensBefore: 0, + fileOps: { read: [], written: [], edited: [] }, + settings: { enabled: true, reserveTokens: 1, keepRecentTokens: 1 }, + }), + insertUsage(usageRow()), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - const fork = await repo.fork(source.metadata, { id: "fork", entryId: CHILD_ID, position: "at" }); + const fork = await repo.fork( + source.metadata, + { id: "fork", entryId: CHILD_ID, position: "at" }, + BACKGROUND_CONTEXT, + ); deepStrictEqual( - (await fork.findEntries({ order: "asc" })).map(({ id }) => id), + (await fork.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id), [ROOT_ID, CHILD_ID], ); - strictEqual(await fork.getLeafId(), CHILD_ID); - deepStrictEqual((await fork.getValue(laneConfig("main")))?.value, configuration); - deepStrictEqual((await fork.getValue(laneState("main")))?.value, idleLaneState); - strictEqual(await fork.getName(), "source name"); - strictEqual(await fork.getValue(applicationValue), undefined); - deepStrictEqual(await fork.readList(applicationList), []); - strictEqual(await fork.getLabel(ROOT_ID), "root label"); - strictEqual(await fork.getLabel(SIBLING_ID), undefined); - strictEqual(await fork.getValue(laneLastResult("main")), undefined); - strictEqual(await fork.getValue(pendingEntry(PENDING_ID)), undefined); - strictEqual(await fork.getValue(operationMeta(OPERATION_ID)), undefined); - strictEqual(await fork.getValue(operationState(OPERATION_ID)), undefined); - strictEqual(await fork.getValue(operationToolArgs(OPERATION_ID, ROOT_ID, 0)), undefined); - strictEqual(await fork.getValue(operationPreparation(OPERATION_ID, OPERATION_ID)), undefined); - const stats = await fork.getStats(); + strictEqual(await fork.getLeafId(BACKGROUND_CONTEXT), CHILD_ID); + deepStrictEqual((await fork.getValue(laneConfig("main"), BACKGROUND_CONTEXT))?.value, configuration); + deepStrictEqual((await fork.getValue(laneState("main"), BACKGROUND_CONTEXT))?.value, idleLaneState); + strictEqual(await fork.getName(BACKGROUND_CONTEXT), "source name"); + strictEqual(await fork.getValue(applicationValue, BACKGROUND_CONTEXT), undefined); + deepStrictEqual(await fork.readList(applicationList, undefined, BACKGROUND_CONTEXT), []); + strictEqual(await fork.getLabel(ROOT_ID, BACKGROUND_CONTEXT), "root label"); + strictEqual(await fork.getLabel(SIBLING_ID, BACKGROUND_CONTEXT), undefined); + strictEqual(await fork.getValue(laneLastResult("main"), BACKGROUND_CONTEXT), undefined); + strictEqual(await fork.getValue(pendingEntry(PENDING_ID), BACKGROUND_CONTEXT), undefined); + strictEqual(await fork.getValue(operationMeta(OPERATION_ID), BACKGROUND_CONTEXT), undefined); + strictEqual(await fork.getValue(operationState(OPERATION_ID), BACKGROUND_CONTEXT), undefined); + strictEqual( + await fork.getValue(operationToolArgs(OPERATION_ID, ROOT_ID, 0), BACKGROUND_CONTEXT), + undefined, + ); + strictEqual( + await fork.getValue(operationPreparation(OPERATION_ID, OPERATION_ID), BACKGROUND_CONTEXT), + undefined, + ); + const stats = await fork.getStats(BACKGROUND_CONTEXT); strictEqual(stats.messageCount, 1); deepStrictEqual(stats.usage, { input: 0, @@ -381,7 +398,7 @@ export function createSessionRepoForkBehaviorConformance { - const source = await repo.create({ id: "source" }); - await source.mutate("main", (mutator) => - mutator.commit([ - insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }), - insertEntry({ - id: CHILD_ID, - parentId: ROOT_ID, - type: "custom", - customType: "child", - }), - setValue(laneLeaf("main"), CHILD_ID), - ]), + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); + await source.mutate( + "main", + (mutator) => + mutator.commit( + [ + insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }), + insertEntry({ + id: CHILD_ID, + parentId: ROOT_ID, + type: "custom", + customType: "child", + }), + setValue(laneLeaf("main"), CHILD_ID), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - const before = await repo.fork(source.metadata, { id: "before", entryId: CHILD_ID, position: "before" }); - strictEqual(await before.getLeafId(), ROOT_ID); + const before = await repo.fork( + source.metadata, + { id: "before", entryId: CHILD_ID, position: "before" }, + BACKGROUND_CONTEXT, + ); + strictEqual(await before.getLeafId(BACKGROUND_CONTEXT), ROOT_ID); deepStrictEqual( - (await before.findEntries({ order: "asc" })).map(({ id }) => id), + (await before.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id), [ROOT_ID], ); - await rejects(repo.fork(source.metadata, { id: "failed", entryId: SIBLING_ID })); - deepStrictEqual((await repo.list()).map(({ id }) => id).sort(), ["before", "source"]); - await Promise.all([source.close(), before.close()]); + await rejects(repo.fork(source.metadata, { id: "failed", entryId: SIBLING_ID }, BACKGROUND_CONTEXT)); + deepStrictEqual((await repo.list(undefined, BACKGROUND_CONTEXT)).map(({ id }) => id).sort(), [ + "before", + "source", + ]); + await Promise.all([source.close(BACKGROUND_CONTEXT), before.close(BACKGROUND_CONTEXT)]); }, ), createCase(factory, "forks", "forks a closed source session", async ({ repo }) => { - const source = await repo.create({ id: "source" }); - await source.mutate("main", (mutator) => - mutator.commit([ - insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }), - setValue(laneLeaf("main"), ROOT_ID), - setValue(applicationValue, "excluded"), - appendList(applicationList, "excluded"), - ]), + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); + await source.mutate( + "main", + (mutator) => + mutator.commit( + [ + insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }), + setValue(laneLeaf("main"), ROOT_ID), + setValue(applicationValue, "excluded"), + appendList(applicationList, "excluded"), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - await source.close(); + await source.close(BACKGROUND_CONTEXT); - const fork = await repo.fork(source.metadata, { id: "fork" }); - strictEqual(await fork.getLeafId(), ROOT_ID); - strictEqual(await fork.getValue(applicationValue), undefined); - deepStrictEqual(await fork.readList(applicationList), []); - await fork.close(); + const fork = await repo.fork(source.metadata, { id: "fork" }, BACKGROUND_CONTEXT); + strictEqual(await fork.getLeafId(BACKGROUND_CONTEXT), ROOT_ID); + strictEqual(await fork.getValue(applicationValue, BACKGROUND_CONTEXT), undefined); + deepStrictEqual(await fork.readList(applicationList, undefined, BACKGROUND_CONTEXT), []); + await fork.close(BACKGROUND_CONTEXT); }), createCase(factory, "forks", "forks the whole configured tree with fresh lane state", async ({ repo }) => { - const source = await repo.create({ id: "source" }); - await source.mutate("main", (mutator) => - mutator.commit([ - insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }), - insertEntry({ id: CHILD_ID, parentId: ROOT_ID, type: "custom", customType: "child" }), - insertEntry({ - id: SIBLING_ID, - parentId: ROOT_ID, - type: "custom", - customType: "sibling", - }), - setValue(laneLeaf("main"), CHILD_ID), - setValue(laneConfig("main"), configuration), - setValue(laneLeaf("review"), SIBLING_ID), - setValue(laneConfig("review"), configuration), - setValue(laneState("review"), idleLaneState), - ]), + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); + await source.mutate( + "main", + (mutator) => + mutator.commit( + [ + insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }), + insertEntry({ id: CHILD_ID, parentId: ROOT_ID, type: "custom", customType: "child" }), + insertEntry({ + id: SIBLING_ID, + parentId: ROOT_ID, + type: "custom", + customType: "sibling", + }), + setValue(laneLeaf("main"), CHILD_ID), + setValue(laneConfig("main"), configuration), + setValue(laneLeaf("review"), SIBLING_ID), + setValue(laneConfig("review"), configuration), + setValue(laneState("review"), idleLaneState), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - const fork = await repo.fork(source.metadata, { id: "fork", scope: "tree" }); + const fork = await repo.fork(source.metadata, { id: "fork", scope: "tree" }, BACKGROUND_CONTEXT); deepStrictEqual( - (await fork.findEntries({ order: "asc" })).map(({ id }) => id), + (await fork.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id), [ROOT_ID, CHILD_ID, SIBLING_ID], ); - strictEqual(await fork.getLeafId(), CHILD_ID); - strictEqual(await fork.view("review").getLeafId(), SIBLING_ID); - deepStrictEqual((await fork.getValue(laneConfig("review")))?.value, configuration); - deepStrictEqual((await fork.getValue(laneState("review")))?.value, idleLaneState); - await Promise.all([source.close(), fork.close()]); + strictEqual(await fork.getLeafId(BACKGROUND_CONTEXT), CHILD_ID); + strictEqual(await fork.view("review").getLeafId(BACKGROUND_CONTEXT), SIBLING_ID); + deepStrictEqual((await fork.getValue(laneConfig("review"), BACKGROUND_CONTEXT))?.value, configuration); + deepStrictEqual((await fork.getValue(laneState("review"), BACKGROUND_CONTEXT))?.value, idleLaneState); + await Promise.all([source.close(BACKGROUND_CONTEXT), fork.close(BACKGROUND_CONTEXT)]); }), ]; } @@ -480,15 +522,15 @@ export function createSessionRepoForkDestinationReservationConformance { - const source = await repo.create({ id: "source" }); + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); const results = await Promise.allSettled([ - repo.create({ id: "destination" }), - repo.fork(source.metadata, { id: "destination" }), + repo.create({ id: "destination" }, BACKGROUND_CONTEXT), + repo.fork(source.metadata, { id: "destination" }, BACKGROUND_CONTEXT), ]); strictEqual(results[0].status, "fulfilled"); strictEqual(results[1].status, "rejected"); - if (results[0].status === "fulfilled") await results[0].value.close(); - await source.close(); + if (results[0].status === "fulfilled") await results[0].value.close(BACKGROUND_CONTEXT); + await source.close(BACKGROUND_CONTEXT); }, ), createCase( @@ -496,15 +538,15 @@ export function createSessionRepoForkDestinationReservationConformance { - const source = await repo.create({ id: "source" }); + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); const results = await Promise.allSettled([ - repo.fork(source.metadata, { id: "destination" }), - repo.create({ id: "destination" }), + repo.fork(source.metadata, { id: "destination" }, BACKGROUND_CONTEXT), + repo.create({ id: "destination" }, BACKGROUND_CONTEXT), ]); strictEqual(results[0].status, "fulfilled"); strictEqual(results[1].status, "rejected"); - if (results[0].status === "fulfilled") await results[0].value.close(); - await source.close(); + if (results[0].status === "fulfilled") await results[0].value.close(BACKGROUND_CONTEXT); + await source.close(BACKGROUND_CONTEXT); }, ), ]; @@ -522,34 +564,46 @@ export function createSessionRepoForkSourceSnapshotConformance { - const source = await repo.create({ id: "source" }); - const firstCommit = source.mutate("main", (mutator) => - mutator.commit([ - insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "first" }), - setValue(laneLeaf("main"), ROOT_ID), - ]), + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); + const firstCommit = source.mutate( + "main", + (mutator) => + mutator.commit( + [ + insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "first" }), + setValue(laneLeaf("main"), ROOT_ID), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - const fork = repo.fork(source.metadata, { id: "fork" }); - const secondCommit = source.mutate("main", (mutator) => - mutator.commit([ - insertEntry({ - id: CHILD_ID, - parentId: ROOT_ID, - type: "custom", - customType: "second", - }), - setValue(laneLeaf("main"), CHILD_ID), - ]), + const fork = repo.fork(source.metadata, { id: "fork" }, BACKGROUND_CONTEXT); + const secondCommit = source.mutate( + "main", + (mutator) => + mutator.commit( + [ + insertEntry({ + id: CHILD_ID, + parentId: ROOT_ID, + type: "custom", + customType: "second", + }), + setValue(laneLeaf("main"), CHILD_ID), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); const [, forked] = await Promise.all([firstCommit, fork]); await secondCommit; - strictEqual(await forked.getLeafId(), ROOT_ID); + strictEqual(await forked.getLeafId(BACKGROUND_CONTEXT), ROOT_ID); deepStrictEqual( - (await forked.findEntries({ order: "asc" })).map(({ id }) => id), + (await forked.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id), [ROOT_ID], ); - await Promise.all([source.close(), forked.close()]); + await Promise.all([source.close(BACKGROUND_CONTEXT), forked.close(BACKGROUND_CONTEXT)]); }, ), ]; diff --git a/packages/agent/src/harness/session/testing/conformance/storage.ts b/packages/agent/src/harness/session/testing/conformance/storage.ts index b6a89abb1..f92915696 100644 --- a/packages/agent/src/harness/session/testing/conformance/storage.ts +++ b/packages/agent/src/harness/session/testing/conformance/storage.ts @@ -1,5 +1,6 @@ import { deepStrictEqual, ok, rejects, strictEqual } from "node:assert/strict"; import type { Usage } from "@earendil-works/pi-ai"; +import { BACKGROUND_CONTEXT } from "../../../context.ts"; import { insertEntry, insertUsage } from "../../commit.ts"; import type { CompactionEntry, CustomEntry, Entry, JsonValue, MessageEntry, NewEntry, Write } from "../../types.ts"; import { @@ -117,26 +118,29 @@ function assertStrictlyIncreasing(values: readonly number[]): void { export function createStorageConformance(factory: () => Promise): readonly ConformanceCase[] { return [ createCase(factory, "transactions", "commits mixed writes atomically in write order", async ({ storage }) => { - const result = await storage.commit([ - insertEntry(userEntry("entry")), - setValue(testName, "session"), - insertUsage({ id: "usage", usage: usage(2, 3), adjustment: false, entryId: "entry" }), - ]); + const result = await storage.commit( + [ + insertEntry(userEntry("entry")), + setValue(testName, "session"), + insertUsage({ id: "usage", usage: usage(2, 3), adjustment: false, entryId: "entry" }), + ], + BACKGROUND_CONTEXT, + ); strictEqual(result.seqs.length, 3); strictEqual(result.firstSeq, result.seqs[0]); assertStrictlyIncreasing(result.seqs); ok(Number.isSafeInteger(result.timestamp) && result.timestamp >= 0); deepStrictEqual( - await storage.getEntries(["entry"]), + await storage.getEntries(["entry"], BACKGROUND_CONTEXT), new Map([["entry", { ...userEntry("entry"), seq: result.seqs[0], timestamp: result.timestamp }]]), ); - deepStrictEqual(await storage.getValue(testName), { + deepStrictEqual(await storage.getValue(testName, BACKGROUND_CONTEXT), { address: testName, value: "session", seq: result.seqs[1], }); - deepStrictEqual(await storage.scanUsage({ order: "asc" }), [ + deepStrictEqual(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT), [ { id: "usage", seq: result.seqs[2], usage: usage(2, 3), adjustment: false, entryId: "entry" }, ]); }), @@ -146,27 +150,30 @@ export function createStorageConformance(factory: () => Promise) "transactions", "rolls back every store when a mixed transaction fails", async ({ storage }) => { - await storage.commit([ - insertEntry(userEntry("root")), - insertUsage({ id: "taken", usage: usage(1, 1), adjustment: false }), - ]); - const entriesBefore = await storage.scanEntries({ order: "asc" }); - const usageBefore = await storage.scanUsage({ order: "asc" }); - const statsBefore = await storage.getStats(); + await storage.commit( + [insertEntry(userEntry("root")), insertUsage({ id: "taken", usage: usage(1, 1), adjustment: false })], + BACKGROUND_CONTEXT, + ); + const entriesBefore = await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT); + const usageBefore = await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT); + const statsBefore = await storage.getStats(BACKGROUND_CONTEXT); await rejects( - storage.commit([ - setValue(testName, "transient"), - insertEntry(customEntry("transient-entry", "root")), - insertUsage({ id: "transient-usage", usage: usage(5, 8), adjustment: true }), - insertEntry(customEntry("taken", "root")), - ]), + storage.commit( + [ + setValue(testName, "transient"), + insertEntry(customEntry("transient-entry", "root")), + insertUsage({ id: "transient-usage", usage: usage(5, 8), adjustment: true }), + insertEntry(customEntry("taken", "root")), + ], + BACKGROUND_CONTEXT, + ), ); - deepStrictEqual(await storage.scanEntries({ order: "asc" }), entriesBefore); - deepStrictEqual(await storage.scanUsage({ order: "asc" }), usageBefore); - deepStrictEqual(await storage.getStats(), statsBefore); - strictEqual(await storage.getValue(testName), undefined); + deepStrictEqual(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT), entriesBefore); + deepStrictEqual(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT), usageBefore); + deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), statsBefore); + strictEqual(await storage.getValue(testName, BACKGROUND_CONTEXT), undefined); }, ), @@ -175,37 +182,51 @@ export function createStorageConformance(factory: () => Promise) "transactions", "preserves overwritten and deleted values when a transaction fails", async ({ storage }) => { - await storage.commit([ - setValue(testValue("overwritten"), "original"), - setValue(testValue("deleted"), { kept: true }), - insertEntry(userEntry("taken")), - ]); - const overwrittenBefore = await storage.getValue(testValue("overwritten")); - const deletedBefore = await storage.getValue(testValue("deleted")); + await storage.commit( + [ + setValue(testValue("overwritten"), "original"), + setValue(testValue("deleted"), { kept: true }), + insertEntry(userEntry("taken")), + ], + BACKGROUND_CONTEXT, + ); + const overwrittenBefore = await storage.getValue(testValue("overwritten"), BACKGROUND_CONTEXT); + const deletedBefore = await storage.getValue(testValue("deleted"), BACKGROUND_CONTEXT); await rejects( - storage.commit([ - setValue(testValue("overwritten"), "transient"), - deleteValue(testValue("deleted")), - insertEntry(customEntry("transient", "taken")), - insertEntry(customEntry("taken", null)), - ]), + storage.commit( + [ + setValue(testValue("overwritten"), "transient"), + deleteValue(testValue("deleted")), + insertEntry(customEntry("transient", "taken")), + insertEntry(customEntry("taken", null)), + ], + BACKGROUND_CONTEXT, + ), ); - deepStrictEqual(await storage.getValue(testValue("overwritten")), overwrittenBefore); - deepStrictEqual(await storage.getValue(testValue("deleted")), deletedBefore); - strictEqual((await storage.getEntries(["transient"])).has("transient"), false); + deepStrictEqual(await storage.getValue(testValue("overwritten"), BACKGROUND_CONTEXT), overwrittenBefore); + deepStrictEqual(await storage.getValue(testValue("deleted"), BACKGROUND_CONTEXT), deletedBefore); + strictEqual((await storage.getEntries(["transient"], BACKGROUND_CONTEXT)).has("transient"), false); }, ), createCase(factory, "transactions", "enforces one shared entry and usage id namespace", async ({ storage }) => { - await storage.commit([ - insertEntry(userEntry("existing-entry")), - insertUsage({ id: "existing-usage", usage: usage(1, 1), adjustment: false }), - ]); + await storage.commit( + [ + insertEntry(userEntry("existing-entry")), + insertUsage({ id: "existing-usage", usage: usage(1, 1), adjustment: false }), + ], + BACKGROUND_CONTEXT, + ); - await rejects(storage.commit([insertUsage({ id: "existing-entry", usage: usage(2, 2), adjustment: false })])); - await rejects(storage.commit([insertEntry(customEntry("existing-usage", null))])); + await rejects( + storage.commit( + [insertUsage({ id: "existing-entry", usage: usage(2, 2), adjustment: false })], + BACKGROUND_CONTEXT, + ), + ); + await rejects(storage.commit([insertEntry(customEntry("existing-usage", null))], BACKGROUND_CONTEXT)); for (const [id, writes] of [ [ @@ -223,12 +244,12 @@ export function createStorageConformance(factory: () => Promise) ], ], ] satisfies Array<[string, Write[]]>) { - await rejects(storage.commit(writes), `Expected duplicate id ${id} to reject`); + await rejects(storage.commit(writes, BACKGROUND_CONTEXT), `Expected duplicate id ${id} to reject`); } - deepStrictEqual(ids(await storage.scanEntries({ order: "asc" })), ["existing-entry"]); + deepStrictEqual(ids(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT)), ["existing-entry"]); deepStrictEqual( - (await storage.scanUsage({ order: "asc" })).map((row) => row.id), + (await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT)).map((row) => row.id), ["existing-usage"], ); }), @@ -238,62 +259,71 @@ export function createStorageConformance(factory: () => Promise) "transactions", "resolves parents only from prior entries and earlier writes", async ({ storage }) => { - await storage.commit([insertEntry(userEntry("root"))]); - await storage.commit([ - insertEntry(customEntry("child", "root")), - insertEntry(customEntry("grandchild", "child")), - ]); - deepStrictEqual(ids(await storage.scanBranch({ start: "grandchild", order: "oldestFirst" })), [ - "root", - "child", - "grandchild", - ]); + await storage.commit([insertEntry(userEntry("root"))], BACKGROUND_CONTEXT); + await storage.commit( + [insertEntry(customEntry("child", "root")), insertEntry(customEntry("grandchild", "child"))], + BACKGROUND_CONTEXT, + ); + deepStrictEqual( + ids(await storage.scanBranch({ start: "grandchild", order: "oldestFirst" }, BACKGROUND_CONTEXT)), + ["root", "child", "grandchild"], + ); await rejects( - storage.commit([ - insertEntry(customEntry("before-parent", "later-parent")), - insertEntry(customEntry("later-parent", "root")), - setValue(entryLabel("before-parent"), "transient"), - ]), + storage.commit( + [ + insertEntry(customEntry("before-parent", "later-parent")), + insertEntry(customEntry("later-parent", "root")), + setValue(entryLabel("before-parent"), "transient"), + ], + BACKGROUND_CONTEXT, + ), + ); + await rejects(storage.commit([insertEntry(customEntry("orphan", "missing"))], BACKGROUND_CONTEXT)); + await storage.commit( + [insertUsage({ id: "usage-is-not-parent", usage: usage(1, 1), adjustment: false })], + BACKGROUND_CONTEXT, + ); + await rejects( + storage.commit([insertEntry(customEntry("usage-child", "usage-is-not-parent"))], BACKGROUND_CONTEXT), ); - await rejects(storage.commit([insertEntry(customEntry("orphan", "missing"))])); - await storage.commit([insertUsage({ id: "usage-is-not-parent", usage: usage(1, 1), adjustment: false })]); - await rejects(storage.commit([insertEntry(customEntry("usage-child", "usage-is-not-parent"))])); deepStrictEqual( - await storage.getEntries(["before-parent", "later-parent", "orphan", "usage-child"]), + await storage.getEntries(["before-parent", "later-parent", "orphan", "usage-child"], BACKGROUND_CONTEXT), new Map(), ); - strictEqual(await storage.getValue(entryLabel("before-parent")), undefined); + strictEqual(await storage.getValue(entryLabel("before-parent"), BACKGROUND_CONTEXT), undefined); }, ), createCase(factory, "transactions", "places pending content under its reserved entry id", async ({ storage }) => { const entry = userEntry("reserved", null, "queued"); - await storage.commit([ - setValue(pendingEntry(entry.id), { type: "message", payload: entry.message }), - setValue(laneLeaf("main"), null), - ]); + await storage.commit( + [ + setValue(pendingEntry(entry.id), { type: "message", payload: entry.message }), + setValue(laneLeaf("main"), null), + ], + BACKGROUND_CONTEXT, + ); - strictEqual((await storage.getEntries([entry.id])).has(entry.id), false); - deepStrictEqual((await storage.getValue(pendingEntry(entry.id)))?.value, { + strictEqual((await storage.getEntries([entry.id], BACKGROUND_CONTEXT)).has(entry.id), false); + deepStrictEqual((await storage.getValue(pendingEntry(entry.id), BACKGROUND_CONTEXT))?.value, { type: "message", payload: entry.message, }); - strictEqual((await storage.getValue(laneLeaf("main")))?.value, null); + strictEqual((await storage.getValue(laneLeaf("main"), BACKGROUND_CONTEXT))?.value, null); - const placement = await storage.commit([ - insertEntry(entry), - deleteValue(pendingEntry(entry.id)), - setValue(laneLeaf("main"), entry.id), - ]); + const placement = await storage.commit( + [insertEntry(entry), deleteValue(pendingEntry(entry.id)), setValue(laneLeaf("main"), entry.id)], + BACKGROUND_CONTEXT, + ); deepStrictEqual( - await storage.getEntries([entry.id]), + await storage.getEntries([entry.id], BACKGROUND_CONTEXT), new Map([[entry.id, { ...entry, seq: placement.seqs[0], timestamp: placement.timestamp }]]), ); - strictEqual(await storage.getValue(pendingEntry(entry.id)), undefined); - deepStrictEqual(await storage.getValue(laneLeaf("main")), { + strictEqual(await storage.getValue(pendingEntry(entry.id), BACKGROUND_CONTEXT), undefined); + deepStrictEqual(await storage.getValue(laneLeaf("main"), BACKGROUND_CONTEXT), { address: laneLeaf("main"), value: entry.id, seq: placement.seqs[2], @@ -305,33 +335,39 @@ export function createStorageConformance(factory: () => Promise) "values", "sets, replaces, deletes, and recreates values without tombstones", async ({ storage }) => { - const first = await storage.commit([ - setValue(testValue("prefix/b"), 1), - setValue(testValue("prefix/a"), 2), - setValue(testValue("other"), 3), - setValue(testValue("prefix/\ue000"), 4), - setValue(testValue("prefix/\u{10000}"), 5), - setValue(testValue("prefix/a"), null), - ]); - deepStrictEqual(await storage.getValue(testValue("prefix/a")), { + const first = await storage.commit( + [ + setValue(testValue("prefix/b"), 1), + setValue(testValue("prefix/a"), 2), + setValue(testValue("other"), 3), + setValue(testValue("prefix/\ue000"), 4), + setValue(testValue("prefix/\u{10000}"), 5), + setValue(testValue("prefix/a"), null), + ], + BACKGROUND_CONTEXT, + ); + deepStrictEqual(await storage.getValue(testValue("prefix/a"), BACKGROUND_CONTEXT), { address: testValue("prefix/a"), value: null, seq: first.seqs[5], }); - const second = await storage.commit([ - deleteValue(testValue("prefix/a")), - deleteValue(testValue("absent")), - setValue(testValue("prefix/a"), "recreated"), - ]); + const second = await storage.commit( + [ + deleteValue(testValue("prefix/a")), + deleteValue(testValue("absent")), + setValue(testValue("prefix/a"), "recreated"), + ], + BACKGROUND_CONTEXT, + ); - deepStrictEqual(await storage.scanValues(testValuePrefix("prefix/")), [ + deepStrictEqual(await storage.scanValues(testValuePrefix("prefix/"), BACKGROUND_CONTEXT), [ { address: testValue("prefix/a"), value: "recreated", seq: second.seqs[2] }, { address: testValue("prefix/b"), value: 1, seq: first.seqs[0] }, { address: testValue("prefix/\ue000"), value: 4, seq: first.seqs[3] }, { address: testValue("prefix/\u{10000}"), value: 5, seq: first.seqs[4] }, ]); - strictEqual(await storage.getValue(testValue("absent")), undefined); + strictEqual(await storage.getValue(testValue("absent"), BACKGROUND_CONTEXT), undefined); }, ), @@ -340,20 +376,26 @@ export function createStorageConformance(factory: () => Promise) "values", "does not change historical stores during value-only commits", async ({ storage }) => { - await storage.commit([ - insertEntry(userEntry("root")), - insertUsage({ id: "historical-usage", usage: usage(2, 3), adjustment: false }), - ]); - const entriesBefore = await storage.scanEntries({ order: "asc" }); - const usageBefore = await storage.scanUsage({ order: "asc" }); - const statsBefore = await storage.getStats(); + await storage.commit( + [ + insertEntry(userEntry("root")), + insertUsage({ id: "historical-usage", usage: usage(2, 3), adjustment: false }), + ], + BACKGROUND_CONTEXT, + ); + const entriesBefore = await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT); + const usageBefore = await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT); + const statsBefore = await storage.getStats(BACKGROUND_CONTEXT); - const result = await storage.commit([setValue(testName, "first"), setValue(testName, "second")]); + const result = await storage.commit( + [setValue(testName, "first"), setValue(testName, "second")], + BACKGROUND_CONTEXT, + ); - deepStrictEqual(await storage.scanEntries({ order: "asc" }), entriesBefore); - deepStrictEqual(await storage.scanUsage({ order: "asc" }), usageBefore); - deepStrictEqual(await storage.getStats(), statsBefore); - deepStrictEqual(await storage.getValue(testName), { + deepStrictEqual(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT), entriesBefore); + deepStrictEqual(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT), usageBefore); + deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), statsBefore); + deepStrictEqual(await storage.getValue(testName, BACKGROUND_CONTEXT), { address: testName, value: "second", seq: result.seqs[1], @@ -363,58 +405,73 @@ export function createStorageConformance(factory: () => Promise) createCase(factory, "lists", "pages appends by global sequence and deletes whole lists", async ({ storage }) => { const address = testList("events"); - strictEqual((await storage.readList(address)).length, 0); - const result = await storage.commit([ - appendList(address, "a"), - setValue(testName, "gap"), - appendList(address, "b"), - appendList(address, "c"), - ]); - deepStrictEqual(await storage.readList(address), [ + strictEqual((await storage.readList(address, undefined, BACKGROUND_CONTEXT)).length, 0); + const result = await storage.commit( + [appendList(address, "a"), setValue(testName, "gap"), appendList(address, "b"), appendList(address, "c")], + BACKGROUND_CONTEXT, + ); + deepStrictEqual(await storage.readList(address, undefined, BACKGROUND_CONTEXT), [ { seq: result.seqs[0], value: "a" }, { seq: result.seqs[2], value: "b" }, { seq: result.seqs[3], value: "c" }, ]); - deepStrictEqual(await storage.readList(address, { limit: 2 }), [ + deepStrictEqual(await storage.readList(address, { limit: 2 }, BACKGROUND_CONTEXT), [ { seq: result.seqs[0], value: "a" }, { seq: result.seqs[2], value: "b" }, ]); - deepStrictEqual(await storage.readList(address, { cursor: { seq: result.seqs[0]! }, limit: 2 }), [ - { seq: result.seqs[2], value: "b" }, - { seq: result.seqs[3], value: "c" }, - ]); - deepStrictEqual(await storage.readList(address, { order: "desc", limit: 2 }), [ + deepStrictEqual( + await storage.readList(address, { cursor: { seq: result.seqs[0]! }, limit: 2 }, BACKGROUND_CONTEXT), + [ + { seq: result.seqs[2], value: "b" }, + { seq: result.seqs[3], value: "c" }, + ], + ); + deepStrictEqual(await storage.readList(address, { order: "desc", limit: 2 }, BACKGROUND_CONTEXT), [ { seq: result.seqs[3], value: "c" }, { seq: result.seqs[2], value: "b" }, ]); deepStrictEqual( - await storage.readList(address, { order: "desc", cursor: { seq: result.seqs[3]! }, limit: 2 }), + await storage.readList( + address, + { order: "desc", cursor: { seq: result.seqs[3]! }, limit: 2 }, + BACKGROUND_CONTEXT, + ), [ { seq: result.seqs[2], value: "b" }, { seq: result.seqs[0], value: "a" }, ], ); - await rejects(storage.readList(address, { limit: 0 })); - await rejects(storage.readList(address, { limit: Number.MAX_VALUE })); + await rejects(storage.readList(address, { limit: 0 }, BACKGROUND_CONTEXT)); + await rejects(storage.readList(address, { limit: Number.MAX_VALUE }, BACKGROUND_CONTEXT)); - await storage.commit([deleteList(address), deleteList(testList("absent")), appendList(address, "new")]); + await storage.commit( + [deleteList(address), deleteList(testList("absent")), appendList(address, "new")], + BACKGROUND_CONTEXT, + ); deepStrictEqual( - (await storage.readList(address)).map(({ value }) => value), + (await storage.readList(address, undefined, BACKGROUND_CONTEXT)).map(({ value }) => value), ["new"], ); }), createCase(factory, "lists", "clamps one read page without limiting list growth", async ({ storage }) => { const address = testList("large"); - await storage.commit(Array.from({ length: 10_001 }, (_, index) => appendList(address, index))); - const firstPage = await storage.readList(address); + await storage.commit( + Array.from({ length: 10_001 }, (_, index) => appendList(address, index)), + BACKGROUND_CONTEXT, + ); + const firstPage = await storage.readList(address, undefined, BACKGROUND_CONTEXT); strictEqual(firstPage.length, 1_000); - strictEqual((await storage.readList(address, { limit: 20_000 })).length, 10_000); + strictEqual((await storage.readList(address, { limit: 20_000 }, BACKGROUND_CONTEXT)).length, 10_000); strictEqual( ( - await storage.readList(address, { - cursor: { seq: firstPage.at(-1)!.seq }, - }) + await storage.readList( + address, + { + cursor: { seq: firstPage.at(-1)!.seq }, + }, + BACKGROUND_CONTEXT, + ) ).length, 1_000, ); @@ -426,34 +483,43 @@ export function createStorageConformance(factory: () => Promise) "commits mixed list writes atomically and rolls them back with siblings", async ({ storage }) => { const address = testList("atomic"); - const committed = await storage.commit([ - insertEntry(userEntry("mixed")), - appendList(address, "kept"), - setValue(testName, "kept"), - insertUsage({ id: "mixed-usage", usage: usage(1, 2), adjustment: false }), + const committed = await storage.commit( + [ + insertEntry(userEntry("mixed")), + appendList(address, "kept"), + setValue(testName, "kept"), + insertUsage({ id: "mixed-usage", usage: usage(1, 2), adjustment: false }), + ], + BACKGROUND_CONTEXT, + ); + deepStrictEqual(await storage.readList(address, undefined, BACKGROUND_CONTEXT), [ + { seq: committed.seqs[1], value: "kept" }, ]); - deepStrictEqual(await storage.readList(address), [{ seq: committed.seqs[1], value: "kept" }]); await rejects( - storage.commit([ - appendList(address, "transient"), - deleteValue(testName), - insertEntry(userEntry("mixed")), - ]), + storage.commit( + [appendList(address, "transient"), deleteValue(testName), insertEntry(userEntry("mixed"))], + BACKGROUND_CONTEXT, + ), ); - deepStrictEqual(await storage.readList(address), [{ seq: committed.seqs[1], value: "kept" }]); - strictEqual((await storage.getValue(testName))?.value, "kept"); + deepStrictEqual(await storage.readList(address, undefined, BACKGROUND_CONTEXT), [ + { seq: committed.seqs[1], value: "kept" }, + ]); + strictEqual((await storage.getValue(testName, BACKGROUND_CONTEXT))?.value, "kept"); }, ), createCase(factory, "entry queries", "stores custom entries with and without data", async ({ storage }) => { - const result = await storage.commit([ - insertEntry({ id: "without-data", parentId: null, type: "custom", customType: "marker" }), - insertEntry(customEntry("with-data", "without-data", "note", { nested: [1, 2] })), - ]); + const result = await storage.commit( + [ + insertEntry({ id: "without-data", parentId: null, type: "custom", customType: "marker" }), + insertEntry(customEntry("with-data", "without-data", "note", { nested: [1, 2] })), + ], + BACKGROUND_CONTEXT, + ); deepStrictEqual( - await storage.getEntries(["without-data", "with-data"]), + await storage.getEntries(["without-data", "with-data"], BACKGROUND_CONTEXT), new Map([ [ "without-data", @@ -483,28 +549,40 @@ export function createStorageConformance(factory: () => Promise) "entry queries", "scans global entries with explicit ranges, filters, orders, and limits", async ({ storage }) => { - const result = await storage.commit([ - insertEntry(userEntry("root")), - insertEntry(customEntry("note-1", "root", "note")), - insertEntry(customEntry("other", "note-1", "other")), - insertEntry(customEntry("note-2", "other", "note")), - insertEntry(userEntry("tail", "note-2")), - ]); + const result = await storage.commit( + [ + insertEntry(userEntry("root")), + insertEntry(customEntry("note-1", "root", "note")), + insertEntry(customEntry("other", "note-1", "other")), + insertEntry(customEntry("note-2", "other", "note")), + insertEntry(userEntry("tail", "note-2")), + ], + BACKGROUND_CONTEXT, + ); deepStrictEqual( ids( - await storage.scanEntries({ - type: "custom", - customType: "note", - fromSeq: result.seqs[1], - toSeq: result.seqs[3], - order: "desc", - }), + await storage.scanEntries( + { + type: "custom", + customType: "note", + fromSeq: result.seqs[1], + toSeq: result.seqs[3], + order: "desc", + }, + BACKGROUND_CONTEXT, + ), ), ["note-2", "note-1"], ); - deepStrictEqual(ids(await storage.scanEntries({ order: "asc", limit: 2 })), ["root", "note-1"]); - deepStrictEqual(ids(await storage.scanEntries({ order: "desc", limit: 2 })), ["tail", "note-2"]); + deepStrictEqual(ids(await storage.scanEntries({ order: "asc", limit: 2 }, BACKGROUND_CONTEXT)), [ + "root", + "note-1", + ]); + deepStrictEqual(ids(await storage.scanEntries({ order: "desc", limit: 2 }, BACKGROUND_CONTEXT)), [ + "tail", + "note-2", + ]); }, ), @@ -513,81 +591,106 @@ export function createStorageConformance(factory: () => Promise) "branch queries", "applies stops before filters and cursors before limits", async ({ storage }) => { - const result = await storage.commit([ - insertEntry(userEntry("root")), - insertEntry(customEntry("marker", "root", "marker")), - insertEntry(userEntry("middle", "marker")), - insertEntry(compactionEntry("compact", "middle")), - insertEntry(customEntry("note", "compact", "note")), - insertEntry(userEntry("leaf", "note")), - ]); + const result = await storage.commit( + [ + insertEntry(userEntry("root")), + insertEntry(customEntry("marker", "root", "marker")), + insertEntry(userEntry("middle", "marker")), + insertEntry(compactionEntry("compact", "middle")), + insertEntry(customEntry("note", "compact", "note")), + insertEntry(userEntry("leaf", "note")), + ], + BACKGROUND_CONTEXT, + ); deepStrictEqual( - ids(await storage.scanBranch({ start: "leaf", stopAtType: "compaction", type: "message" })), + ids( + await storage.scanBranch( + { start: "leaf", stopAtType: "compaction", type: "message" }, + BACKGROUND_CONTEXT, + ), + ), ["leaf"], ); deepStrictEqual( ids( - await storage.scanBranch({ - start: "leaf", - order: "oldestFirst", - stopAtId: "middle", - type: "custom", - }), + await storage.scanBranch( + { + start: "leaf", + order: "oldestFirst", + stopAtId: "middle", + type: "custom", + }, + BACKGROUND_CONTEXT, + ), ), ["marker"], ); deepStrictEqual( ids( - await storage.scanBranch({ - start: "leaf", - order: "newestFirst", - cursor: { seq: result.seqs[4] }, - limit: 2, - }), + await storage.scanBranch( + { + start: "leaf", + order: "newestFirst", + cursor: { seq: result.seqs[4] }, + limit: 2, + }, + BACKGROUND_CONTEXT, + ), ), ["compact", "middle"], ); deepStrictEqual( ids( - await storage.scanBranch({ - start: "leaf", - order: "oldestFirst", - cursor: { seq: result.seqs[1] }, - limit: 2, - }), + await storage.scanBranch( + { + start: "leaf", + order: "oldestFirst", + cursor: { seq: result.seqs[1] }, + limit: 2, + }, + BACKGROUND_CONTEXT, + ), ), ["middle", "compact"], ); - deepStrictEqual(ids(await storage.scanBranch({ start: "leaf", stopAtId: "leaf", type: "custom" })), []); - deepStrictEqual(ids(await storage.scanBranch({ start: "leaf", customType: "note" })), ["note"]); - await rejects(storage.scanBranch({ start: "missing" })); + deepStrictEqual( + ids(await storage.scanBranch({ start: "leaf", stopAtId: "leaf", type: "custom" }, BACKGROUND_CONTEXT)), + [], + ); + deepStrictEqual(ids(await storage.scanBranch({ start: "leaf", customType: "note" }, BACKGROUND_CONTEXT)), [ + "note", + ]); + await rejects(storage.scanBranch({ start: "missing" }, BACKGROUND_CONTEXT)); }, ), createCase(factory, "branch queries", "returns branch structure without payload fields", async ({ storage }) => { - const result = await storage.commit([ - insertEntry(userEntry("root")), - insertEntry(customEntry("child", "root", "note")), - ]); + const result = await storage.commit( + [insertEntry(userEntry("root")), insertEntry(customEntry("child", "root", "note"))], + BACKGROUND_CONTEXT, + ); - deepStrictEqual(await storage.scanBranchStructure({ start: "child", order: "oldestFirst" }), [ - { - id: "root", - parentId: null, - seq: result.seqs[0], - timestamp: result.timestamp, - type: "message", - }, - { - id: "child", - parentId: "root", - seq: result.seqs[1], - timestamp: result.timestamp, - type: "custom", - customType: "note", - }, - ]); + deepStrictEqual( + await storage.scanBranchStructure({ start: "child", order: "oldestFirst" }, BACKGROUND_CONTEXT), + [ + { + id: "root", + parentId: null, + seq: result.seqs[0], + timestamp: result.timestamp, + type: "message", + }, + { + id: "child", + parentId: "root", + seq: result.seqs[1], + timestamp: result.timestamp, + type: "custom", + customType: "note", + }, + ], + ); }), createCase( @@ -595,25 +698,31 @@ export function createStorageConformance(factory: () => Promise) "usage and stats", "scans the usage ledger with explicit ranges, orders, and limits", async ({ storage }) => { - const result = await storage.commit([ - insertUsage({ id: "usage-1", usage: usage(1, 1), adjustment: false }), - setValue(testName, "sequence gap"), - insertUsage({ id: "usage-2", usage: usage(2, 2), adjustment: false }), - insertUsage({ id: "usage-3", usage: usage(3, 3), adjustment: true }), - ]); + const result = await storage.commit( + [ + insertUsage({ id: "usage-1", usage: usage(1, 1), adjustment: false }), + setValue(testName, "sequence gap"), + insertUsage({ id: "usage-2", usage: usage(2, 2), adjustment: false }), + insertUsage({ id: "usage-3", usage: usage(3, 3), adjustment: true }), + ], + BACKGROUND_CONTEXT, + ); deepStrictEqual( - (await storage.scanUsage({ fromSeq: result.seqs[1], toSeq: result.seqs[2], order: "asc" })).map( - (row) => row.id, - ), + ( + await storage.scanUsage( + { fromSeq: result.seqs[1], toSeq: result.seqs[2], order: "asc" }, + BACKGROUND_CONTEXT, + ) + ).map((row) => row.id), ["usage-2"], ); deepStrictEqual( - (await storage.scanUsage({ order: "desc", limit: 2 })).map((row) => row.id), + (await storage.scanUsage({ order: "desc", limit: 2 }, BACKGROUND_CONTEXT)).map((row) => row.id), ["usage-3", "usage-2"], ); deepStrictEqual( - (await storage.scanUsage({ order: "asc", limit: 2 })).map((row) => row.id), + (await storage.scanUsage({ order: "asc", limit: 2 }, BACKGROUND_CONTEXT)).map((row) => row.id), ["usage-1", "usage-2"], ); }, @@ -624,22 +733,28 @@ export function createStorageConformance(factory: () => Promise) "usage and stats", "keeps stats equal to message count and ledger totals", async ({ storage }) => { - deepStrictEqual(await storage.getStats(), { messageCount: 0, usage: zeroUsage() }); + deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), { messageCount: 0, usage: zeroUsage() }); const firstUsage = usage(2, 3, { cacheWrite1h: 4, reasoning: 1 }); - await storage.commit([ - insertEntry(userEntry("message")), - insertUsage({ id: "usage-1", usage: firstUsage, adjustment: false }), - ]); - deepStrictEqual(await storage.getStats(), { messageCount: 1, usage: firstUsage }); + 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 }); const secondUsage = usage(5, 7, { cacheWrite1h: 6, reasoning: 2 }); - await storage.commit([ - insertEntry(customEntry("custom", "message")), - insertEntry(compactionEntry("compaction", "custom")), - insertUsage({ id: "usage-2", usage: secondUsage, adjustment: true }), - ]); - deepStrictEqual(await storage.getStats(), { + await storage.commit( + [ + insertEntry(customEntry("custom", "message")), + insertEntry(compactionEntry("compaction", "custom")), + insertUsage({ id: "usage-2", usage: secondUsage, adjustment: true }), + ], + BACKGROUND_CONTEXT, + ); + deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), { messageCount: 1, usage: { input: 7, @@ -666,12 +781,12 @@ export function createStorageConformance(factory: () => Promise) "serialization", "serializes back-to-back commits in admission order", async ({ storage }) => { - const first = storage.commit([insertEntry(userEntry("first"))]); - const second = storage.commit([insertEntry(userEntry("second", "first"))]); + const first = storage.commit([insertEntry(userEntry("first"))], BACKGROUND_CONTEXT); + const second = storage.commit([insertEntry(userEntry("second", "first"))], BACKGROUND_CONTEXT); const [firstResult, secondResult] = await Promise.all([first, second]); ok(firstResult.seqs[0]! < secondResult.seqs[0]!); - deepStrictEqual(ids(await storage.scanEntries({ order: "asc" })), ["first", "second"]); + deepStrictEqual(ids(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT)), ["first", "second"]); }, ), @@ -680,25 +795,25 @@ export function createStorageConformance(factory: () => Promise) "lifecycle", "seals admission, drains admitted commits, and closes idempotently", async ({ storage }) => { - const admitted = storage.commit([insertEntry(userEntry("admitted"))]); - const firstClose = storage.close(); - const secondClose = storage.close(); + const admitted = storage.commit([insertEntry(userEntry("admitted"))], BACKGROUND_CONTEXT); + const firstClose = storage.close(BACKGROUND_CONTEXT); + const secondClose = storage.close(BACKGROUND_CONTEXT); - await rejects(storage.getStats()); - await rejects(storage.commit([])); + await rejects(storage.getStats(BACKGROUND_CONTEXT)); + await rejects(storage.commit([], BACKGROUND_CONTEXT)); strictEqual((await admitted).seqs.length, 1); await Promise.all([firstClose, secondClose]); const rejectedReads = [ - storage.getEntries([]), - storage.getValue(testName), - storage.scanValues(testName), - storage.readList(testList("events")), - storage.scanBranch({ start: "admitted" }), - storage.scanBranchStructure({ start: "admitted" }), - storage.scanEntries({ order: "asc" }), - storage.scanUsage({ order: "asc" }), - storage.getStats(), + storage.getEntries([], BACKGROUND_CONTEXT), + storage.getValue(testName, BACKGROUND_CONTEXT), + storage.scanValues(testName, BACKGROUND_CONTEXT), + storage.readList(testList("events"), undefined, BACKGROUND_CONTEXT), + storage.scanBranch({ start: "admitted" }, BACKGROUND_CONTEXT), + storage.scanBranchStructure({ start: "admitted" }, BACKGROUND_CONTEXT), + storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT), + storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT), + storage.getStats(BACKGROUND_CONTEXT), ]; for (const read of rejectedReads) await rejects(read); }, diff --git a/packages/agent/src/harness/session/testing/instrumented-storage.ts b/packages/agent/src/harness/session/testing/instrumented-storage.ts index 8790161d1..34cc41aba 100644 --- a/packages/agent/src/harness/session/testing/instrumented-storage.ts +++ b/packages/agent/src/harness/session/testing/instrumented-storage.ts @@ -1,3 +1,4 @@ +import type { Context } from "../../context.ts"; import type { CommitResult, Write } from "../types.ts"; import { StorageDecorator } from "./storage-decorator.ts"; @@ -13,8 +14,8 @@ export class InstrumentedStorage extends StorageDecorator { this.commitAttempts.length = 0; } - override commit(writes: Write[]): Promise { + override commit(writes: Write[], context: Context): Promise { this.commitAttempts.push(writes); - return this.delegate.commit(writes); + return this.delegate.commit(writes, context); } } diff --git a/packages/agent/src/harness/session/testing/storage-decorator.ts b/packages/agent/src/harness/session/testing/storage-decorator.ts index fa8140d07..d22e6a775 100644 --- a/packages/agent/src/harness/session/testing/storage-decorator.ts +++ b/packages/agent/src/harness/session/testing/storage-decorator.ts @@ -1,3 +1,4 @@ +import type { Context } from "../../context.ts"; import type { CommitResult, Entry, @@ -20,47 +21,51 @@ export class StorageDecorator implements Storage { this.delegate = delegate; } - commit(writes: Write[]): Promise { - return this.delegate.commit(writes); + commit(writes: Write[], context: Context): Promise { + return this.delegate.commit(writes, context); } - getEntries(ids: string[]): Promise> { - return this.delegate.getEntries(ids); + getEntries(ids: string[], context: Context): Promise> { + return this.delegate.getEntries(ids, context); } - getValue(address: Value): Promise | undefined> { - return this.delegate.getValue(address); + getValue(address: Value, context: Context): Promise | undefined> { + return this.delegate.getValue(address, context); } - scanValues(prefix: Value): Promise[]> { - return this.delegate.scanValues(prefix); + scanValues(prefix: Value, context: Context): Promise[]> { + return this.delegate.scanValues(prefix, context); } - readList(address: ValueList, options?: ListReadOptions): Promise[]> { - return this.delegate.readList(address, options); + readList( + address: ValueList, + options: ListReadOptions | undefined, + context: Context, + ): Promise[]> { + return this.delegate.readList(address, options, context); } - scanBranch(query: StorageBranchScan): Promise { - return this.delegate.scanBranch(query); + scanBranch(query: StorageBranchScan, context: Context): Promise { + return this.delegate.scanBranch(query, context); } - scanBranchStructure(query: StorageBranchScan): Promise { - return this.delegate.scanBranchStructure(query); + scanBranchStructure(query: StorageBranchScan, context: Context): Promise { + return this.delegate.scanBranchStructure(query, context); } - scanEntries(query: EntryScan): Promise { - return this.delegate.scanEntries(query); + scanEntries(query: EntryScan, context: Context): Promise { + return this.delegate.scanEntries(query, context); } - scanUsage(query: UsageScan): Promise { - return this.delegate.scanUsage(query); + scanUsage(query: UsageScan, context: Context): Promise { + return this.delegate.scanUsage(query, context); } - getStats(): Promise { - return this.delegate.getStats(); + getStats(context: Context): Promise { + return this.delegate.getStats(context); } - close(): Promise { - return this.delegate.close(); + close(context: Context): Promise { + return this.delegate.close(context); } } diff --git a/packages/agent/src/harness/session/types.ts b/packages/agent/src/harness/session/types.ts index a0b96aa7e..fc5066bc9 100644 --- a/packages/agent/src/harness/session/types.ts +++ b/packages/agent/src/harness/session/types.ts @@ -2,6 +2,7 @@ import type { AssistantMessage, StopReason, Usage } from "@earendil-works/pi-ai" import type { AgentMessage, QueueMode, ThinkingLevel } from "../../types.ts"; import type { BranchPreparation } from "../compaction/branch-summarization.ts"; import type { CompactionPreparation, CompactionSettings } from "../compaction/compaction.ts"; +import type { Context } from "../context.ts"; import type { AgentHarnessStreamOptions } from "../types.ts"; import type { ListElement, ListReadOptions, ListWrite, StoredValue, Value, ValueList, ValueWrite } from "./values.ts"; @@ -54,7 +55,10 @@ export interface CustomEntry extends EntryBase { } /** Convert an application-defined custom entry into model context. */ -export type EntryProjector = (entry: CustomEntry) => AgentMessage[] | undefined | Promise; +export type EntryProjector = ( + entry: CustomEntry, + context: Context, +) => AgentMessage[] | undefined | Promise; export type Entry = MessageEntry | CompactionEntry | BranchSummaryEntry | CustomEntry; @@ -436,17 +440,21 @@ export interface SessionStats { } export interface Storage { - commit(writes: Write[]): Promise; - getEntries(ids: string[]): Promise>; - getValue(address: Value): Promise | undefined>; - scanValues(prefix: Value): Promise[]>; - readList(address: ValueList, options?: ListReadOptions): Promise[]>; - scanBranch(query: StorageBranchScan): Promise; - scanBranchStructure(query: StorageBranchScan): Promise; - scanEntries(query: EntryScan): Promise; - scanUsage(query: UsageScan): Promise; - getStats(): Promise; - close(): Promise; + commit(writes: Write[], context: Context): Promise; + getEntries(ids: string[], context: Context): Promise>; + getValue(address: Value, context: Context): Promise | undefined>; + scanValues(prefix: Value, context: Context): Promise[]>; + readList( + address: ValueList, + options: ListReadOptions | undefined, + context: Context, + ): Promise[]>; + scanBranch(query: StorageBranchScan, context: Context): Promise; + scanBranchStructure(query: StorageBranchScan, context: Context): Promise; + scanEntries(query: EntryScan, context: Context): Promise; + scanUsage(query: UsageScan, context: Context): Promise; + getStats(context: Context): Promise; + close(context: Context): Promise; } export interface SessionMetadata { @@ -471,49 +479,66 @@ export interface EntryQuery { } export interface SessionReader { - getEntries(ids: string[]): Promise>; - getValue(address: Value): Promise | undefined>; - scanValues(prefix: Value): Promise[]>; - readList(address: ValueList, options?: ListReadOptions): Promise[]>; + getEntries(ids: string[], context: Context): Promise>; + getValue(address: Value, context: Context): Promise | undefined>; + scanValues(prefix: Value, context: Context): Promise[]>; + readList( + address: ValueList, + options: ListReadOptions | undefined, + context: Context, + ): Promise[]>; } /** Callback-scoped write capability bound to one lane. */ export interface SessionMutator extends SessionReader { readonly lane: string; /** The mutation callback's sole commit. A second attempt rejects. */ - commit(writes: Write[]): Promise; + commit(writes: Write[], context: Context): Promise; } export interface SessionTree { - getLeafId(): Promise; - getEntry(id: string): Promise; - getStats(): Promise; - getValue(address: Value): Promise | undefined>; - scanValues(prefix: Value): Promise[]>; - readList(address: ValueList, options?: ListReadOptions): Promise[]>; - setValue(address: Value, next: NoInfer): Promise; - deleteValue(address: Value): Promise; - appendList(address: ValueList, element: NoInfer): Promise; - deleteList(address: ValueList): Promise; - getName(): Promise; - setName(name: string | undefined): Promise; - getLabel(targetId: string): Promise; - setLabel(targetId: string, label: string | undefined): Promise; - findEntries(query?: EntryQuery): Promise; - findEntry(query?: EntryQuery): Promise; - findEntriesOnBranch(query?: BranchScan): Promise; - findEntryOnBranch(query?: BranchScan): Promise; - appendMessage(message: AgentMessage): Promise; - appendCustomEntry(customType: string, data?: JsonValue): Promise; + getLeafId(context: Context): Promise; + getEntry(id: string, context: Context): Promise; + getStats(context: Context): Promise; + getValue(address: Value, context: Context): Promise | undefined>; + scanValues(prefix: Value, context: Context): Promise[]>; + readList( + address: ValueList, + options: ListReadOptions | undefined, + context: Context, + ): Promise[]>; + setValue(address: Value, next: NoInfer, context: Context): Promise; + deleteValue(address: Value, context: Context): Promise; + appendList(address: ValueList, element: NoInfer, context: Context): Promise; + deleteList(address: ValueList, context: Context): Promise; + getName(context: Context): Promise; + setName(name: string | undefined, context: Context): Promise; + getLabel(targetId: string, context: Context): Promise; + setLabel(targetId: string, label: string | undefined, context: Context): Promise; + findEntries(query: EntryQuery | undefined, context: Context): Promise; + findEntry(query: EntryQuery | undefined, context: Context): Promise; + findEntriesOnBranch(query: BranchScan | undefined, context: Context): Promise; + findEntryOnBranch(query: BranchScan | undefined, context: Context): Promise; + appendMessage(message: AgentMessage, context: Context): Promise; + appendCustomEntry(customType: string, data: JsonValue | undefined, context: Context): Promise; } export interface Session extends SessionTree, SessionReader { readonly metadata: TMetadata; readonly idGenerator: IdGenerator; view(lane: string): SessionTree; - mutate(lane: string, mutation: (mutator: SessionMutator) => T | Promise): Promise; - createLane(name: string, at: string | null, configuration: LaneConfiguration): Promise; - close(): Promise; + mutate( + lane: string, + mutation: (mutator: SessionMutator, context: Context) => T | Promise, + context: Context, + ): Promise; + createLane( + name: string, + at: string | null, + configuration: LaneConfiguration, + context: Context, + ): Promise; + close(context: Context): Promise; } export interface SessionCreateOptions { @@ -556,9 +581,9 @@ export interface SessionRepo< TCreateOptions extends { id?: string; parentSessionId?: string } = SessionCreateOptions, TListOptions = void, > { - create(options: TCreateOptions): Promise>; - open(metadata: TMetadata): Promise>; - list(options?: TListOptions): Promise; - delete(metadata: TMetadata): Promise; - fork(source: TMetadata, options: ForkOptions): Promise>; + create(options: TCreateOptions, context: Context): Promise>; + open(metadata: TMetadata, context: Context): Promise>; + list(options: TListOptions | undefined, context: Context): Promise; + delete(metadata: TMetadata, context: Context): Promise; + fork(source: TMetadata, options: ForkOptions, context: Context): Promise>; } diff --git a/packages/agent/test/harness/branch-summarization.test.ts b/packages/agent/test/harness/branch-summarization.test.ts index 913e582d7..09bbe74bf 100644 --- a/packages/agent/test/harness/branch-summarization.test.ts +++ b/packages/agent/test/harness/branch-summarization.test.ts @@ -1,6 +1,7 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { describe, expect, it } from "vitest"; import { collectEntriesForBranchSummary } from "../../src/harness/compaction/branch-summarization.ts"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import type { Entry, MessageEntry, SessionTree } from "../../src/harness/session/index.ts"; function message(text: string): AgentMessage { @@ -40,7 +41,7 @@ describe("v4 branch summarization", () => { const target = messageEntry("target", common.id, "target", 5); const session = branchReader([root, common, abandoned1, abandoned2, target]); - const result = await collectEntriesForBranchSummary(session, abandoned2.id, target.id); + const result = await collectEntriesForBranchSummary(session, abandoned2.id, target.id, BACKGROUND_CONTEXT); expect(result.commonAncestorId).toBe(common.id); expect(result.entries.map((entry) => entry.id)).toEqual([abandoned1.id, abandoned2.id]); expect(result.entries.some((entry) => entry.id === root.id)).toBe(false); @@ -48,9 +49,11 @@ describe("v4 branch summarization", () => { it("returns no entries when there was no previous leaf", async () => { const target = messageEntry("target", null, "target", 1); - expect(await collectEntriesForBranchSummary(branchReader([target]), null, target.id)).toEqual({ - entries: [], - commonAncestorId: null, - }); + expect(await collectEntriesForBranchSummary(branchReader([target]), null, target.id, BACKGROUND_CONTEXT)).toEqual( + { + entries: [], + commonAncestorId: null, + }, + ); }); }); diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index bc8949f13..9db0e385e 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -28,6 +28,7 @@ import { serializeConversation, shouldCompact, } from "../../src/harness/compaction/compaction.ts"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import { buildSessionContext } from "../../src/harness/session/context.ts"; import type { BranchSummaryEntry, @@ -333,7 +334,7 @@ describe("harness compaction", () => { ]); const u3 = createMessageEntry(createUserMessage("3"), compaction.id); const a3 = createMessageEntry(createAssistantMessage("c"), u3.id); - const loaded = await buildSessionContext([u1, a1, u2, a2, compaction, u3, a3]); + const loaded = await buildSessionContext([u1, a1, u2, a2, compaction, u3, a3], undefined, BACKGROUND_CONTEXT); expect(loaded).toHaveLength(5); expect(loaded[0]?.role).toBe("compactionSummary"); expect(loaded.map((message) => message.role)).toEqual([ @@ -358,7 +359,9 @@ describe("harness compaction", () => { expect(preparation).toBeDefined(); expect(preparation?.previousSummary).toBe("First summary"); expect(preparation?.retainedTail.length).toBeGreaterThan(0); - expect(preparation?.tokensBefore).toBe(estimateContextTokens(await buildSessionContext(pathEntries)).tokens); + expect(preparation?.tokensBefore).toBe( + estimateContextTokens(await buildSessionContext(pathEntries, undefined, BACKGROUND_CONTEXT)).tokens, + ); }); it("carries a previous compaction's retained tail into the next preparation", () => { diff --git a/packages/agent/test/harness/execution-primitives.test.ts b/packages/agent/test/harness/execution-primitives.test.ts index f5bf888f8..c10d5481f 100644 --- a/packages/agent/test/harness/execution-primitives.test.ts +++ b/packages/agent/test/harness/execution-primitives.test.ts @@ -1,9 +1,17 @@ import { fauxProvider } from "@earendil-works/pi-ai"; import { describe, expect, it, vi } from "vitest"; +import { + BACKGROUND_CONTEXT, + createContextKey, + withAbortSignal, + withContextValue, + withTelemetryContext, +} from "../../src/harness/context.ts"; import { HarnessEventBus } from "../../src/harness/events.ts"; import { BreakpointBarrier } from "../../src/harness/execution/breakpoint.ts"; import { AbortRequested, OperationEffectGate } from "../../src/harness/execution/effect-gate.ts"; import { HookRegistry } from "../../src/harness/hooks.ts"; +import { InMemoryTelemetryContext } from "../../src/index.ts"; function deferred(): { promise: Promise; resolve(): void } { let resolvePromise: (() => void) | undefined; @@ -126,6 +134,7 @@ describe("HookRegistry", () => { resources: {}, }, new OperationEffectGate(), + BACKGROUND_CONTEXT, ); expect(result).toMatchObject({ @@ -156,10 +165,10 @@ describe("HookRegistry", () => { const closed = new Error("closed"); const closedGate = new OperationEffectGate(); closedGate.close(closed); - expect(() => hooks.runWithGate("before_run", event, closedGate)).toThrow(closed); + expect(() => hooks.runWithGate("before_run", event, closedGate, BACKGROUND_CONTEXT)).toThrow(closed); expect(calls).toEqual([]); - const running = hooks.runWithGate("before_run", event, new OperationEffectGate()); + const running = hooks.runWithGate("before_run", event, new OperationEffectGate(), BACKGROUND_CONTEXT); expect(calls).toEqual(["first:start"]); hooks.on("before_run", () => { calls.push("late"); @@ -171,6 +180,77 @@ describe("HookRegistry", () => { expect(calls).toEqual(["first:start", "first:end", "second"]); }); + it("rejects pre-aborted invocations and combines admitted hook cancellation with the effect gate", async () => { + const event = { operation: "run" as const, lane: "main", runId: "run" }; + const preAbortedHooks = new HookRegistry(() => {}); + const handler = vi.fn(); + preAbortedHooks.on("before_drive", handler); + const preAbortedGate = new OperationEffectGate(); + const invocationController = new AbortController(); + const abortReason = new Error("invocation aborted"); + invocationController.abort(abortReason); + + expect(() => + preAbortedHooks.runWithGate( + "before_drive", + event, + preAbortedGate, + withAbortSignal(invocationController.signal, BACKGROUND_CONTEXT), + ), + ).toThrow(abortReason); + expect(handler).not.toHaveBeenCalled(); + expect(() => preAbortedGate.assertOpen()).not.toThrow(); + + const admittedHooks = new HookRegistry(() => {}); + const admittedGate = new OperationEffectGate(); + let admittedSignal: AbortSignal | undefined; + admittedHooks.on("before_drive", (_event, context) => { + admittedSignal = context.abortSignal; + }); + await admittedHooks.runWithGate("before_drive", event, admittedGate, BACKGROUND_CONTEXT); + expect(admittedSignal?.aborted).toBe(false); + + const gateReason = new Error("gate closed"); + admittedGate.close(gateReason); + expect(admittedSignal?.aborted).toBe(true); + expect(admittedSignal?.reason).toBe(gateReason); + }); + + it("passes tool handlers a child context of the active hook span", async () => { + const telemetry = new InMemoryTelemetryContext(); + const valueKey = createContextKey("hook.test.value"); + const hooks = new HookRegistry(() => {}); + let receivedValue: string | undefined; + hooks.on("before_tool", async (_event, context) => { + receivedValue = context.value(valueKey); + await context.telemetryContext.startSpan({ name: "handler.child" }, () => undefined); + return undefined; + }); + + await telemetry.startSpan({ name: "invocation" }, (invocationSpan) => + hooks.runToolWithGate( + "before_tool", + { + lane: "main", + runId: "run", + toolCallId: "call", + toolName: "tool", + args: {}, + }, + new OperationEffectGate(), + withContextValue(valueKey, "preserved", withTelemetryContext(invocationSpan, BACKGROUND_CONTEXT)), + ), + ); + + expect(receivedValue).toBe("preserved"); + const spans = telemetry.getSpans(); + const invocationSpan = spans.find((span) => span.name === "invocation"); + const hookSpan = spans.find((span) => span.name === "pi.harness.hook"); + const handlerSpan = spans.find((span) => span.name === "handler.child"); + expect(hookSpan?.parentId).toBe(invocationSpan?.id); + expect(handlerSpan?.parentId).toBe(hookSpan?.id); + }); + it("treats registration ids as optional metadata and fails before_drive closed", async () => { const errors: Error[] = []; const hooks = new HookRegistry((error) => { @@ -188,7 +268,12 @@ describe("HookRegistry", () => { hooks.on("before_drive", later, { id: "duplicate" }); await expect( - hooks.runWithGate("before_drive", { lane: "main", runId: "run", operation: "run" }, new OperationEffectGate()), + hooks.runWithGate( + "before_drive", + { lane: "main", runId: "run", operation: "run" }, + new OperationEffectGate(), + BACKGROUND_CONTEXT, + ), ).rejects.toBe(failure); expect(later).not.toHaveBeenCalled(); expect(errors).toEqual([failure]); @@ -221,6 +306,7 @@ describe("HookRegistry", () => { systemPrompt: "base", }, new OperationEffectGate(), + BACKGROUND_CONTEXT, ); expect(result).toEqual({ messages, systemPrompt: "final" }); @@ -247,6 +333,7 @@ describe("HookRegistry", () => { streamOptions: { headers: { a: "1", b: "2" }, metadata: { x: 1 } }, }, new OperationEffectGate(), + BACKGROUND_CONTEXT, ); expect(result).toEqual({ streamOptions: { @@ -272,6 +359,7 @@ describe("HookRegistry", () => { isError: true, }, new OperationEffectGate(), + BACKGROUND_CONTEXT, ); expect(result).toEqual({ content: [{ type: "text", text: "patched" }], isError: false }); }); @@ -299,6 +387,7 @@ describe("HookRegistry", () => { }, }, new OperationEffectGate(), + BACKGROUND_CONTEXT, ); expect(result).toEqual({ decline: false, summary: selected }); expect(errors).toHaveLength(1); @@ -322,11 +411,13 @@ describe("HookRegistry", () => { const abortFirstGate = new OperationEffectGate(); const cancellation = Promise.resolve(); abortFirstGate.beginAbort(cancellation); - expect(() => hooks.runWithGate("before_drive", event, abortFirstGate)).toThrow(AbortRequested); + expect(() => hooks.runWithGate("before_drive", event, abortFirstGate, BACKGROUND_CONTEXT)).toThrow( + AbortRequested, + ); expect(calls).toEqual([]); const startFirstGate = new OperationEffectGate(); - const running = hooks.runWithGate("before_drive", event, startFirstGate); + const running = hooks.runWithGate("before_drive", event, startFirstGate, BACKGROUND_CONTEXT); expect(calls).toEqual(["first:start"]); startFirstGate.beginAbort(cancellation); startFirstGate.signalAbort(); @@ -339,18 +430,18 @@ describe("HookRegistry", () => { describe("HarnessEventBus", () => { it("buffers between snapshot and start, then delivers each event once in order", async () => { const bus = new HarnessEventBus(); - const watcher = bus.watch({ leafId: null }, () => true); - await bus.emit({ type: "run_start", runId: "one", lane: "main" }); + const watcher = bus.watch({ leafId: null }, () => true, BACKGROUND_CONTEXT); + await bus.emit({ type: "run_start", runId: "one", lane: "main" }, BACKGROUND_CONTEXT); const seen: string[] = []; watcher.start((event) => { seen.push(`${event.type}:${"runId" in event ? event.runId : ""}`); }); - await bus.emit({ type: "run_start", runId: "two", lane: "main" }); + await bus.emit({ type: "run_start", runId: "two", lane: "main" }, BACKGROUND_CONTEXT); await new Promise((resolve) => setTimeout(resolve, 0)); expect(watcher.snapshot).toEqual({ leafId: null }); expect(seen).toEqual(["run_start:one", "run_start:two"]); watcher.unsubscribe(); - await bus.emit({ type: "run_start", runId: "three", lane: "main" }); + await bus.emit({ type: "run_start", runId: "three", lane: "main" }, BACKGROUND_CONTEXT); expect(seen).toHaveLength(2); }); @@ -363,13 +454,16 @@ describe("HarnessEventBus", () => { bus.on("config_update", (event) => { if (event.property === "activeTools") observed = event.value; }); - await bus.emit({ - type: "config_update", - property: "activeTools", - value: ["read"], - previous: [], - lane: "main", - }); + await bus.emit( + { + type: "config_update", + property: "activeTools", + value: ["read"], + previous: [], + lane: "main", + }, + BACKGROUND_CONTEXT, + ); expect(observed).toEqual(["read"]); }); @@ -387,9 +481,9 @@ describe("HarnessEventBus", () => { seen.push(`${event.runId}:end`); }); - const one = bus.emit({ type: "run_start", runId: "one", lane: "main" }); + const one = bus.emit({ type: "run_start", runId: "one", lane: "main" }, BACKGROUND_CONTEXT); await started.promise; - const two = bus.emit({ type: "run_start", runId: "two", lane: "main" }); + const two = bus.emit({ type: "run_start", runId: "two", lane: "main" }, BACKGROUND_CONTEXT); await new Promise((resolve) => setTimeout(resolve, 0)); expect(seen).toEqual(["one:start"]); release.resolve(); @@ -403,7 +497,7 @@ describe("HarnessEventBus", () => { bus.on("handler_error", (event) => { failures.push(event.error); }); - const watcher = bus.watch({}, (event) => event.type === "run_start"); + const watcher = bus.watch({}, (event) => event.type === "run_start", BACKGROUND_CONTEXT); const seen: string[] = []; watcher.start((event) => { if (event.type !== "run_start") return; @@ -411,8 +505,8 @@ describe("HarnessEventBus", () => { seen.push(event.runId); }); - await bus.emit({ type: "run_start", runId: "one", lane: "main" }); - await bus.emit({ type: "run_start", runId: "two", lane: "main" }); + await bus.emit({ type: "run_start", runId: "one", lane: "main" }, BACKGROUND_CONTEXT); + await bus.emit({ type: "run_start", runId: "two", lane: "main" }, BACKGROUND_CONTEXT); await new Promise((resolve) => setTimeout(resolve, 0)); expect(seen).toEqual(["two"]); expect(failures).toEqual(["watcher failed"]); @@ -427,7 +521,7 @@ describe("HarnessEventBus", () => { bus.on("handler_error", (event) => { failures.push(event.error); }); - await bus.emit({ type: "run_start", runId: "run", lane: "main" }); + await bus.emit({ type: "run_start", runId: "run", lane: "main" }, BACKGROUND_CONTEXT); expect(failures).toEqual(["listener failed"]); }); }); diff --git a/packages/agent/test/harness/instrumented-storage.test.ts b/packages/agent/test/harness/instrumented-storage.test.ts index fabbfe450..c9eec4463 100644 --- a/packages/agent/test/harness/instrumented-storage.test.ts +++ b/packages/agent/test/harness/instrumented-storage.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it, vi } from "vitest"; +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"; @@ -47,10 +48,10 @@ describe("InstrumentedStorage", () => { const firstTransaction: Write[] = [storedValues.setValue(storedValues.sessionName, "first")]; const secondTransaction: Write[] = [storedValues.setValue(storedValues.sessionName, "second")]; - const firstCommit = storage.commit(firstTransaction); + const firstCommit = storage.commit(firstTransaction, BACKGROUND_CONTEXT); expect(firstCommit).toBe(delegate.lastCommit); expect(storage.getCommitAttempts()).toEqual([firstTransaction]); - const secondCommit = storage.commit(secondTransaction); + const secondCommit = storage.commit(secondTransaction, BACKGROUND_CONTEXT); expect(secondCommit).toBe(delegate.lastCommit); expect(storage.getCommitAttempts()).toEqual([firstTransaction, secondTransaction]); @@ -61,74 +62,87 @@ describe("InstrumentedStorage", () => { const secondResult = { firstSeq: 2, seqs: [2], timestamp: 20 }; delegate.resolveNextCommit(secondResult); expect(await secondCommit).toBe(secondResult); - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); }); it("records the transaction reference passed to the delegate", async () => { const delegate = new ControlledCommitStorage(); const storage = new InstrumentedStorage(delegate); const transaction: Write[] = [storedValues.setValue(storedValues.sessionName, "value")]; - const commit = storage.commit(transaction); + const commit = storage.commit(transaction, BACKGROUND_CONTEXT); expect(storage.getCommitAttempts()[0]).toBe(transaction); delegate.resolveNextCommit({ firstSeq: 1, seqs: [1], timestamp: 10 }); await commit; - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); }); it("clears recorded attempts between phases without affecting the delegate", async () => { const delegate = new MemoryStorage({ now: () => 100 }); const storage = new InstrumentedStorage(delegate); - await storage.commit([storedValues.setValue(storedValues.sessionName, "first")]); + await storage.commit([storedValues.setValue(storedValues.sessionName, "first")], BACKGROUND_CONTEXT); storage.clearCommitAttempts(); expect(storage.getCommitAttempts()).toEqual([]); - expect(await storage.getValue(storedValues.sessionName)).toMatchObject({ value: "first" }); + expect(await storage.getValue(storedValues.sessionName, BACKGROUND_CONTEXT)).toMatchObject({ value: "first" }); const secondTransaction: Write[] = [storedValues.setValue(storedValues.sessionName, "second")]; - await storage.commit(secondTransaction); + await storage.commit(secondTransaction, BACKGROUND_CONTEXT); expect(storage.getCommitAttempts()).toEqual([secondTransaction]); - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); }); it("delegates every read and query without recording synthetic writes", async () => { const delegate = new MemoryStorage({ now: () => 100 }); const storage = new InstrumentedStorage(delegate); const events = storedValues.list("test.events"); - await storage.commit([ - sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "note" }), - storedValues.setValue(storedValues.sessionName, "session"), - storedValues.appendList(events, "event"), - sessionWrites.insertUsage({ - id: "usage", - adjustment: false, - 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: "root", parentId: null, type: "custom", customType: "note" }), + storedValues.setValue(storedValues.sessionName, "session"), + storedValues.appendList(events, "event"), + sessionWrites.insertUsage({ + id: "usage", + adjustment: false, + usage: { + input: 1, + output: 2, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 3, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }), + ], + BACKGROUND_CONTEXT, + ); - expect(await storage.getEntries(["root"])).toEqual(await delegate.getEntries(["root"])); - expect(await storage.getValue(storedValues.sessionName)).toEqual( - await delegate.getValue(storedValues.sessionName), + expect(await storage.getEntries(["root"], BACKGROUND_CONTEXT)).toEqual( + await delegate.getEntries(["root"], BACKGROUND_CONTEXT), ); - expect(await storage.scanValues(storedValues.sessionName)).toEqual( - await delegate.scanValues(storedValues.sessionName), + expect(await storage.getValue(storedValues.sessionName, BACKGROUND_CONTEXT)).toEqual( + await delegate.getValue(storedValues.sessionName, BACKGROUND_CONTEXT), ); - expect(await storage.readList(events)).toEqual(await delegate.readList(events)); - expect(await storage.scanBranch({ start: "root" })).toEqual(await delegate.scanBranch({ start: "root" })); - expect(await storage.scanBranchStructure({ start: "root" })).toEqual( - await delegate.scanBranchStructure({ start: "root" }), + expect(await storage.scanValues(storedValues.sessionName, BACKGROUND_CONTEXT)).toEqual( + await delegate.scanValues(storedValues.sessionName, BACKGROUND_CONTEXT), ); - expect(await storage.scanEntries({ order: "asc" })).toEqual(await delegate.scanEntries({ order: "asc" })); - expect(await storage.scanUsage({ order: "asc" })).toEqual(await delegate.scanUsage({ order: "asc" })); - expect(await storage.getStats()).toEqual(await delegate.getStats()); + expect(await storage.readList(events, undefined, BACKGROUND_CONTEXT)).toEqual( + await delegate.readList(events, undefined, BACKGROUND_CONTEXT), + ); + expect(await storage.scanBranch({ start: "root" }, BACKGROUND_CONTEXT)).toEqual( + await delegate.scanBranch({ start: "root" }, BACKGROUND_CONTEXT), + ); + expect(await storage.scanBranchStructure({ start: "root" }, BACKGROUND_CONTEXT)).toEqual( + await delegate.scanBranchStructure({ start: "root" }, BACKGROUND_CONTEXT), + ); + expect(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT)).toEqual( + await delegate.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT), + ); + expect(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT)).toEqual( + await delegate.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT), + ); + expect(await storage.getStats(BACKGROUND_CONTEXT)).toEqual(await delegate.getStats(BACKGROUND_CONTEXT)); expect(storage.getCommitAttempts()).toHaveLength(1); - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); }); it("records list appends without reading the target list", async () => { @@ -137,25 +151,28 @@ describe("InstrumentedStorage", () => { const storage = new InstrumentedStorage(delegate); const events = storedValues.list("test.events"); - await storage.commit([storedValues.appendList(events, "event")]); + await storage.commit([storedValues.appendList(events, "event")], BACKGROUND_CONTEXT); expect(readList).not.toHaveBeenCalled(); expect(storage.getCommitAttempts()).toEqual([ [{ kind: "list", op: "append", namespace: "test.events", key: "", value: "event" }], ]); - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); }); it("delegates close idempotence and admitted commit draining", async () => { const delegate = new MemoryStorage({ now: () => 100 }); const storage = new InstrumentedStorage(delegate); - const admitted = storage.commit([storedValues.setValue(storedValues.sessionName, "admitted")]); + const admitted = storage.commit( + [storedValues.setValue(storedValues.sessionName, "admitted")], + BACKGROUND_CONTEXT, + ); - const firstClose = storage.close(); - const secondClose = storage.close(); + const firstClose = storage.close(BACKGROUND_CONTEXT); + const secondClose = storage.close(BACKGROUND_CONTEXT); await admitted; await Promise.all([firstClose, secondClose]); - await expect(storage.getStats()).rejects.toThrow("MemoryStorage is closed"); + await expect(storage.getStats(BACKGROUND_CONTEXT)).rejects.toThrow("MemoryStorage is closed"); expect(storage.getCommitAttempts()).toHaveLength(1); }); }); diff --git a/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts b/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts index 53f20580f..623360175 100644 --- a/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts +++ b/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts @@ -1,4 +1,5 @@ import { describe, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; import { type JsonlSessionMetadata, JsonlSessionRepo } from "../../src/harness/session/jsonl/index.ts"; import { @@ -33,20 +34,28 @@ async function createConformanceRepo() { jsonlRepo = new JsonlSessionRepo({ fileSystem, sessionsRoot: "sessions", now: () => NOW }); return { create: (options: { id?: string; parentSessionId?: string }) => - jsonlRepo.create({ ...options, cwd: CONFORMANCE_CWD }), - open: (metadata: JsonlSessionMetadata) => jsonlRepo.open(metadata), - list: () => jsonlRepo.list({ cwd: CONFORMANCE_CWD }), - delete: (metadata: JsonlSessionMetadata) => jsonlRepo.delete(metadata), - fork: (source: JsonlSessionMetadata, options: ForkOptions) => jsonlRepo.fork(source, options), + jsonlRepo.create({ ...options, cwd: CONFORMANCE_CWD }, BACKGROUND_CONTEXT), + open: (metadata: JsonlSessionMetadata) => jsonlRepo.open(metadata, BACKGROUND_CONTEXT), + list: () => jsonlRepo.list({ cwd: CONFORMANCE_CWD }, BACKGROUND_CONTEXT), + delete: (metadata: JsonlSessionMetadata) => jsonlRepo.delete(metadata, BACKGROUND_CONTEXT), + fork: (source: JsonlSessionMetadata, options: ForkOptions) => jsonlRepo.fork(source, options, BACKGROUND_CONTEXT), }; } registerConformance("JsonlSessionRepo conformance", [ - ...createSessionRepoLifecycleConformance(createConformanceRepo, () => jsonlRepo.close()), - ...createSessionRepoOwnershipConformance(createConformanceRepo, () => jsonlRepo.close()), - ...createSessionRepoMessageConformance(createConformanceRepo, () => jsonlRepo.close()), - ...createSessionRepoForkBehaviorConformance(createConformanceRepo, () => jsonlRepo.close()), + ...createSessionRepoLifecycleConformance(createConformanceRepo, () => + jsonlRepo.close(BACKGROUND_CONTEXT), + ), + ...createSessionRepoOwnershipConformance(createConformanceRepo, () => + jsonlRepo.close(BACKGROUND_CONTEXT), + ), + ...createSessionRepoMessageConformance(createConformanceRepo, () => + jsonlRepo.close(BACKGROUND_CONTEXT), + ), + ...createSessionRepoForkBehaviorConformance(createConformanceRepo, () => + jsonlRepo.close(BACKGROUND_CONTEXT), + ), ...createSessionRepoForkSourceSnapshotConformance(createConformanceRepo, () => - jsonlRepo.close(), + jsonlRepo.close(BACKGROUND_CONTEXT), ), ]); diff --git a/packages/agent/test/harness/jsonl-session-repo.test.ts b/packages/agent/test/harness/jsonl-session-repo.test.ts index fa17b84a6..79bdbe413 100644 --- a/packages/agent/test/harness/jsonl-session-repo.test.ts +++ b/packages/agent/test/harness/jsonl-session-repo.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; import { JSONL_STORAGE_VERSION, JsonlSessionRepo } from "../../src/harness/session/jsonl/index.ts"; import { getOrThrow } from "../../src/harness/types.ts"; @@ -10,7 +11,10 @@ describe("JsonlSessionRepo cwd-scoped lifecycle", () => { it("persists metadata and filters discovery by cwd", async () => { const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() }); const repo = new JsonlSessionRepo({ fileSystem, sessionsRoot: "sessions", now: () => NOW }); - const session = await repo.create({ id: "child", cwd: "/workspace", parentSessionId: "parent" }); + const session = await repo.create( + { id: "child", cwd: "/workspace", parentSessionId: "parent" }, + BACKGROUND_CONTEXT, + ); const metadata = session.metadata; expect(metadata).toMatchObject({ @@ -23,10 +27,10 @@ describe("JsonlSessionRepo cwd-scoped lifecycle", () => { expect(metadata.path).toContain("/sessions/--workspace--/"); expect(metadata.path.endsWith("_child.jsonl")).toBe(true); expect(Number.isFinite(metadata.modifiedAt)).toBe(true); - await session.close(); + await session.close(BACKGROUND_CONTEXT); - expect(await repo.list({ cwd: "/other" })).toEqual([]); - expect(await repo.list({ cwd: "/workspace" })).toEqual([metadata]); + expect(await repo.list({ cwd: "/other" }, BACKGROUND_CONTEXT)).toEqual([]); + expect(await repo.list({ cwd: "/workspace" }, BACKGROUND_CONTEXT)).toEqual([metadata]); const firstLine = getOrThrow(await fileSystem.readTextLines(metadata.path, { maxLines: 1 }))[0]; expect(JSON.parse(firstLine!)).toEqual({ v: 4, @@ -37,7 +41,7 @@ describe("JsonlSessionRepo cwd-scoped lifecycle", () => { cwd: "/workspace", parentSessionId: "parent", }); - await repo.close(); + await repo.close(BACKGROUND_CONTEXT); }); it("discovers legacy v3 session files without rewriting them", async () => { @@ -65,7 +69,7 @@ describe("JsonlSessionRepo cwd-scoped lifecycle", () => { getOrThrow(await fileSystem.writeFile(path, content)); const before = getOrThrow(await fileSystem.readTextFile(path)); - const [metadata] = await repo.list({ cwd: "/workspace" }); + const [metadata] = await repo.list({ cwd: "/workspace" }, BACKGROUND_CONTEXT); const after = getOrThrow(await fileSystem.readTextFile(path)); expect(metadata).toMatchObject({ @@ -78,26 +82,26 @@ describe("JsonlSessionRepo cwd-scoped lifecycle", () => { }); expect(Number.isFinite(metadata?.modifiedAt)).toBe(true); expect(after).toBe(before); - await repo.close(); + await repo.close(BACKGROUND_CONTEXT); }); it("keeps fork destinations claimed until close and rejects deleting open sessions", async () => { const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() }); const repo = new JsonlSessionRepo({ fileSystem, sessionsRoot: "sessions", now: () => NOW }); - const source = await repo.create({ id: "source", cwd: "/workspace" }); - const fork = await repo.fork(source.metadata, { id: "fork" }); + const source = await repo.create({ id: "source", cwd: "/workspace" }, BACKGROUND_CONTEXT); + const fork = await repo.fork(source.metadata, { id: "fork" }, BACKGROUND_CONTEXT); - await expect(repo.open(fork.metadata)).rejects.toThrow("already open"); - await expect(repo.delete(fork.metadata)).rejects.toThrow("open"); - await fork.close(); + await expect(repo.open(fork.metadata, BACKGROUND_CONTEXT)).rejects.toThrow("already open"); + await expect(repo.delete(fork.metadata, BACKGROUND_CONTEXT)).rejects.toThrow("open"); + await fork.close(BACKGROUND_CONTEXT); - const reopened = await repo.open(fork.metadata); - await reopened.close(); - await repo.delete(fork.metadata); - await expect(repo.open(fork.metadata)).rejects.toThrow("does not exist"); + const reopened = await repo.open(fork.metadata, BACKGROUND_CONTEXT); + await reopened.close(BACKGROUND_CONTEXT); + await repo.delete(fork.metadata, BACKGROUND_CONTEXT); + await expect(repo.open(fork.metadata, BACKGROUND_CONTEXT)).rejects.toThrow("does not exist"); - await source.close(); - await repo.close(); + await source.close(BACKGROUND_CONTEXT); + await repo.close(BACKGROUND_CONTEXT); }); it("rejects concurrent creates for the same working-directory id", async () => { @@ -105,35 +109,40 @@ describe("JsonlSessionRepo cwd-scoped lifecycle", () => { const repo = new JsonlSessionRepo({ fileSystem, sessionsRoot: "sessions", now: () => NOW }); const results = await Promise.allSettled([ - repo.create({ id: "session", cwd: "/workspace" }), - repo.create({ id: "session", cwd: "/workspace" }), + repo.create({ id: "session", cwd: "/workspace" }, BACKGROUND_CONTEXT), + repo.create({ id: "session", cwd: "/workspace" }, BACKGROUND_CONTEXT), ]); expect(results.filter((result) => result.status === "fulfilled")).toHaveLength(1); expect(results.filter((result) => result.status === "rejected")).toHaveLength(1); - expect(await repo.list({ cwd: "/workspace" })).toHaveLength(1); + expect(await repo.list({ cwd: "/workspace" }, BACKGROUND_CONTEXT)).toHaveLength(1); for (const result of results) { - if (result.status === "fulfilled") await result.value.close(); + if (result.status === "fulfilled") await result.value.close(BACKGROUND_CONTEXT); } - await repo.close(); + await repo.close(BACKGROUND_CONTEXT); }); it("allows the same id to be active in different working directories", async () => { const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() }); const repo = new JsonlSessionRepo({ fileSystem, sessionsRoot: "sessions", now: () => NOW }); - const first = await repo.create({ id: "shared", cwd: "/workspace-a" }); - const second = await repo.create({ id: "shared", cwd: "/workspace-b" }); + const first = await repo.create({ id: "shared", cwd: "/workspace-a" }, BACKGROUND_CONTEXT); + const second = await repo.create({ id: "shared", cwd: "/workspace-b" }, BACKGROUND_CONTEXT); expect(first.metadata.path).not.toBe(second.metadata.path); - await expect(repo.create({ id: "shared", cwd: "/workspace-a" })).rejects.toThrow("already exists"); - expect((await repo.list()).map(({ cwd, id }) => ({ cwd, id }))).toEqual([ + await expect(repo.create({ id: "shared", cwd: "/workspace-a" }, BACKGROUND_CONTEXT)).rejects.toThrow( + "already exists", + ); + expect((await repo.list(undefined, BACKGROUND_CONTEXT)).map(({ cwd, id }) => ({ cwd, id }))).toEqual([ { cwd: "/workspace-a", id: "shared" }, { cwd: "/workspace-b", id: "shared" }, ]); - await Promise.all([first.close(), second.close()]); - const reopened = await Promise.all([repo.open(first.metadata), repo.open(second.metadata)]); - await Promise.all(reopened.map((session) => session.close())); - await repo.close(); + await Promise.all([first.close(BACKGROUND_CONTEXT), second.close(BACKGROUND_CONTEXT)]); + const reopened = await Promise.all([ + repo.open(first.metadata, BACKGROUND_CONTEXT), + repo.open(second.metadata, BACKGROUND_CONTEXT), + ]); + await Promise.all(reopened.map((session) => session.close(BACKGROUND_CONTEXT))); + await repo.close(BACKGROUND_CONTEXT); }); }); diff --git a/packages/agent/test/harness/jsonl-storage-conformance.test.ts b/packages/agent/test/harness/jsonl-storage-conformance.test.ts index 3cf3e9f3c..525a007a8 100644 --- a/packages/agent/test/harness/jsonl-storage-conformance.test.ts +++ b/packages/agent/test/harness/jsonl-storage-conformance.test.ts @@ -1,4 +1,5 @@ import { describe, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; import { JSONL_FORMAT_VERSION, JsonlStorage } from "../../src/harness/session/jsonl/index.ts"; import { @@ -36,10 +37,11 @@ registerConformance( createdAt: NOW, cwd: "/workspace", }, + BACKGROUND_CONTEXT, ); return { storage, - [Symbol.asyncDispose]: () => storage.close(), + [Symbol.asyncDispose]: () => storage.close(BACKGROUND_CONTEXT), } satisfies StorageFixture; }), ); diff --git a/packages/agent/test/harness/jsonl-storage.test.ts b/packages/agent/test/harness/jsonl-storage.test.ts index 1501940b5..34dcff9af 100644 --- a/packages/agent/test/harness/jsonl-storage.test.ts +++ b/packages/agent/test/harness/jsonl-storage.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; import * as sessionWrites from "../../src/harness/session/commit.ts"; import { JSONL_FORMAT_VERSION, JsonlStorage, type JsonlStorageHeader } from "../../src/harness/session/jsonl/index.ts"; @@ -62,20 +63,25 @@ describe("JsonlStorage snapshot creation", () => { it("atomically publishes, advances, and reopens a prepared snapshot", async () => { const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() }); const options = { fileSystem, path: "fork.jsonl", now: () => NOW }; - const storage = await JsonlStorage.createFromSnapshot(options, header("fork"), preparedSnapshot()); + const storage = await JsonlStorage.createFromSnapshot( + options, + header("fork"), + preparedSnapshot(), + BACKGROUND_CONTEXT, + ); const lines = getOrThrow(await fileSystem.readTextFile("fork.jsonl")) .trimEnd() .split("\n"); expect(lines.slice(1).every((line) => !Array.isArray(JSON.parse(line)))).toBe(true); - expect((await storage.getEntries(["root"])).get("root")?.timestamp).toBe(NOW); - expect((await storage.getValue(storedValues.laneLeaf("main")))?.value).toBe("root"); - expect((await storage.getValue(storedValues.laneState("main")))?.value).toEqual({ + expect((await storage.getEntries(["root"], BACKGROUND_CONTEXT)).get("root")?.timestamp).toBe(NOW); + expect((await storage.getValue(storedValues.laneLeaf("main"), BACKGROUND_CONTEXT))?.value).toBe("root"); + expect((await storage.getValue(storedValues.laneState("main"), BACKGROUND_CONTEXT))?.value).toEqual({ currentOperationId: null, pendingNextRun: [], }); - expect((await storage.getValue(storedValues.sessionName))?.value).toBe("forked"); - expect(await storage.getStats()).toEqual({ + expect((await storage.getValue(storedValues.sessionName, BACKGROUND_CONTEXT))?.value).toBe("forked"); + expect(await storage.getStats(BACKGROUND_CONTEXT)).toEqual({ messageCount: 1, usage: { input: 0, @@ -86,17 +92,22 @@ describe("JsonlStorage snapshot creation", () => { cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, }); - expect(await storage.scanUsage({ order: "asc" })).toEqual([]); - expect(await storage.readList(storedValues.list("test.events"))).toEqual([{ seq: 5, value: "event" }]); - expect((await storage.commit([storedValues.setValue(storedValues.sessionName, "updated")])).firstSeq).toBe(9); + expect(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT)).toEqual([]); + expect(await storage.readList(storedValues.list("test.events"), undefined, BACKGROUND_CONTEXT)).toEqual([ + { seq: 5, value: "event" }, + ]); + expect( + (await storage.commit([storedValues.setValue(storedValues.sessionName, "updated")], BACKGROUND_CONTEXT)) + .firstSeq, + ).toBe(9); expect(getOrThrow(await fileSystem.exists("fork.jsonl.tmp"))).toBe(false); - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); - const reopened = await JsonlStorage.open(options); - expect((await reopened.getEntries(["root"])).has("root")).toBe(true); - expect((await reopened.getValue(storedValues.sessionName))?.value).toBe("updated"); - expect((await reopened.commit([])).firstSeq).toBe(10); - await reopened.close(); + const reopened = await JsonlStorage.open(options, BACKGROUND_CONTEXT); + expect((await reopened.getEntries(["root"], BACKGROUND_CONTEXT)).has("root")).toBe(true); + expect((await reopened.getValue(storedValues.sessionName, BACKGROUND_CONTEXT))?.value).toBe("updated"); + expect((await reopened.commit([], BACKGROUND_CONTEXT)).firstSeq).toBe(10); + await reopened.close(BACKGROUND_CONTEXT); }); it("removes the temporary file when publication fails", async () => { @@ -108,6 +119,7 @@ describe("JsonlStorage snapshot creation", () => { { fileSystem, path: "blocked.jsonl", now: () => NOW }, header("blocked"), preparedSnapshot(), + BACKGROUND_CONTEXT, ), ).rejects.toThrow("Failed to publish JSONL storage"); expect(getOrThrow(await fileSystem.exists("blocked.jsonl.tmp"))).toBe(false); @@ -120,46 +132,52 @@ describe("JsonlStorage persistence", () => { const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() }); const options = { fileSystem, path: "list-delete.jsonl", now: () => NOW }; const events = storedValues.list("test.events"); - const storage = await JsonlStorage.create(options, header("list-delete")); - await storage.commit([storedValues.appendList(events, "first"), storedValues.appendList(events, "second")]); - await storage.commit([storedValues.deleteList(events)]); - await storage.close(); + const storage = await JsonlStorage.create(options, header("list-delete"), BACKGROUND_CONTEXT); + await storage.commit( + [storedValues.appendList(events, "first"), storedValues.appendList(events, "second")], + BACKGROUND_CONTEXT, + ); + await storage.commit([storedValues.deleteList(events)], BACKGROUND_CONTEXT); + await storage.close(BACKGROUND_CONTEXT); - const reopened = await JsonlStorage.open(options); - expect(await reopened.readList(events)).toEqual([]); - const recreated = await reopened.commit([storedValues.appendList(events, "after")]); + const reopened = await JsonlStorage.open(options, BACKGROUND_CONTEXT); + expect(await reopened.readList(events, undefined, BACKGROUND_CONTEXT)).toEqual([]); + const recreated = await reopened.commit([storedValues.appendList(events, "after")], BACKGROUND_CONTEXT); expect(recreated.firstSeq).toBe(4); - expect(await reopened.readList(events)).toEqual([{ seq: 4, value: "after" }]); - await reopened.close(); + expect(await reopened.readList(events, undefined, BACKGROUND_CONTEXT)).toEqual([{ seq: 4, value: "after" }]); + await reopened.close(BACKGROUND_CONTEXT); }); it("writes one line per transaction and replays stamped state", async () => { const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() }); const options = { fileSystem, path: "session.jsonl", now: () => NOW }; - const storage = await JsonlStorage.create(options, header("round-trip")); - const committed = await storage.commit([ - sessionWrites.insertEntry({ - id: "root", - parentId: null, - type: "message", - message: { role: "user", content: "hello", timestamp: 1 }, - }), - storedValues.setValue(storedValues.laneLeaf("main"), "root"), - sessionWrites.insertUsage({ - id: "usage", - entryId: "root", - adjustment: false, - 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([storedValues.setValue(storedValues.sessionName, "name")]); + const storage = await JsonlStorage.create(options, header("round-trip"), BACKGROUND_CONTEXT); + const committed = await storage.commit( + [ + sessionWrites.insertEntry({ + id: "root", + parentId: null, + type: "message", + message: { role: "user", content: "hello", timestamp: 1 }, + }), + storedValues.setValue(storedValues.laneLeaf("main"), "root"), + sessionWrites.insertUsage({ + id: "usage", + entryId: "root", + adjustment: false, + usage: { + input: 1, + output: 2, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 3, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }), + ], + BACKGROUND_CONTEXT, + ); + await storage.commit([storedValues.setValue(storedValues.sessionName, "name")], BACKGROUND_CONTEXT); const lines = getOrThrow(await fileSystem.readTextFile("session.jsonl")) .trimEnd() @@ -167,10 +185,10 @@ describe("JsonlStorage persistence", () => { expect(JSON.parse(lines[0]!)).toEqual(header("round-trip")); expect(JSON.parse(lines[1]!)).toHaveLength(3); expect(Array.isArray(JSON.parse(lines[2]!))).toBe(false); - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); - const reopened = await JsonlStorage.open(options); - expect((await reopened.getEntries(["root"])).get("root")).toEqual({ + const reopened = await JsonlStorage.open(options, BACKGROUND_CONTEXT); + expect((await reopened.getEntries(["root"], BACKGROUND_CONTEXT)).get("root")).toEqual({ id: "root", parentId: null, type: "message", @@ -178,15 +196,15 @@ describe("JsonlStorage persistence", () => { seq: committed.seqs[0], timestamp: committed.timestamp, }); - expect(await reopened.getValue(storedValues.laneLeaf("main"))).toEqual({ + expect(await reopened.getValue(storedValues.laneLeaf("main"), BACKGROUND_CONTEXT)).toEqual({ address: storedValues.laneLeaf("main"), value: "root", seq: committed.seqs[1], }); - expect((await reopened.scanUsage({ order: "asc" })).map(({ id, seq }) => ({ id, seq }))).toEqual([ - { id: "usage", seq: committed.seqs[2] }, - ]); - expect(await reopened.getStats()).toEqual({ + 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({ messageCount: 1, usage: { input: 1, @@ -197,9 +215,9 @@ describe("JsonlStorage persistence", () => { cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, }); - const next = await reopened.commit([]); + const next = await reopened.commit([], BACKGROUND_CONTEXT); expect(next.firstSeq).toBe(5); - await reopened.close(); + await reopened.close(BACKGROUND_CONTEXT); }); }); @@ -207,17 +225,23 @@ describe("JsonlStorage snapshots", () => { it("captures one serialized boundary between commits", async () => { const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() }); const options = { fileSystem, path: "session.jsonl", now: () => NOW }; - const storage = await JsonlStorage.create(options, header("snapshot")); + const storage = await JsonlStorage.create(options, header("snapshot"), BACKGROUND_CONTEXT); - const firstCommit = storage.commit([ - sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "root" }), - storedValues.setValue(storedValues.laneLeaf("main"), "root"), - ]); - const snapshot = storage.snapshot(); - const secondCommit = storage.commit([ - sessionWrites.insertEntry({ id: "child", parentId: "root", type: "custom", customType: "child" }), - storedValues.setValue(storedValues.laneLeaf("main"), "child"), - ]); + const firstCommit = storage.commit( + [ + sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "root" }), + storedValues.setValue(storedValues.laneLeaf("main"), "root"), + ], + BACKGROUND_CONTEXT, + ); + const snapshot = storage.snapshot(BACKGROUND_CONTEXT); + const secondCommit = storage.commit( + [ + sessionWrites.insertEntry({ id: "child", parentId: "root", type: "custom", customType: "child" }), + storedValues.setValue(storedValues.laneLeaf("main"), "child"), + ], + BACKGROUND_CONTEXT, + ); await firstCommit; const captured = await snapshot; @@ -229,9 +253,9 @@ describe("JsonlStorage snapshots", () => { ({ address }) => address.namespace === storedValues.laneLeaf("main").namespace && address.key === "main", )?.value, ).toBe("root"); - expect((await storage.getEntries(["child"])).has("child")).toBe(true); - expect((await storage.getValue(storedValues.laneLeaf("main")))?.value).toBe("child"); - await storage.close(); + expect((await storage.getEntries(["child"], BACKGROUND_CONTEXT)).has("child")).toBe(true); + expect((await storage.getValue(storedValues.laneLeaf("main"), BACKGROUND_CONTEXT))?.value).toBe("child"); + await storage.close(BACKGROUND_CONTEXT); }); }); @@ -248,9 +272,9 @@ describe("JsonlStorage torn tail", () => { async function seed() { const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() }); const options = { fileSystem, path: "session.jsonl" as const, now: () => NOW }; - const storage = await JsonlStorage.create(options, header("torn")); - await storage.commit([entryWrite("kept")]); - await storage.close(); + const storage = await JsonlStorage.create(options, header("torn"), BACKGROUND_CONTEXT); + await storage.commit([entryWrite("kept")], BACKGROUND_CONTEXT); + await storage.close(BACKGROUND_CONTEXT); const prefix = getOrThrow(await fileSystem.readTextFile("session.jsonl")); return { fileSystem, options, prefix }; } @@ -270,16 +294,16 @@ describe("JsonlStorage torn tail", () => { }), ); - const reopened = await JsonlStorage.open(options); - expect((await reopened.getEntries(["kept", "torn"])).has("torn")).toBe(false); - expect((await reopened.getEntries(["kept"])).get("kept")?.id).toBe("kept"); + const reopened = await JsonlStorage.open(options, BACKGROUND_CONTEXT); + expect((await reopened.getEntries(["kept", "torn"], BACKGROUND_CONTEXT)).has("torn")).toBe(false); + expect((await reopened.getEntries(["kept"], BACKGROUND_CONTEXT)).get("kept")?.id).toBe("kept"); expect(getOrThrow(await fileSystem.readTextFile("session.jsonl"))).toBe(prefix); expect(getOrThrow(await fileSystem.exists("session.jsonl.tmp"))).toBe(false); - const next = await reopened.commit([entryWrite("after")]); + const next = await reopened.commit([entryWrite("after")], BACKGROUND_CONTEXT); expect(next.firstSeq).toBe(2); - expect((await reopened.getEntries(["after"])).get("after")?.seq).toBe(2); - await reopened.close(); + expect((await reopened.getEntries(["after"], BACKGROUND_CONTEXT)).get("after")?.seq).toBe(2); + await reopened.close(BACKGROUND_CONTEXT); }); it("discards a torn array line wholly, including list elements", async () => { @@ -302,12 +326,12 @@ describe("JsonlStorage torn tail", () => { ]), ); - const reopened = await JsonlStorage.open(options); - expect((await reopened.getEntries(["torn-a"])).has("torn-a")).toBe(false); - expect(await reopened.getValue(storedValues.sessionName)).toBeUndefined(); - expect(await reopened.readList(events)).toEqual([]); + const reopened = await JsonlStorage.open(options, BACKGROUND_CONTEXT); + expect((await reopened.getEntries(["torn-a"], BACKGROUND_CONTEXT)).has("torn-a")).toBe(false); + expect(await reopened.getValue(storedValues.sessionName, BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await reopened.readList(events, undefined, BACKGROUND_CONTEXT)).toEqual([]); expect(getOrThrow(await fileSystem.readTextFile("session.jsonl"))).toBe(prefix); - await reopened.close(); + await reopened.close(BACKGROUND_CONTEXT); }); it("rejects a malformed interior line without rewriting", async () => { @@ -315,7 +339,7 @@ describe("JsonlStorage torn tail", () => { const corrupted = `${prefix}not-json\n${JSON.stringify(storedValues.setValue(storedValues.sessionName, "after"))}\n`; await fileSystem.writeFile("session.jsonl", corrupted); - await expect(JsonlStorage.open(options)).rejects.toThrow(/line 3/); + await expect(JsonlStorage.open(options, BACKGROUND_CONTEXT)).rejects.toThrow(/line 3/); expect(getOrThrow(await fileSystem.readTextFile("session.jsonl"))).toBe(corrupted); expect(getOrThrow(await fileSystem.exists("session.jsonl.tmp"))).toBe(false); }); @@ -333,7 +357,7 @@ describe("JsonlStorage torn tail", () => { })}\n`; await fileSystem.writeFile("session.jsonl", corrupted); - await expect(JsonlStorage.open(options)).rejects.toThrow(/line 3/); + await expect(JsonlStorage.open(options, BACKGROUND_CONTEXT)).rejects.toThrow(/line 3/); expect(getOrThrow(await fileSystem.readTextFile("session.jsonl"))).toBe(corrupted); }); @@ -342,7 +366,7 @@ describe("JsonlStorage torn tail", () => { const corrupted = `${prefix}not-json\n`; await fileSystem.writeFile("session.jsonl", corrupted); - await expect(JsonlStorage.open(options)).rejects.toThrow(/line 3/); + await expect(JsonlStorage.open(options, BACKGROUND_CONTEXT)).rejects.toThrow(/line 3/); expect(getOrThrow(await fileSystem.readTextFile("session.jsonl"))).toBe(corrupted); }); @@ -351,7 +375,7 @@ describe("JsonlStorage torn tail", () => { const corrupted = `${prefix}${JSON.stringify({ kind: "nope", seq: 2 })}\n`; await fileSystem.writeFile("session.jsonl", corrupted); - await expect(JsonlStorage.open(options)).rejects.toThrow(/line 3/); + await expect(JsonlStorage.open(options, BACKGROUND_CONTEXT)).rejects.toThrow(/line 3/); expect(getOrThrow(await fileSystem.readTextFile("session.jsonl"))).toBe(corrupted); }); @@ -360,6 +384,6 @@ describe("JsonlStorage torn tail", () => { const options = { fileSystem, path: "session.jsonl", now: () => NOW }; await fileSystem.writeFile("session.jsonl", JSON.stringify(header("torn")).slice(0, -4)); - await expect(JsonlStorage.open(options)).rejects.toThrow(/missing header/); + await expect(JsonlStorage.open(options, BACKGROUND_CONTEXT)).rejects.toThrow(/missing header/); }); }); diff --git a/packages/agent/test/harness/memory-conformance.test.ts b/packages/agent/test/harness/memory-conformance.test.ts index a27d04dd0..27d9aa6ab 100644 --- a/packages/agent/test/harness/memory-conformance.test.ts +++ b/packages/agent/test/harness/memory-conformance.test.ts @@ -1,4 +1,5 @@ import { describe, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import { MemorySessionRepo } from "../../src/harness/session/index.ts"; import { MemoryStorage } from "../../src/harness/session/memory.ts"; import { @@ -28,7 +29,7 @@ registerConformance( const storage = new MemoryStorage({ now: () => NOW }); return Promise.resolve({ storage, - [Symbol.asyncDispose]: () => storage.close(), + [Symbol.asyncDispose]: () => storage.close(BACKGROUND_CONTEXT), }); }), ); @@ -41,6 +42,6 @@ registerConformance( memoryRepo = new MemorySessionRepo({ now: () => NOW }); return Promise.resolve(memoryRepo); }, - () => memoryRepo.close(), + () => memoryRepo.close(BACKGROUND_CONTEXT), ), ); diff --git a/packages/agent/test/harness/memory-session-repo.test.ts b/packages/agent/test/harness/memory-session-repo.test.ts index 2ac69a369..bb38cccb8 100644 --- a/packages/agent/test/harness/memory-session-repo.test.ts +++ b/packages/agent/test/harness/memory-session-repo.test.ts @@ -1,4 +1,5 @@ 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 * as storedValues from "../../src/harness/session/values.ts"; @@ -11,51 +12,61 @@ function uuidTimestamp(id: string): number { describe("MemorySessionRepo metadata", () => { it("uses its injected clock for generated session identity and metadata", async () => { const repo = new MemorySessionRepo({ now: () => NOW }); - const session = await repo.create({}); + const session = await repo.create({}, BACKGROUND_CONTEXT); expect(session.metadata.createdAt).toBe(NOW); expect(uuidTimestamp(session.metadata.id)).toBe(NOW); - await Promise.all([session.close(), repo.close()]); + await Promise.all([session.close(BACKGROUND_CONTEXT), repo.close(BACKGROUND_CONTEXT)]); }); it("returns a fresh facade after close while retaining one session and storage", async () => { const repo = new MemorySessionRepo({ now: () => NOW }); - const first = await repo.create({ id: "session" }); + const first = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); const firstView = first.view("main"); - const admittedWrite = first.setName("preserved"); + const admittedWrite = first.setName("preserved", BACKGROUND_CONTEXT); - await expect(repo.open(first.metadata)).rejects.toThrow("already open"); - await Promise.all([admittedWrite, first.close()]); - await expect(first.getName()).rejects.toThrow("Session is closed"); - await expect(firstView.getName()).rejects.toThrow("Session is closed"); + await expect(repo.open(first.metadata, BACKGROUND_CONTEXT)).rejects.toThrow("already open"); + await Promise.all([admittedWrite, first.close(BACKGROUND_CONTEXT)]); + await expect(first.getName(BACKGROUND_CONTEXT)).rejects.toThrow("Session is closed"); + await expect(firstView.getName(BACKGROUND_CONTEXT)).rejects.toThrow("Session is closed"); - const second = await repo.open(first.metadata); + const second = await repo.open(first.metadata, BACKGROUND_CONTEXT); expect(second).not.toBe(first); - expect(await second.getName()).toBe("preserved"); - await second.close(); - await repo.close(); + expect(await second.getName(BACKGROUND_CONTEXT)).toBe("preserved"); + await second.close(BACKGROUND_CONTEXT); + await repo.close(BACKGROUND_CONTEXT); }); it("captures fork options before waiting for its snapshot boundary", async () => { const repo = new MemorySessionRepo({ now: () => NOW }); - const source = await repo.create({ id: "source" }); + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); const rootId = "00000000-0000-7000-8000-000000000001"; const childId = "00000000-0000-7000-8000-000000000002"; - const commit = source.mutate("main", (mutator) => - mutator.commit([ - sessionWrites.insertEntry({ id: rootId, parentId: null, type: "custom", customType: "root" }), - sessionWrites.insertEntry({ id: childId, parentId: rootId, type: "custom", customType: "child" }), - storedValues.setValue(storedValues.laneLeaf("main"), childId), - ]), + const commit = source.mutate( + "main", + (mutator) => + mutator.commit( + [ + sessionWrites.insertEntry({ id: rootId, parentId: null, type: "custom", customType: "root" }), + sessionWrites.insertEntry({ id: childId, parentId: rootId, type: "custom", customType: "child" }), + storedValues.setValue(storedValues.laneLeaf("main"), childId), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); const options = { id: "fork", entryId: childId, position: "before" as "before" | "at" }; - const fork = repo.fork(source.metadata, options); + const fork = repo.fork(source.metadata, options, BACKGROUND_CONTEXT); options.entryId = rootId; options.position = "at"; await commit; const forked = await fork; - expect(await forked.getLeafId()).toBe(rootId); - await Promise.all([source.close(), forked.close(), repo.close()]); + expect(await forked.getLeafId(BACKGROUND_CONTEXT)).toBe(rootId); + await Promise.all([ + source.close(BACKGROUND_CONTEXT), + forked.close(BACKGROUND_CONTEXT), + repo.close(BACKGROUND_CONTEXT), + ]); }); }); diff --git a/packages/agent/test/harness/memory-storage.test.ts b/packages/agent/test/harness/memory-storage.test.ts index d58b54810..e7fab23fe 100644 --- a/packages/agent/test/harness/memory-storage.test.ts +++ b/packages/agent/test/harness/memory-storage.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "vitest"; +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 * as storedValues from "../../src/harness/session/values.ts"; @@ -10,20 +11,26 @@ describe("MemoryStorage", () => { let timestamp = NOW; const storage = new MemoryStorage({ now: () => timestamp++ }); - const first = await storage.commit([ - sessionWrites.insertEntry({ - id: "first", - parentId: null, - type: "custom", - customType: "note", - }), - storedValues.setValue(storedValues.sessionName, "first"), - ]); - const second = await storage.commit([storedValues.setValue(storedValues.sessionName, "second")]); + const first = await storage.commit( + [ + sessionWrites.insertEntry({ + id: "first", + parentId: null, + type: "custom", + customType: "note", + }), + storedValues.setValue(storedValues.sessionName, "first"), + ], + BACKGROUND_CONTEXT, + ); + const second = await storage.commit( + [storedValues.setValue(storedValues.sessionName, "second")], + BACKGROUND_CONTEXT, + ); expect(first.timestamp).toBe(NOW); expect(second.timestamp).toBe(NOW + 1); - expect((await storage.getEntries(["first"])).get("first")?.timestamp).toBe(first.timestamp); - await storage.close(); + expect((await storage.getEntries(["first"], BACKGROUND_CONTEXT)).get("first")?.timestamp).toBe(first.timestamp); + await storage.close(BACKGROUND_CONTEXT); }); }); diff --git a/packages/agent/test/harness/runtime2/harness.test.ts b/packages/agent/test/harness/runtime2/harness.test.ts index d2f82c2b3..cd99864ef 100644 --- a/packages/agent/test/harness/runtime2/harness.test.ts +++ b/packages/agent/test/harness/runtime2/harness.test.ts @@ -3,6 +3,7 @@ import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; import { AgentHarness, HarnessClosed, HarnessFault } from "../../../src/harness/agent-harness.ts"; import { DEFAULT_COMPACTION_SETTINGS } from "../../../src/harness/compaction/compaction.ts"; +import { BACKGROUND_CONTEXT } from "../../../src/harness/context.ts"; import type { Result } from "../../../src/harness/result.ts"; import { createAgentHarness, Harness } from "../../../src/harness/runtime2/harness.ts"; import { Lane } from "../../../src/harness/runtime2/lane.ts"; @@ -44,12 +45,15 @@ function tool(name: string): AgentHarnessTool async function createSession(): Promise { const repo = new MemorySessionRepo(); repos.push(repo); - return repo.create({}); + return repo.create({}, BACKGROUND_CONTEXT); } async function configureMain(session: Session, configuration: LaneConfiguration = configuredMain): Promise { - await session.mutate("main", (mutator) => - mutator.commit([storedValues.setValue(storedValues.laneConfig("main"), configuration)]), + await session.mutate( + "main", + (mutator) => + mutator.commit([storedValues.setValue(storedValues.laneConfig("main"), configuration)], BACKGROUND_CONTEXT), + BACKGROUND_CONTEXT, ); } @@ -59,11 +63,17 @@ async function createStorageSession(storage: MemoryStorage): Promise { storage, ); standaloneSessions.push(session); - await session.mutate("main", (mutator) => - mutator.commit([ - storedValues.setValue(storedValues.laneLeaf("main"), null), - storedValues.setValue(storedValues.laneState("main"), { currentOperationId: null, pendingNextRun: [] }), - ]), + await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + storedValues.setValue(storedValues.laneLeaf("main"), null), + storedValues.setValue(storedValues.laneState("main"), { currentOperationId: null, pendingNextRun: [] }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); return session; } @@ -90,15 +100,21 @@ function operationMeta(session: Session, intent: OperationMeta["intent"], starte } async function installOperation(session: Session, meta: OperationMeta, state: OperationState): Promise { - await session.mutate(meta.lane, (mutator) => - mutator.commit([ - storedValues.setValue(storedValues.operationMeta(meta.operationId), meta), - storedValues.setValue(storedValues.operationState(meta.operationId), state), - storedValues.setValue(storedValues.laneState(meta.lane), { - currentOperationId: meta.operationId, - pendingNextRun: [], - }), - ]), + await session.mutate( + meta.lane, + (mutator) => + mutator.commit( + [ + storedValues.setValue(storedValues.operationMeta(meta.operationId), meta), + storedValues.setValue(storedValues.operationState(meta.operationId), state), + storedValues.setValue(storedValues.laneState(meta.lane), { + currentOperationId: meta.operationId, + pendingNextRun: [], + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); } @@ -123,15 +139,15 @@ function runState(triggerEntryId: string): RunState { } afterEach(async () => { - for (const repo of repos.splice(0)) await repo.close(); - for (const session of standaloneSessions.splice(0)) await session.close(); + for (const repo of repos.splice(0)) await repo.close(BACKGROUND_CONTEXT); + for (const session of standaloneSessions.splice(0)) await session.close(BACKGROUND_CONTEXT); }); describe("runtime2 AgentHarness", () => { it("is selected by the public AgentHarness constructor", async () => { const session = await createSession(); - const { harness } = await AgentHarness.create(modelOptions(session)); + const { harness } = await AgentHarness.create(modelOptions(session), BACKGROUND_CONTEXT); expect(harness).toBeInstanceOf(Harness); }); @@ -141,7 +157,7 @@ describe("runtime2 AgentHarness", () => { const duplicate = tool("duplicate"); await expect( - createAgentHarness({ ...modelOptions(session), tools: [duplicate, duplicate] }), + createAgentHarness({ ...modelOptions(session), tools: [duplicate, duplicate] }, BACKGROUND_CONTEXT), ).rejects.toBeInstanceOf(TypeError); }); @@ -152,31 +168,34 @@ describe("runtime2 AgentHarness", () => { const initialStreamOptions = { timeoutMs: 10 }; const initialRetry = { enabled: true, maxRetries: 2, baseDelayMs: 20 }; const initialCompaction = { enabled: false, reserveTokens: 100, keepRecentTokens: 200 }; - const { harness } = await createAgentHarness({ - ...modelOptions(session), - tools: initialTools, - resources: initialResources, - streamOptions: initialStreamOptions, - retry: initialRetry, - compaction: initialCompaction, - steeringMode: "one-at-a-time", - followUpMode: "one-at-a-time", - toolExecution: "sequential", - }); + const { harness } = await createAgentHarness( + { + ...modelOptions(session), + tools: initialTools, + resources: initialResources, + streamOptions: initialStreamOptions, + retry: initialRetry, + compaction: initialCompaction, + steeringMode: "one-at-a-time", + followUpMode: "one-at-a-time", + toolExecution: "sequential", + }, + BACKGROUND_CONTEXT, + ); - expect(await harness.getTools()).toBe(initialTools); - expect(await harness.getResources()).toBe(initialResources); - expect(await harness.getStreamOptions()).toBe(initialStreamOptions); - expect(await harness.getRetryPolicy()).toBe(initialRetry); - expect(await harness.getCompactionSettings()).toBe(initialCompaction); - expect(await harness.getSteeringMode()).toBe("one-at-a-time"); - expect(await harness.getFollowUpMode()).toBe("one-at-a-time"); + expect(await harness.getTools(BACKGROUND_CONTEXT)).toBe(initialTools); + expect(await harness.getResources(BACKGROUND_CONTEXT)).toBe(initialResources); + expect(await harness.getStreamOptions(BACKGROUND_CONTEXT)).toBe(initialStreamOptions); + expect(await harness.getRetryPolicy(BACKGROUND_CONTEXT)).toBe(initialRetry); + expect(await harness.getCompactionSettings(BACKGROUND_CONTEXT)).toBe(initialCompaction); + expect(await harness.getSteeringMode(BACKGROUND_CONTEXT)).toBe("one-at-a-time"); + expect(await harness.getFollowUpMode(BACKGROUND_CONTEXT)).toBe("one-at-a-time"); const events: string[] = []; let retryAtEvent: typeof initialRetry | undefined; harness.events.on("config_update", async (event) => { events.push(event.property); - if (event.property === "retryPolicy") retryAtEvent = await harness.getRetryPolicy(); + if (event.property === "retryPolicy") retryAtEvent = await harness.getRetryPolicy(BACKGROUND_CONTEXT); }); const nextTools = [tool("next")]; const nextResources = { promptTemplates: [] }; @@ -184,21 +203,21 @@ describe("runtime2 AgentHarness", () => { const nextRetry = { enabled: false, maxRetries: 0, baseDelayMs: 0 }; const nextCompaction = { enabled: true, reserveTokens: 300, keepRecentTokens: 400 }; - await harness.setTools(nextTools); - await harness.setResources(nextResources); - await harness.setStreamOptions(nextStreamOptions); - await harness.setRetryPolicy(nextRetry); - await harness.setCompactionSettings(nextCompaction); - await harness.setSteeringMode("all"); - await harness.setFollowUpMode("all"); + await harness.setTools(nextTools, BACKGROUND_CONTEXT); + await harness.setResources(nextResources, BACKGROUND_CONTEXT); + await harness.setStreamOptions(nextStreamOptions, BACKGROUND_CONTEXT); + await harness.setRetryPolicy(nextRetry, BACKGROUND_CONTEXT); + await harness.setCompactionSettings(nextCompaction, BACKGROUND_CONTEXT); + await harness.setSteeringMode("all", BACKGROUND_CONTEXT); + await harness.setFollowUpMode("all", BACKGROUND_CONTEXT); - expect(await harness.getTools()).toBe(nextTools); - expect(await harness.getResources()).toBe(nextResources); - expect(await harness.getStreamOptions()).toBe(nextStreamOptions); - expect(await harness.getRetryPolicy()).toBe(nextRetry); - expect(await harness.getCompactionSettings()).toBe(nextCompaction); - expect(await harness.getSteeringMode()).toBe("all"); - expect(await harness.getFollowUpMode()).toBe("all"); + expect(await harness.getTools(BACKGROUND_CONTEXT)).toBe(nextTools); + expect(await harness.getResources(BACKGROUND_CONTEXT)).toBe(nextResources); + expect(await harness.getStreamOptions(BACKGROUND_CONTEXT)).toBe(nextStreamOptions); + expect(await harness.getRetryPolicy(BACKGROUND_CONTEXT)).toBe(nextRetry); + expect(await harness.getCompactionSettings(BACKGROUND_CONTEXT)).toBe(nextCompaction); + expect(await harness.getSteeringMode(BACKGROUND_CONTEXT)).toBe("all"); + expect(await harness.getFollowUpMode(BACKGROUND_CONTEXT)).toBe("all"); expect(retryAtEvent).toBe(nextRetry); expect(events).toEqual([ "tools", @@ -212,25 +231,28 @@ describe("runtime2 AgentHarness", () => { const penultimateResources = { skills: [] }; const finalResources = { promptTemplates: [] }; - await Promise.all([harness.setResources(penultimateResources), harness.setResources(finalResources)]); - expect(await harness.getResources()).toBe(finalResources); + await Promise.all([ + harness.setResources(penultimateResources, BACKGROUND_CONTEXT), + harness.setResources(finalResources, BACKGROUND_CONTEXT), + ]); + expect(await harness.getResources(BACKGROUND_CONTEXT)).toBe(finalResources); expect(events.slice(-2)).toEqual(["resources", "resources"]); }); it("emits lane config changes after committed memory publication", async () => { const session = await createSession(); const options = modelOptions(session); - const { harness } = await createAgentHarness(options); + const { harness } = await createAgentHarness(options, BACKGROUND_CONTEXT); const events: unknown[] = []; let thinkingAtEvent: string | undefined; harness.events.on("config_update", async (event) => { events.push(event); - if (event.property === "thinkingLevel") thinkingAtEvent = await harness.getThinkingLevel(); + if (event.property === "thinkingLevel") thinkingAtEvent = await harness.getThinkingLevel(BACKGROUND_CONTEXT); }); - await harness.setModel(options.model); - await harness.setThinkingLevel("high"); - await harness.setActiveTools(["read"]); + await harness.setModel(options.model, BACKGROUND_CONTEXT); + await harness.setThinkingLevel("high", BACKGROUND_CONTEXT); + await harness.setActiveTools(["read"], BACKGROUND_CONTEXT); const model = { provider: options.model.provider, modelId: options.model.id }; expect(events).toEqual([ @@ -244,61 +266,81 @@ describe("runtime2 AgentHarness", () => { it("rejects invalid config without publication or fault", async () => { const session = await createSession(); const initialTools = [tool("initial")]; - const { harness } = await createAgentHarness({ ...modelOptions(session), tools: initialTools }); + const { harness } = await createAgentHarness( + { ...modelOptions(session), tools: initialTools }, + BACKGROUND_CONTEXT, + ); const events: string[] = []; harness.events.on("config_update", (event) => { events.push(event.property); }); const duplicate = tool("duplicate"); - expect(() => harness.setTools([duplicate, duplicate])).toThrow(TypeError); - expect(() => harness.setRetryPolicy({ enabled: true, maxRetries: -1, baseDelayMs: 0 })).toThrow(RangeError); - expect(() => harness.setCompactionSettings({ enabled: true, reserveTokens: -1, keepRecentTokens: 0 })).toThrow( - RangeError, - ); + expect(() => harness.setTools([duplicate, duplicate], BACKGROUND_CONTEXT)).toThrow(TypeError); + expect(() => + harness.setRetryPolicy({ enabled: true, maxRetries: -1, baseDelayMs: 0 }, BACKGROUND_CONTEXT), + ).toThrow(RangeError); + expect(() => + harness.setCompactionSettings({ enabled: true, reserveTokens: -1, keepRecentTokens: 0 }, BACKGROUND_CONTEXT), + ).toThrow(RangeError); - expect(await harness.getTools()).toBe(initialTools); - expect(await harness.getRetryPolicy()).toEqual({ enabled: true, maxRetries: 3, baseDelayMs: 1_000 }); - expect(await harness.getCompactionSettings()).toBe(DEFAULT_COMPACTION_SETTINGS); + expect(await harness.getTools(BACKGROUND_CONTEXT)).toBe(initialTools); + expect(await harness.getRetryPolicy(BACKGROUND_CONTEXT)).toEqual({ + enabled: true, + maxRetries: 3, + baseDelayMs: 1_000, + }); + expect(await harness.getCompactionSettings(BACKGROUND_CONTEXT)).toBe(DEFAULT_COMPACTION_SETTINGS); expect(events).toEqual([]); - expect(await harness.lanes()).toHaveLength(1); + expect(await harness.lanes(BACKGROUND_CONTEXT)).toHaveLength(1); }); it("validates initial config before writing durable state", async () => { const invalidRetrySession = await createSession(); await expect( - createAgentHarness({ - ...modelOptions(invalidRetrySession), - retry: { enabled: true, maxRetries: -1, baseDelayMs: 0 }, - }), + createAgentHarness( + { + ...modelOptions(invalidRetrySession), + retry: { enabled: true, maxRetries: -1, baseDelayMs: 0 }, + }, + BACKGROUND_CONTEXT, + ), ).rejects.toBeInstanceOf(RangeError); - expect(await invalidRetrySession.getValue(storedValues.laneConfig("main"))).toBeUndefined(); + expect(await invalidRetrySession.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT)).toBeUndefined(); const invalidCompactionSession = await createSession(); await expect( - createAgentHarness({ - ...modelOptions(invalidCompactionSession), - compaction: { enabled: true, reserveTokens: -1, keepRecentTokens: 0 }, - }), + createAgentHarness( + { + ...modelOptions(invalidCompactionSession), + compaction: { enabled: true, reserveTokens: -1, keepRecentTokens: 0 }, + }, + BACKGROUND_CONTEXT, + ), ).rejects.toBeInstanceOf(RangeError); - expect(await invalidCompactionSession.getValue(storedValues.laneConfig("main"))).toBeUndefined(); + expect( + await invalidCompactionSession.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT), + ).toBeUndefined(); }); it("seeds main and returns the concrete harness as its main lane", async () => { const session = await createSession(); const options = modelOptions(session); - const { harness, suspended } = await createAgentHarness({ - ...options, - thinkingLevel: "high", - activeToolNames: ["read"], - }); + const { harness, suspended } = await createAgentHarness( + { + ...options, + thinkingLevel: "high", + activeToolNames: ["read"], + }, + BACKGROUND_CONTEXT, + ); expect(harness).toBeInstanceOf(Harness); - expect(await harness.lane("main")).toBe(harness); - expect(await harness.getLeafId()).toBeNull(); + expect(await harness.lane("main", BACKGROUND_CONTEXT)).toBe(harness); + expect(await harness.getLeafId(BACKGROUND_CONTEXT)).toBeNull(); expect(suspended).toEqual([]); - expect((await session.getValue(storedValues.laneConfig("main")))?.value).toEqual({ + expect((await session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT))?.value).toEqual({ model: { provider: options.model.provider, modelId: options.model.id }, thinkingLevel: "high", activeToolNames: ["read"], @@ -307,15 +349,18 @@ describe("runtime2 AgentHarness", () => { it("appends idle entries through owned lane state", async () => { const session = await createSession(); - const { harness } = await createAgentHarness(modelOptions(session)); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); - const entryId = await harness.sessionTree.appendMessage({ role: "user", content: "hello", timestamp: 1 }); + const entryId = await harness.sessionTree.appendMessage( + { role: "user", content: "hello", timestamp: 1 }, + BACKGROUND_CONTEXT, + ); - expect(await harness.getLeafId()).toBe(entryId); - expect((await session.getValue(storedValues.laneLeaf("main")))?.value).toBe(entryId); - expect(await harness.sessionTree.findEntriesOnBranch({ order: "oldestFirst" })).toMatchObject([ - { id: entryId, parentId: null, message: { role: "user", content: "hello" } }, - ]); + expect(await harness.getLeafId(BACKGROUND_CONTEXT)).toBe(entryId); + expect((await session.getValue(storedValues.laneLeaf("main"), BACKGROUND_CONTEXT))?.value).toBe(entryId); + expect(await harness.sessionTree.findEntriesOnBranch({ order: "oldestFirst" }, BACKGROUND_CONTEXT)).toMatchObject( + [{ id: entryId, parentId: null, message: { role: "user", content: "hello" } }], + ); }); it("restores every configured lane without replacing its configuration", async () => { @@ -325,25 +370,31 @@ describe("runtime2 AgentHarness", () => { thinkingLevel: "medium", activeToolNames: [], }; - await session.mutate("main", (mutator) => - mutator.commit([storedValues.setValue(storedValues.laneConfig("main"), configuredMain)]), + await session.mutate( + "main", + (mutator) => + mutator.commit( + [storedValues.setValue(storedValues.laneConfig("main"), configuredMain)], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - await session.createLane("worker", null, workerConfiguration); + await session.createLane("worker", null, workerConfiguration, BACKGROUND_CONTEXT); const before = await Promise.all([ - session.getValue(storedValues.laneConfig("main")), - session.getValue(storedValues.laneConfig("worker")), + session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT), + session.getValue(storedValues.laneConfig("worker"), BACKGROUND_CONTEXT), ]); - const { harness } = await createAgentHarness(modelOptions(session)); - const worker = await harness.lane("worker"); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); + const worker = await harness.lane("worker", BACKGROUND_CONTEXT); expect(worker).toBeInstanceOf(Lane); - expect(await worker?.getLeafId()).toBeNull(); - expect((await harness.lanes()).map((lane) => lane.name).sort()).toEqual(["main", "worker"]); + expect(await worker?.getLeafId(BACKGROUND_CONTEXT)).toBeNull(); + expect((await harness.lanes(BACKGROUND_CONTEXT)).map((lane) => lane.name).sort()).toEqual(["main", "worker"]); expect( await Promise.all([ - session.getValue(storedValues.laneConfig("main")), - session.getValue(storedValues.laneConfig("worker")), + session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT), + session.getValue(storedValues.laneConfig("worker"), BACKGROUND_CONTEXT), ]), ).toEqual(before); }); @@ -351,33 +402,37 @@ describe("runtime2 AgentHarness", () => { it("creates a lane from the captured seed and publishes it after commit", async () => { const session = await createSession(); const options = modelOptions(session); - const { harness } = await createAgentHarness({ - ...options, - thinkingLevel: "high", - activeToolNames: ["read"], - }); - const anchor = await harness.sessionTree.appendCustomEntry("anchor"); - await harness.setThinkingLevel("low"); + const { harness } = await createAgentHarness( + { + ...options, + thinkingLevel: "high", + activeToolNames: ["read"], + }, + BACKGROUND_CONTEXT, + ); + const anchor = await harness.sessionTree.appendCustomEntry("anchor", undefined, BACKGROUND_CONTEXT); + await harness.setThinkingLevel("low", BACKGROUND_CONTEXT); let durableConfigurationAtEvent: LaneConfiguration | undefined; let laneVisibleAtEvent = false; harness.events.on("lane_created", async () => { - durableConfigurationAtEvent = (await session.getValue(storedValues.laneConfig("worker")))?.value; - laneVisibleAtEvent = (await harness.lane("worker")) !== undefined; + durableConfigurationAtEvent = (await session.getValue(storedValues.laneConfig("worker"), BACKGROUND_CONTEXT)) + ?.value; + laneVisibleAtEvent = (await harness.lane("worker", BACKGROUND_CONTEXT)) !== undefined; }); - const worker = unwrap(await harness.createLane("worker", anchor)); + const worker = unwrap(await harness.createLane("worker", anchor, BACKGROUND_CONTEXT)); - expect(await harness.lane("worker")).toBe(worker); - expect(await worker.getLeafId()).toBe(anchor); - expect(await worker.getThinkingLevel()).toBe("high"); - expect(await worker.getActiveTools()).toEqual(["read"]); + expect(await harness.lane("worker", BACKGROUND_CONTEXT)).toBe(worker); + expect(await worker.getLeafId(BACKGROUND_CONTEXT)).toBe(anchor); + expect(await worker.getThinkingLevel(BACKGROUND_CONTEXT)).toBe("high"); + expect(await worker.getActiveTools(BACKGROUND_CONTEXT)).toEqual(["read"]); expect(durableConfigurationAtEvent).toEqual({ model: { provider: options.model.provider, modelId: options.model.id }, thinkingLevel: "high", activeToolNames: ["read"], }); expect(laneVisibleAtEvent).toBe(true); - expect((await session.getValue(storedValues.laneState("worker")))?.value).toEqual({ + expect((await session.getValue(storedValues.laneState("worker"), BACKGROUND_CONTEXT))?.value).toEqual({ currentOperationId: null, pendingNextRun: [], }); @@ -385,34 +440,34 @@ describe("runtime2 AgentHarness", () => { it("maps expected lane creation failures without faulting", async () => { const session = await createSession(); - const { harness } = await createAgentHarness(modelOptions(session)); - unwrap(await harness.createLane("worker", null)); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); + unwrap(await harness.createLane("worker", null, BACKGROUND_CONTEXT)); - expect(await harness.createLane("worker", null)).toMatchObject({ + expect(await harness.createLane("worker", null, BACKGROUND_CONTEXT)).toMatchObject({ ok: false, error: { _tag: "LaneExists", lane: "worker" }, }); - expect(await harness.createLane("", null)).toMatchObject({ + expect(await harness.createLane("", null, BACKGROUND_CONTEXT)).toMatchObject({ ok: false, error: { _tag: "InvalidLane", lane: "", reason: "lane name must not be empty" }, }); - expect(await harness.createLane("missing", "unknown-entry")).toMatchObject({ + expect(await harness.createLane("missing", "unknown-entry", BACKGROUND_CONTEXT)).toMatchObject({ ok: false, error: { _tag: "UnknownTarget", targetId: "unknown-entry" }, }); - expect((await harness.lanes()).map((lane) => lane.name).sort()).toEqual(["main", "worker"]); + expect((await harness.lanes(BACKGROUND_CONTEXT)).map((lane) => lane.name).sort()).toEqual(["main", "worker"]); }); it("faults without publishing a lane when creation commit fails", async () => { const { session, storage } = await createFailingSession(); - const { harness } = await createAgentHarness(modelOptions(session)); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); if (!(harness instanceof Harness)) throw new Error("missing runtime2 harness"); const failure = new Error("create lane failed"); storage.failure = failure; let rejected: unknown; try { - await harness.createLane("worker", null); + await harness.createLane("worker", null, BACKGROUND_CONTEXT); } catch (error) { rejected = error; } @@ -420,30 +475,30 @@ describe("runtime2 AgentHarness", () => { expect(rejected).toBeInstanceOf(HarnessFault); if (!(rejected instanceof HarnessFault)) throw new Error("missing harness fault"); expect(rejected.cause).toBe(failure); - await expect(harness.createLane("later", null)).rejects.toBe(rejected); + await expect(harness.createLane("later", null, BACKGROUND_CONTEXT)).rejects.toBe(rejected); expect(harness.lanesByName.has("worker")).toBe(false); - expect(await session.getValue(storedValues.laneConfig("worker"))).toBeUndefined(); - expect(await session.getValue(storedValues.laneLeaf("worker"))).toBeUndefined(); - expect(await session.getValue(storedValues.laneState("worker"))).toBeUndefined(); + expect(await session.getValue(storedValues.laneConfig("worker"), BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await session.getValue(storedValues.laneLeaf("worker"), BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await session.getValue(storedValues.laneState("worker"), BACKGROUND_CONTEXT)).toBeUndefined(); }); it("returns Closed when lane creation starts after close", async () => { const session = await createSession(); - const { harness } = await createAgentHarness(modelOptions(session)); - await harness.close(); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); + await harness.close(BACKGROUND_CONTEXT); - expect(await harness.createLane("late", null)).toMatchObject({ + expect(await harness.createLane("late", null, BACKGROUND_CONTEXT)).toMatchObject({ ok: false, error: { _tag: "Closed" }, }); - await expect(harness.getTools()).rejects.toBeInstanceOf(HarnessClosed); - await expect(harness.setResources({})).rejects.toBeInstanceOf(HarnessClosed); + await expect(harness.getTools(BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(HarnessClosed); + await expect(harness.setResources({}, BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(HarnessClosed); }); it("publishes an admitted lane creation before close finishes", async () => { const storage = new ControlledMemoryStorage(); const session = await createStorageSession(storage); - const { harness } = await createAgentHarness(modelOptions(session)); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); if (!(harness instanceof Harness)) throw new Error("missing runtime2 harness"); const commitStarted = deferred(); const releaseCommit = deferred(); @@ -451,16 +506,16 @@ describe("runtime2 AgentHarness", () => { commitStarted.resolve(); await releaseCommit.promise; }; - const creating = harness.createLane("worker", null); + const creating = harness.createLane("worker", null, BACKGROUND_CONTEXT); await commitStarted.promise; - const closing = harness.close(); + const closing = harness.close(BACKGROUND_CONTEXT); releaseCommit.resolve(); const worker = unwrap(await creating); await closing; expect(harness.lanesByName.get("worker")).toBe(worker); - await expect(worker.getLeafId()).rejects.toBeInstanceOf(HarnessClosed); + await expect(worker.getLeafId(BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(HarnessClosed); }); it("classifies assistant-ready suspensions from captured identities", async () => { @@ -468,7 +523,7 @@ describe("runtime2 AgentHarness", () => { const options = modelOptions(session); await configureMain(session); const prompt = { role: "user" as const, content: "hello", timestamp: 1 }; - const promptId = await session.appendMessage(prompt); + const promptId = await session.appendMessage(prompt, BACKGROUND_CONTEXT); const meta = operationMeta(session, { kind: "run", promptEntryIds: [promptId] }, 2); const { operationId } = meta; const captured: LaneConfiguration = { @@ -495,7 +550,10 @@ describe("runtime2 AgentHarness", () => { }, }); - const { harness, suspended } = await createAgentHarness({ ...options, tools: [tool("available")] }); + const { harness, suspended } = await createAgentHarness( + { ...options, tools: [tool("available")] }, + BACKGROUND_CONTEXT, + ); const expected = { lane: "main", operationId, @@ -507,10 +565,10 @@ describe("runtime2 AgentHarness", () => { }; expect(suspended).toEqual([expected]); - expect((await harness.inspectExecution()).current?.suspended).toBe(suspended[0]); + expect((await harness.inspectExecution(BACKGROUND_CONTEXT)).current?.suspended).toBe(suspended[0]); const getEntries = vi.spyOn(session, "getEntries"); - await harness.inspectExecution(); - await harness.lanes(); + await harness.inspectExecution(BACKGROUND_CONTEXT); + await harness.lanes(BACKGROUND_CONTEXT); expect(getEntries).not.toHaveBeenCalled(); }); @@ -533,7 +591,7 @@ describe("runtime2 AgentHarness", () => { }, }); - const { harness, suspended } = await createAgentHarness(options); + const { harness, suspended } = await createAgentHarness(options, BACKGROUND_CONTEXT); expect(suspended).toEqual([ { @@ -546,7 +604,7 @@ describe("runtime2 AgentHarness", () => { missing: { tools: ["missing-tool"] }, }, ]); - expect((await harness.inspectExecution()).current?.suspended).toBe(suspended[0]); + expect((await harness.inspectExecution(BACKGROUND_CONTEXT)).current?.suspended).toBe(suspended[0]); }); it("restores deferred suspension handles from their source entries", async () => { @@ -554,7 +612,7 @@ describe("runtime2 AgentHarness", () => { const options = modelOptions(session); await configureMain(session); const prompt = { role: "user" as const, content: "defer", timestamp: 1 }; - const promptId = await session.appendMessage(prompt); + const promptId = await session.appendMessage(prompt, BACKGROUND_CONTEXT); const handle = { provider: options.model.provider, modelId: options.model.id, @@ -563,6 +621,7 @@ describe("runtime2 AgentHarness", () => { }; const sourceId = await session.appendMessage( fauxAssistantMessage([], { stopReason: "deferred", deferred: handle }), + BACKGROUND_CONTEXT, ); const meta = operationMeta(session, { kind: "run", promptEntryIds: [promptId] }, 2); const { operationId } = meta; @@ -581,7 +640,7 @@ describe("runtime2 AgentHarness", () => { }, }); - const { harness, suspended } = await createAgentHarness(options); + const { harness, suspended } = await createAgentHarness(options, BACKGROUND_CONTEXT); expect(suspended).toEqual([ { @@ -594,7 +653,7 @@ describe("runtime2 AgentHarness", () => { deferred: handle, }, ]); - expect((await harness.inspectExecution()).current?.suspended).toBe(suspended[0]); + expect((await harness.inspectExecution(BACKGROUND_CONTEXT)).current?.suspended).toBe(suspended[0]); }); it("faults harness creation when restored suspension payloads are invalid", async () => { @@ -616,7 +675,7 @@ describe("runtime2 AgentHarness", () => { }, }); - await expect(createAgentHarness(modelOptions(session))).rejects.toMatchObject({ + await expect(createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT)).rejects.toMatchObject({ name: "HarnessFault", cause: { name: "SessionInvariantError", message: "Deferred suspension source is invalid" }, }); @@ -628,7 +687,7 @@ describe("runtime2 AgentHarness", () => { operationMeta(missingPromptSession, { kind: "run", promptEntryIds: ["missing-prompt"] }), runState(missingPromptSession.idGenerator.next()), ); - await expect(createAgentHarness(modelOptions(missingPromptSession))).rejects.toMatchObject({ + await expect(createAgentHarness(modelOptions(missingPromptSession), BACKGROUND_CONTEXT)).rejects.toMatchObject({ name: "HarnessFault", cause: { name: "SessionInvariantError", message: "Prompt entry missing-prompt is missing" }, }); @@ -644,15 +703,15 @@ describe("runtime2 AgentHarness", () => { control: { status: "cancel_requested", requestedAt: 2, drainedSteer: [], drainedFollowUp: [] }, }); - const { harness, suspended } = await createAgentHarness(modelOptions(session)); + const { harness, suspended } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); expect(suspended).toHaveLength(1); - expect(await harness.inspectExecution()).toEqual({ + expect(await harness.inspectExecution(BACKGROUND_CONTEXT)).toEqual({ lane: "main", leafId: null, current: { id: operationId, kind: "run", status: "aborting", startedAt: 1 }, }); - expect(await harness.lanes()).toEqual([ + expect(await harness.lanes(BACKGROUND_CONTEXT)).toEqual([ { name: "main", leafId: null, operation: { id: operationId, kind: "run", status: "aborting" } }, ]); }); @@ -667,15 +726,26 @@ describe("runtime2 AgentHarness", () => { oldLeafId: null, leafId: null, }; - await session.mutate("main", (mutator) => - mutator.commit([ - storedValues.setValue(storedValues.laneConfig("main"), configuredMain), - storedValues.setValue(storedValues.laneLastResult("main"), lastResult), - ]), + await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + storedValues.setValue(storedValues.laneConfig("main"), configuredMain), + storedValues.setValue(storedValues.laneLastResult("main"), lastResult), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - const { harness } = await createAgentHarness(modelOptions(session)); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); - expect(await harness.inspectExecution()).toEqual({ lane: "main", leafId: null, current: null, lastResult }); + expect(await harness.inspectExecution(BACKGROUND_CONTEXT)).toEqual({ + lane: "main", + leafId: null, + current: null, + lastResult, + }); }); it("persists application values through idle and active lane facades", async () => { @@ -684,73 +754,79 @@ describe("runtime2 AgentHarness", () => { const meta = operationMeta(session, { kind: "run", promptEntryIds: [] }); const { operationId } = meta; await installOperation(session, meta, runState(session.idGenerator.next())); - const { harness } = await createAgentHarness(modelOptions(session)); - const worker = unwrap(await harness.createLane("worker", null)); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); + const worker = unwrap(await harness.createLane("worker", null, BACKGROUND_CONTEXT)); const updates: unknown[] = []; harness.events.on("value_update", (event) => { updates.push(event); }); - await worker.sessionTree.setName("idle-worker"); - await harness.sessionTree.setName("active-main"); - await harness.sessionTree.setLabel("target", "label"); - await harness.sessionTree.setValue(applicationValue, { durable: true }); + await worker.sessionTree.setName("idle-worker", BACKGROUND_CONTEXT); + await harness.sessionTree.setName("active-main", BACKGROUND_CONTEXT); + await harness.sessionTree.setLabel("target", "label", BACKGROUND_CONTEXT); + await harness.sessionTree.setValue(applicationValue, { durable: true }, BACKGROUND_CONTEXT); - expect(await session.getName()).toBe("active-main"); - expect(await session.getLabel("target")).toBe("label"); - expect((await session.getValue(applicationValue))?.value).toEqual({ durable: true }); + expect(await session.getName(BACKGROUND_CONTEXT)).toBe("active-main"); + expect(await session.getLabel("target", BACKGROUND_CONTEXT)).toBe("label"); + expect((await session.getValue(applicationValue, BACKGROUND_CONTEXT))?.value).toEqual({ durable: true }); expect(updates).toEqual([ { type: "value_update", value: "session_name", name: "idle-worker" }, { type: "value_update", value: "session_name", name: "active-main" }, { type: "value_update", value: "entry_label", targetId: "target", label: "label" }, ]); - expect((await harness.inspectExecution()).current?.id).toBe(operationId); + expect((await harness.inspectExecution(BACKGROUND_CONTEXT)).current?.id).toBe(operationId); }); it("reports restored open operations without activating them", async () => { const repo = new MemorySessionRepo(); repos.push(repo); - const session = await repo.create({}); + const session = await repo.create({}, BACKGROUND_CONTEXT); await configureMain(session); const meta = operationMeta(session, { kind: "run", promptEntryIds: [] }); const { operationId } = meta; const state = runState(session.idGenerator.next()); await installOperation(session, meta, state); - const { harness, suspended } = await createAgentHarness(modelOptions(session)); + const { harness, suspended } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); expect(suspended).toEqual([ { lane: "main", operationId, kind: "run", startedAt: 1, prompt: [], reason: "crash" }, ]); - expect(await harness.inspectExecution()).toEqual({ + expect(await harness.inspectExecution(BACKGROUND_CONTEXT)).toEqual({ lane: "main", leafId: null, current: { id: operationId, kind: "run", status: "suspended", startedAt: 1, suspended: suspended[0] }, }); - expect(await harness.lanes()).toEqual([ + expect(await harness.lanes(BACKGROUND_CONTEXT)).toEqual([ { name: "main", leafId: null, operation: { id: operationId, kind: "run", status: "suspended" } }, ]); - const pendingId = await harness.sessionTree.appendCustomEntry("note", { text: "queued" }); + const pendingId = await harness.sessionTree.appendCustomEntry("note", { text: "queued" }, BACKGROUND_CONTEXT); if (!(harness instanceof Harness)) throw new Error("missing runtime2 harness"); const updatedState = harness.state.operation?.state; if (updatedState?.kind !== "run") throw new Error("missing updated run state"); expect(updatedState.inbox.writes).toEqual([pendingId]); - expect((await session.getValue(storedValues.pendingEntry(pendingId)))?.value).toEqual({ + expect((await session.getValue(storedValues.pendingEntry(pendingId), BACKGROUND_CONTEXT))?.value).toEqual({ type: "custom", customType: "note", payload: { text: "queued" }, }); - expect((await session.getValue(storedValues.laneLeaf("main")))?.value).toBeNull(); + expect((await session.getValue(storedValues.laneLeaf("main"), BACKGROUND_CONTEXT))?.value).toBeNull(); - const closing = harness.close(); - expect(harness.close()).toBe(closing); + const closing = harness.close(BACKGROUND_CONTEXT); + expect(harness.close(BACKGROUND_CONTEXT)).toBe(closing); await closing; - await expect(harness.lanes()).rejects.toBeInstanceOf(HarnessClosed); - await expect(harness.inspectExecution()).rejects.toBeInstanceOf(HarnessClosed); - const reopened = await repo.open(session.metadata); - expect((await reopened.getValue(storedValues.laneState("main")))?.value.currentOperationId).toBe(operationId); - expect((await reopened.getValue(storedValues.operationMeta(operationId)))?.value).toEqual(meta); - expect((await reopened.getValue(storedValues.operationState(operationId)))?.value).toEqual(updatedState); + await expect(harness.lanes(BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(HarnessClosed); + await expect(harness.inspectExecution(BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(HarnessClosed); + const reopened = await repo.open(session.metadata, BACKGROUND_CONTEXT); + expect( + (await reopened.getValue(storedValues.laneState("main"), BACKGROUND_CONTEXT))?.value.currentOperationId, + ).toBe(operationId); + expect((await reopened.getValue(storedValues.operationMeta(operationId), BACKGROUND_CONTEXT))?.value).toEqual( + meta, + ); + expect((await reopened.getValue(storedValues.operationState(operationId), BACKGROUND_CONTEXT))?.value).toEqual( + updatedState, + ); }); it("uses owned state after creation and starts no option callbacks", async () => { @@ -758,34 +834,41 @@ describe("runtime2 AgentHarness", () => { const forbidden = vi.fn(() => { throw new Error("effect started"); }); - const { harness } = await createAgentHarness({ - ...modelOptions(session), - toolContext: forbidden, - systemPrompt: forbidden, - toProviderMessages: forbidden, - entryProjectors: { forbidden }, - }); + const { harness } = await createAgentHarness( + { + ...modelOptions(session), + toolContext: forbidden, + systemPrompt: forbidden, + toProviderMessages: forbidden, + entryProjectors: { forbidden }, + }, + BACKGROUND_CONTEXT, + ); const mutate = vi.spyOn(session, "mutate"); const getValue = vi.spyOn(session, "getValue"); const scanValues = vi.spyOn(session, "scanValues"); const readList = vi.spyOn(session, "readList"); - expect(await harness.lane("main")).toBe(harness); - expect(await harness.lanes()).toHaveLength(1); - expect(await harness.getLeafId()).toBeNull(); - expect(await harness.getLastResult()).toBeUndefined(); + expect(await harness.lane("main", BACKGROUND_CONTEXT)).toBe(harness); + expect(await harness.lanes(BACKGROUND_CONTEXT)).toHaveLength(1); + expect(await harness.getLeafId(BACKGROUND_CONTEXT)).toBeNull(); + expect(await harness.getLastResult(BACKGROUND_CONTEXT)).toBeUndefined(); expect(forbidden).not.toHaveBeenCalled(); expect(mutate).not.toHaveBeenCalled(); expect(getValue).not.toHaveBeenCalled(); expect(scanValues).not.toHaveBeenCalled(); expect(readList).not.toHaveBeenCalled(); - expect(await harness.getTools()).toEqual([]); - expect(await harness.getResources()).toEqual({}); - expect(await harness.getStreamOptions()).toEqual({}); - expect(await harness.getRetryPolicy()).toEqual({ enabled: true, maxRetries: 3, baseDelayMs: 1_000 }); - expect(await harness.getCompactionSettings()).toBe(DEFAULT_COMPACTION_SETTINGS); - expect(await harness.getSteeringMode()).toBe("all"); - expect(await harness.getFollowUpMode()).toBe("all"); + expect(await harness.getTools(BACKGROUND_CONTEXT)).toEqual([]); + expect(await harness.getResources(BACKGROUND_CONTEXT)).toEqual({}); + expect(await harness.getStreamOptions(BACKGROUND_CONTEXT)).toEqual({}); + expect(await harness.getRetryPolicy(BACKGROUND_CONTEXT)).toEqual({ + enabled: true, + maxRetries: 3, + baseDelayMs: 1_000, + }); + expect(await harness.getCompactionSettings(BACKGROUND_CONTEXT)).toBe(DEFAULT_COMPACTION_SETTINGS); + expect(await harness.getSteeringMode(BACKGROUND_CONTEXT)).toBe("all"); + expect(await harness.getFollowUpMode(BACKGROUND_CONTEXT)).toBe("all"); }); it("rejects append during a structural operation without faulting", async () => { @@ -801,21 +884,21 @@ describe("runtime2 AgentHarness", () => { phase: { kind: "ready_to_commit" }, }; await installOperation(session, meta, state); - const { harness, suspended } = await createAgentHarness(modelOptions(session)); + const { harness, suspended } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); expect(suspended).toEqual([{ lane: "main", operationId, kind: "navigation", startedAt: 1, reason: "crash" }]); - await expect(harness.sessionTree.appendMessage({ role: "user", content: "later", timestamp: 2 })).rejects.toThrow( - `Cannot append while structural operation ${operationId} is active`, - ); - expect(await harness.lanes()).toHaveLength(1); - await harness.setThinkingLevel("high"); + await expect( + harness.sessionTree.appendMessage({ role: "user", content: "later", timestamp: 2 }, BACKGROUND_CONTEXT), + ).rejects.toThrow(`Cannot append while structural operation ${operationId} is active`); + expect(await harness.lanes(BACKGROUND_CONTEXT)).toHaveLength(1); + await harness.setThinkingLevel("high", BACKGROUND_CONTEXT); }); it("faults queued lane work before releasing a failed mutation", async () => { const { session, storage } = await createFailingSession(); - const { harness } = await createAgentHarness(modelOptions(session)); + const { harness } = await createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT); if (!(harness instanceof Harness)) throw new Error("missing runtime2 harness"); - unwrap(await harness.createLane("worker", null)); + unwrap(await harness.createLane("worker", null, BACKGROUND_CONTEXT)); const worker = harness.lanesByName.get("worker"); if (worker === undefined) throw new Error("missing runtime2 worker lane"); const initialConfiguration = worker.state.configuration; @@ -824,8 +907,8 @@ describe("runtime2 AgentHarness", () => { harness.events.on("config_update", configUpdate); const failure = new Error("commit failed"); storage.failure = failure; - const failed = worker.setThinkingLevel("high"); - const queued = worker.setActiveTools(["read"]); + const failed = worker.setThinkingLevel("high", BACKGROUND_CONTEXT); + const queued = worker.setActiveTools(["read"], BACKGROUND_CONTEXT); let rejected: unknown; try { @@ -840,28 +923,36 @@ describe("runtime2 AgentHarness", () => { expect(await faultEvent).toMatchObject({ type: "fault", code: "harness_fault" }); expect(configUpdate).not.toHaveBeenCalled(); await expect(queued).rejects.toBe(rejected); - await expect(harness.getLeafId()).rejects.toBe(rejected); - await expect(worker.getThinkingLevel()).rejects.toBe(rejected); - await expect(worker.inspectExecution()).rejects.toBe(rejected); - await expect(harness.getTools()).rejects.toBe(rejected); - await expect(harness.setResources({})).rejects.toBe(rejected); - expect(harness.fault(new Error("later"))).toBe(rejected); + await expect(harness.getLeafId(BACKGROUND_CONTEXT)).rejects.toBe(rejected); + await expect(worker.getThinkingLevel(BACKGROUND_CONTEXT)).rejects.toBe(rejected); + await expect(worker.inspectExecution(BACKGROUND_CONTEXT)).rejects.toBe(rejected); + await expect(harness.getTools(BACKGROUND_CONTEXT)).rejects.toBe(rejected); + await expect(harness.setResources({}, BACKGROUND_CONTEXT)).rejects.toBe(rejected); + expect(harness.fault(new Error("later"), BACKGROUND_CONTEXT)).toBe(rejected); expect(worker.state.configuration).toBe(initialConfiguration); - expect((await session.getValue(storedValues.laneConfig("worker")))?.value).toEqual(initialConfiguration); - await harness.close(); + expect((await session.getValue(storedValues.laneConfig("worker"), BACKGROUND_CONTEXT))?.value).toEqual( + initialConfiguration, + ); + await harness.close(BACKGROUND_CONTEXT); }); it("wraps initialization invariant failures as harness faults", async () => { const session = await createSession(); - await session.mutate("main", (mutator) => - mutator.commit([ - storedValues.setValue(storedValues.laneState("main"), { - currentOperationId: session.idGenerator.next(), - pendingNextRun: [], - }), - ]), + await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + storedValues.setValue(storedValues.laneState("main"), { + currentOperationId: session.idGenerator.next(), + pendingNextRun: [], + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - await expect(createAgentHarness(modelOptions(session))).rejects.toBeInstanceOf(HarnessFault); + await expect(createAgentHarness(modelOptions(session), BACKGROUND_CONTEXT)).rejects.toBeInstanceOf(HarnessFault); }); }); diff --git a/packages/agent/test/harness/runtime2/lane.test.ts b/packages/agent/test/harness/runtime2/lane.test.ts index ac52f7146..dc1686ad3 100644 --- a/packages/agent/test/harness/runtime2/lane.test.ts +++ b/packages/agent/test/harness/runtime2/lane.test.ts @@ -1,6 +1,7 @@ import { type Api, createModels, fauxProvider, type Model } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it } from "vitest"; import { HarnessClosed } from "../../../src/harness/agent-harness.ts"; +import { BACKGROUND_CONTEXT } from "../../../src/harness/context.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"; @@ -27,12 +28,18 @@ async function createLane(): Promise<{ storage, ); sessions.push(session); - 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: null, pendingNextRun: [] }), - ]), + 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: null, pendingNextRun: [] }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); const faux = fauxProvider(); const models = createModels(); @@ -42,7 +49,7 @@ async function createLane(): Promise<{ "main", session, models, - await restoreLane(session, "main"), + await restoreLane(session, "main", BACKGROUND_CONTEXT), (cause) => (cause instanceof Error ? cause : new Error(String(cause))), () => Promise.resolve(), ), @@ -66,11 +73,11 @@ function setThinkingLevel( next, materialize: () => thinkingLevel, }; - }); + }, BACKGROUND_CONTEXT); } afterEach(async () => { - for (const session of sessions.splice(0)) await session.close(); + for (const session of sessions.splice(0)) await session.close(BACKGROUND_CONTEXT); }); describe("runtime2 Lane commands", () => { @@ -78,14 +85,14 @@ describe("runtime2 Lane commands", () => { const { lane, model, session } = await createLane(); const activeToolNames = ["read"]; - await lane.setModel(model); - await lane.setThinkingLevel("high"); - await lane.setActiveTools(activeToolNames); + await lane.setModel(model, BACKGROUND_CONTEXT); + await lane.setThinkingLevel("high", BACKGROUND_CONTEXT); + await lane.setActiveTools(activeToolNames, BACKGROUND_CONTEXT); - expect(await lane.getModel()).toBe(model); - expect(await lane.getThinkingLevel()).toBe("high"); - expect(await lane.getActiveTools()).toBe(activeToolNames); - expect((await session.getValue(storedValues.laneConfig("main")))?.value).toEqual({ + expect(await lane.getModel(BACKGROUND_CONTEXT)).toBe(model); + expect(await lane.getThinkingLevel(BACKGROUND_CONTEXT)).toBe("high"); + expect(await lane.getActiveTools(BACKGROUND_CONTEXT)).toBe(activeToolNames); + expect((await session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT))?.value).toEqual({ model: { provider: model.provider, modelId: model.id }, thinkingLevel: "high", activeToolNames, @@ -101,9 +108,9 @@ describe("runtime2 Lane commands", () => { await releaseCommit.promise; }; - const modelUpdate = lane.setModel(model); + const modelUpdate = lane.setModel(model, BACKGROUND_CONTEXT); await commitStarted.promise; - const thinkingUpdate = lane.setThinkingLevel("high"); + const thinkingUpdate = lane.setThinkingLevel("high", BACKGROUND_CONTEXT); releaseCommit.resolve(); await Promise.all([modelUpdate, thinkingUpdate]); @@ -119,12 +126,12 @@ describe("runtime2 Lane commands", () => { const completion = deferred(); let completed = false; const joined = lane - .command(() => ({ kind: "return", result: completion.promise })) + .command(() => ({ kind: "return", result: completion.promise }), BACKGROUND_CONTEXT) .then(() => { completed = true; }); - await lane.setThinkingLevel("high"); + await lane.setThinkingLevel("high", BACKGROUND_CONTEXT); expect(completed).toBe(false); completion.resolve(); await joined; @@ -134,9 +141,11 @@ describe("runtime2 Lane commands", () => { const { lane } = await createLane(); const rejection = new Error("declined"); - await expect(lane.command(() => ({ kind: "reject", error: rejection }))).rejects.toBe(rejection); - expect(await lane.getLeafId()).toBeNull(); - await lane.setThinkingLevel("high"); + await expect(lane.command(() => ({ kind: "reject", error: rejection }), BACKGROUND_CONTEXT)).rejects.toBe( + rejection, + ); + expect(await lane.getLeafId(BACKGROUND_CONTEXT)).toBeNull(); + await lane.setThinkingLevel("high", BACKGROUND_CONTEXT); }); it("passes bounded reads and commit metadata through the serialized command", async () => { @@ -145,7 +154,7 @@ describe("runtime2 Lane commands", () => { let memoryPublished = false; const commit = await lane.command(async (state, reader) => { - storedConfiguration = (await reader.getValue(storedValues.laneConfig("main")))?.value; + storedConfiguration = (await reader.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT))?.value; const configuration: LaneConfiguration = { ...state.configuration, thinkingLevel: "high" }; const next = { ...state, configuration }; return { @@ -157,7 +166,7 @@ describe("runtime2 Lane commands", () => { return commit; }, }; - }); + }, BACKGROUND_CONTEXT); expect(storedConfiguration).toEqual(configuration); expect(memoryPublished).toBe(true); @@ -178,11 +187,13 @@ describe("runtime2 Lane commands", () => { next: { ...state, configuration }, materialize: async () => undefined, }; - }), + }, BACKGROUND_CONTEXT), ).rejects.toThrow("Lane command materialize() must be synchronous"); expect(lane.state.configuration.thinkingLevel).toBe("high"); - expect((await session.getValue(storedValues.laneConfig("main")))?.value.thinkingLevel).toBe("high"); + expect((await session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT))?.value.thinkingLevel).toBe( + "high", + ); }); it("rejects work after sealing while an admitted commit finishes", async () => { @@ -199,8 +210,8 @@ describe("runtime2 Lane commands", () => { const closed = new HarnessClosed(); lane.seal(closed); - await expect(lane.getLeafId()).rejects.toBe(closed); - await expect(lane.setThinkingLevel("low")).rejects.toBe(closed); + await expect(lane.getLeafId(BACKGROUND_CONTEXT)).rejects.toBe(closed); + await expect(lane.setThinkingLevel("low", BACKGROUND_CONTEXT)).rejects.toBe(closed); releaseCommit.resolve(); await admitted; expect(lane.state.configuration.thinkingLevel).toBe("high"); @@ -222,7 +233,9 @@ describe("runtime2 Lane commands", () => { releaseCommit.resolve(); expect(await command).toBe("high"); expect(lane.state.configuration.thinkingLevel).toBe("high"); - expect((await session.getValue(storedValues.laneConfig("main")))?.value.thinkingLevel).toBe("high"); + expect((await session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT))?.value.thinkingLevel).toBe( + "high", + ); }); it("preserves memory when the durable commit fails", async () => { @@ -235,7 +248,9 @@ describe("runtime2 Lane commands", () => { await expect(setThinkingLevel(lane, "high")).rejects.toBe(failure); expect(lane.state.configuration.thinkingLevel).toBe("off"); - expect((await session.getValue(storedValues.laneConfig("main")))?.value.thinkingLevel).toBe("off"); + expect((await session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT))?.value.thinkingLevel).toBe( + "off", + ); }); it("plans queued commands from the latest committed memory", async () => { diff --git a/packages/agent/test/harness/runtime2/restore.test.ts b/packages/agent/test/harness/runtime2/restore.test.ts index d2b94ae09..cdea7e278 100644 --- a/packages/agent/test/harness/runtime2/restore.test.ts +++ b/packages/agent/test/harness/runtime2/restore.test.ts @@ -1,5 +1,6 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { DEFAULT_COMPACTION_SETTINGS } from "../../../src/harness/compaction/compaction.ts"; +import { BACKGROUND_CONTEXT } from "../../../src/harness/context.ts"; import { restoreLane, restoreSession } from "../../../src/harness/runtime2/restore.ts"; import { MemorySessionRepo, MemoryStorage } from "../../../src/harness/session/memory.ts"; import type { LaneConfiguration, OperationMeta, RunState, Session } from "../../../src/harness/session/types.ts"; @@ -15,9 +16,12 @@ const configuration: LaneConfiguration = { async function createSession(): Promise { const repo = new MemorySessionRepo(); repos.push(repo); - const session = await repo.create({}); - await session.mutate("main", (mutator) => - mutator.commit([storedValues.setValue(storedValues.laneConfig("main"), configuration)]), + const session = await repo.create({}, BACKGROUND_CONTEXT); + await session.mutate( + "main", + (mutator) => + mutator.commit([storedValues.setValue(storedValues.laneConfig("main"), configuration)], BACKGROUND_CONTEXT), + BACKGROUND_CONTEXT, ); return session; } @@ -43,7 +47,7 @@ function runState(triggerEntryId: string): RunState { } afterEach(async () => { - for (const repo of repos.splice(0)) await repo.close(); + for (const repo of repos.splice(0)) await repo.close(BACKGROUND_CONTEXT); }); describe("runtime2 lane restore", () => { @@ -56,11 +60,14 @@ describe("runtime2 lane restore", () => { oldLeafId: null, leafId: null, }; - await session.mutate("main", (mutator) => - mutator.commit([storedValues.setValue(storedValues.laneLastResult("main"), result)]), + await session.mutate( + "main", + (mutator) => + mutator.commit([storedValues.setValue(storedValues.laneLastResult("main"), result)], BACKGROUND_CONTEXT), + BACKGROUND_CONTEXT, ); - const state = await restoreLane(session, "main"); + const state = await restoreLane(session, "main", BACKGROUND_CONTEXT); expect(state).toEqual({ leafId: null, @@ -83,18 +90,24 @@ describe("runtime2 lane restore", () => { intent: { kind: "run", promptEntryIds: [] }, }; const state = runState(missingTriggerId); - await session.mutate("main", (mutator) => - mutator.commit([ - storedValues.setValue(storedValues.operationMeta(operationId), meta), - storedValues.setValue(storedValues.operationState(operationId), state), - storedValues.setValue(storedValues.laneState("main"), { - currentOperationId: operationId, - pendingNextRun: [], - }), - ]), + await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + storedValues.setValue(storedValues.operationMeta(operationId), meta), + storedValues.setValue(storedValues.operationState(operationId), state), + storedValues.setValue(storedValues.laneState("main"), { + currentOperationId: operationId, + pendingNextRun: [], + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - const restored = await restoreLane(session, "main"); + const restored = await restoreLane(session, "main", BACKGROUND_CONTEXT); expect(restored.operation).toEqual({ meta, state }); }); @@ -105,9 +118,9 @@ describe("runtime2 lane restore", () => { [storedValues.laneState("main").namespace, storedValues.deleteValue(storedValues.laneState("main"))], ] as const)("requires %s", async (namespace, write) => { const session = await createSession(); - await session.mutate("main", (mutator) => mutator.commit([write])); + await session.mutate("main", (mutator) => mutator.commit([write], BACKGROUND_CONTEXT), BACKGROUND_CONTEXT); - await expect(restoreLane(session, "main")).rejects.toThrow(`missing ${namespace.slice(3)}`); + await expect(restoreLane(session, "main", BACKGROUND_CONTEXT)).rejects.toThrow(`missing ${namespace.slice(3)}`); }); it.each([storedValues.operationMeta("").namespace, storedValues.operationState("").namespace] as const)( @@ -123,22 +136,30 @@ describe("runtime2 lane restore", () => { intent: { kind: "run", promptEntryIds: [] }, }; const state = runState(session.idGenerator.next()); - await session.mutate("main", (mutator) => - mutator.commit([ - ...(namespace === storedValues.operationMeta(operationId).namespace - ? [] - : [storedValues.setValue(storedValues.operationMeta(operationId), meta)]), - ...(namespace === storedValues.operationState(operationId).namespace - ? [] - : [storedValues.setValue(storedValues.operationState(operationId), state)]), - storedValues.setValue(storedValues.laneState("main"), { - currentOperationId: operationId, - pendingNextRun: [], - }), - ]), + await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + ...(namespace === storedValues.operationMeta(operationId).namespace + ? [] + : [storedValues.setValue(storedValues.operationMeta(operationId), meta)]), + ...(namespace === storedValues.operationState(operationId).namespace + ? [] + : [storedValues.setValue(storedValues.operationState(operationId), state)]), + storedValues.setValue(storedValues.laneState("main"), { + currentOperationId: operationId, + pendingNextRun: [], + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - await expect(restoreLane(session, "main")).rejects.toThrow(`missing ${namespace.slice(3)}`); + await expect(restoreLane(session, "main", BACKGROUND_CONTEXT)).rejects.toThrow( + `missing ${namespace.slice(3)}`, + ); }, ); @@ -149,7 +170,7 @@ describe("runtime2 lane restore", () => { thinkingLevel: "high", activeToolNames: ["read"], }; - await session.createLane("worker", null, workerConfiguration); + await session.createLane("worker", null, workerConfiguration, BACKGROUND_CONTEXT); const operationId = session.idGenerator.next(); const meta: OperationMeta = { operationId, @@ -159,27 +180,33 @@ describe("runtime2 lane restore", () => { intent: { kind: "run", promptEntryIds: [] }, }; const state = runState(session.idGenerator.next()); - await session.mutate("worker", (mutator) => - mutator.commit([ - storedValues.setValue(storedValues.operationMeta(operationId), meta), - storedValues.setValue(storedValues.operationState(operationId), state), - storedValues.setValue(storedValues.laneState("worker"), { - currentOperationId: operationId, - pendingNextRun: [], - }), - ]), + await session.mutate( + "worker", + (mutator) => + mutator.commit( + [ + storedValues.setValue(storedValues.operationMeta(operationId), meta), + storedValues.setValue(storedValues.operationState(operationId), state), + storedValues.setValue(storedValues.laneState("worker"), { + currentOperationId: operationId, + pendingNextRun: [], + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); const before = { - leaves: await session.scanValues(storedValues.laneLeafInventoryPrefix()), - mainConfiguration: await session.getValue(storedValues.laneConfig("main")), - workerConfiguration: await session.getValue(storedValues.laneConfig("worker")), - mainState: await session.getValue(storedValues.laneState("main")), - workerState: await session.getValue(storedValues.laneState("worker")), + leaves: await session.scanValues(storedValues.laneLeafInventoryPrefix(), BACKGROUND_CONTEXT), + mainConfiguration: await session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT), + workerConfiguration: await session.getValue(storedValues.laneConfig("worker"), BACKGROUND_CONTEXT), + mainState: await session.getValue(storedValues.laneState("main"), BACKGROUND_CONTEXT), + workerState: await session.getValue(storedValues.laneState("worker"), BACKGROUND_CONTEXT), }; const mutate = vi.spyOn(session, "mutate"); const readList = vi.spyOn(MemoryStorage.prototype, "readList"); - const lanes = await restoreSession(session); + const lanes = await restoreSession(session, BACKGROUND_CONTEXT); expect([...lanes.keys()].sort()).toEqual(["main", "worker"]); expect(lanes.get("main")?.configuration).toEqual(configuration); @@ -190,21 +217,23 @@ describe("runtime2 lane restore", () => { expect(mutate.mock.calls.map(([lane]) => lane).sort()).toEqual(["main", "worker"]); expect(readList).not.toHaveBeenCalled(); expect({ - leaves: await session.scanValues(storedValues.laneLeafInventoryPrefix()), - mainConfiguration: await session.getValue(storedValues.laneConfig("main")), - workerConfiguration: await session.getValue(storedValues.laneConfig("worker")), - mainState: await session.getValue(storedValues.laneState("main")), - workerState: await session.getValue(storedValues.laneState("worker")), + leaves: await session.scanValues(storedValues.laneLeafInventoryPrefix(), BACKGROUND_CONTEXT), + mainConfiguration: await session.getValue(storedValues.laneConfig("main"), BACKGROUND_CONTEXT), + workerConfiguration: await session.getValue(storedValues.laneConfig("worker"), BACKGROUND_CONTEXT), + mainState: await session.getValue(storedValues.laneState("main"), BACKGROUND_CONTEXT), + workerState: await session.getValue(storedValues.laneState("worker"), BACKGROUND_CONTEXT), }).toEqual(before); readList.mockRestore(); }); it("requires main in the trusted lane inventory", async () => { const session = await createSession(); - await session.mutate("main", (mutator) => - mutator.commit([storedValues.deleteValue(storedValues.laneLeaf("main"))]), + await session.mutate( + "main", + (mutator) => mutator.commit([storedValues.deleteValue(storedValues.laneLeaf("main"))], BACKGROUND_CONTEXT), + BACKGROUND_CONTEXT, ); - await expect(restoreSession(session)).rejects.toThrow("Session is missing main lane"); + await expect(restoreSession(session, BACKGROUND_CONTEXT)).rejects.toThrow("Session is missing main lane"); }); }); diff --git a/packages/agent/test/harness/runtime2/test-utils.ts b/packages/agent/test/harness/runtime2/test-utils.ts index a2179b689..b1e4883d3 100644 --- a/packages/agent/test/harness/runtime2/test-utils.ts +++ b/packages/agent/test/harness/runtime2/test-utils.ts @@ -1,3 +1,4 @@ +import { BACKGROUND_CONTEXT } from "../../../src/harness/context.ts"; import { MemoryStorage } from "../../../src/harness/session/memory.ts"; import type { CommitResult, Write } from "../../../src/harness/session/types.ts"; @@ -7,7 +8,7 @@ export class FailingMemoryStorage extends MemoryStorage { override commit(writes: Write[]) { const failure = this.failure; this.failure = undefined; - return failure === undefined ? super.commit(writes) : Promise.reject(failure); + return failure === undefined ? super.commit(writes, BACKGROUND_CONTEXT) : Promise.reject(failure); } } @@ -18,7 +19,7 @@ export class ControlledMemoryStorage extends MemoryStorage { const beforeCommit = this.beforeNextCommit; this.beforeNextCommit = undefined; await beforeCommit?.(); - return super.commit(writes); + return super.commit(writes, BACKGROUND_CONTEXT); } } diff --git a/packages/agent/test/harness/session-context.test.ts b/packages/agent/test/harness/session-context.test.ts index 2e5cace42..87c4efdb4 100644 --- a/packages/agent/test/harness/session-context.test.ts +++ b/packages/agent/test/harness/session-context.test.ts @@ -1,5 +1,6 @@ import type { AssistantMessage } from "@earendil-works/pi-ai"; import { describe, expect, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import { buildSessionContext } from "../../src/harness/session/context.ts"; import type { BranchSummaryEntry, @@ -65,7 +66,12 @@ describe("session context projection", () => { messageEntry("length", "deferred", 7, length), ]; - expect(await buildSessionContext(entries)).toEqual([user, stopped, toolUse, length]); + expect(await buildSessionContext(entries, undefined, BACKGROUND_CONTEXT)).toEqual([ + user, + stopped, + toolUse, + length, + ]); }); it("projects branch summaries in branch order", async () => { @@ -82,7 +88,7 @@ describe("session context projection", () => { }; const after = messageEntry("after", summary.id, 3, userMessage("after summary")); - expect(await buildSessionContext([before, summary, after])).toEqual([ + expect(await buildSessionContext([before, summary, after], undefined, BACKGROUND_CONTEXT)).toEqual([ userMessage("before summary"), { role: "branchSummary", @@ -114,7 +120,7 @@ describe("session context projection", () => { fromHook: false, }; - expect(await buildSessionContext([compaction])).toEqual([ + expect(await buildSessionContext([compaction], undefined, BACKGROUND_CONTEXT)).toEqual([ { role: "compactionSummary", summary: "summary", tokensBefore: 100, timestamp: NOW }, user, stopped, @@ -149,7 +155,9 @@ describe("session context projection", () => { }; const after = messageEntry("after", latest.id, 5, userMessage("after latest")); - expect(await buildSessionContext([beforeFirst, first, between, latest, after])).toEqual([ + expect( + await buildSessionContext([beforeFirst, first, between, latest, after], undefined, BACKGROUND_CONTEXT), + ).toEqual([ { role: "compactionSummary", summary: "latest summary", tokensBefore: 200, timestamp: NOW }, userMessage("latest tail"), userMessage("after latest"), @@ -202,19 +210,23 @@ describe("session context projection", () => { }; const projectedIds: string[] = []; - const messages = await buildSessionContext([oldCustom, compaction, syncCustom, omittedCustom, asyncCustom], { - entryProjectors: { - sync: (entry) => { - projectedIds.push(entry.id); - return [userMessage(`projected:${entry.id}`)]; - }, - async: async (entry) => { - await Promise.resolve(); - projectedIds.push(entry.id); - return [userMessage(`projected:${entry.id}`)]; + const messages = await buildSessionContext( + [oldCustom, compaction, syncCustom, omittedCustom, asyncCustom], + { + entryProjectors: { + sync: (entry) => { + projectedIds.push(entry.id); + return [userMessage(`projected:${entry.id}`)]; + }, + async: async (entry) => { + await Promise.resolve(); + projectedIds.push(entry.id); + return [userMessage(`projected:${entry.id}`)]; + }, }, }, - }); + BACKGROUND_CONTEXT, + ); expect(projectedIds).toEqual([syncCustom.id, asyncCustom.id]); expect(messages).toEqual([ @@ -236,9 +248,13 @@ describe("session context projection", () => { const failure = new Error("projector failed"); await expect( - buildSessionContext([custom], { - entryProjectors: { broken: async () => Promise.reject(failure) }, - }), + buildSessionContext( + [custom], + { + entryProjectors: { broken: async () => Promise.reject(failure) }, + }, + BACKGROUND_CONTEXT, + ), ).rejects.toBe(failure); }); }); diff --git a/packages/agent/test/harness/session-create-lane.test.ts b/packages/agent/test/harness/session-create-lane.test.ts index 0db19ce49..9d3bd8250 100644 --- a/packages/agent/test/harness/session-create-lane.test.ts +++ b/packages/agent/test/harness/session-create-lane.test.ts @@ -1,3 +1,4 @@ +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; // Implements Part 8, slice 2. // Full state validation is R1; harness-wide close admission is R6. @@ -40,7 +41,7 @@ function rootTransaction(): Write[] { } function commitSession(session: Session, transaction: Write[]) { - return session.mutate("main", (mutator) => mutator.commit(transaction)); + return session.mutate("main", (mutator) => mutator.commit(transaction, BACKGROUND_CONTEXT), BACKGROUND_CONTEXT); } function expectedLaneWrites(name: string, at: string | null, value: LaneConfiguration): Write[] { @@ -69,7 +70,9 @@ class RejectingCommitStorage extends StorageDecorator { rejection: Error | undefined; override commit(transaction: Write[]) { - return this.rejection === undefined ? super.commit(transaction) : Promise.reject(this.rejection); + return this.rejection === undefined + ? super.commit(transaction, BACKGROUND_CONTEXT) + : Promise.reject(this.rejection); } } @@ -87,7 +90,7 @@ class BlockingCommitStorage extends StorageDecorator { this.admittedGate.resolve(); await this.releaseGate.promise; } - return super.commit(transaction); + return super.commit(transaction, BACKGROUND_CONTEXT); } release(): void { @@ -106,20 +109,26 @@ describe("StorageBackedSession.createLane", () => { await commitSession(session, rootTransaction()); storage.clearCommitAttempts(); - const rooted = await session.createLane("rooted", ROOT_ID, configuration); - const empty = await session.createLane("empty", null, configuration); + const rooted = await session.createLane("rooted", ROOT_ID, configuration, BACKGROUND_CONTEXT); + const empty = await session.createLane("empty", null, configuration, BACKGROUND_CONTEXT); expect(storage.getCommitAttempts()).toEqual([ expectedLaneWrites("rooted", ROOT_ID, configuration), expectedLaneWrites("empty", null, configuration), ]); - expect(await rooted.getLeafId()).toBe(ROOT_ID); - expect((await rooted.findEntriesOnBranch()).map((entry: Entry) => entry.id)).toEqual([ROOT_ID]); - expect(await empty.getLeafId()).toBeNull(); - expect(await empty.findEntriesOnBranch()).toEqual([]); - expect(await session.getValue(storedValues.laneConfig("rooted"))).toMatchObject({ value: configuration }); - expect(await session.getValue(storedValues.laneState("rooted"))).toMatchObject({ value: idleLaneState }); - await session.close(); + expect(await rooted.getLeafId(BACKGROUND_CONTEXT)).toBe(ROOT_ID); + expect((await rooted.findEntriesOnBranch(undefined, BACKGROUND_CONTEXT)).map((entry: Entry) => entry.id)).toEqual( + [ROOT_ID], + ); + expect(await empty.getLeafId(BACKGROUND_CONTEXT)).toBeNull(); + expect(await empty.findEntriesOnBranch(undefined, BACKGROUND_CONTEXT)).toEqual([]); + expect(await session.getValue(storedValues.laneConfig("rooted"), BACKGROUND_CONTEXT)).toMatchObject({ + value: configuration, + }); + expect(await session.getValue(storedValues.laneState("rooted"), BACKGROUND_CONTEXT)).toMatchObject({ + value: idleLaneState, + }); + await session.close(BACKGROUND_CONTEXT); }); it("passes configuration directly to storage", async () => { @@ -133,11 +142,11 @@ describe("StorageBackedSession.createLane", () => { activeToolNames: [...configuration.activeToolNames], }; - await session.createLane("captured", ROOT_ID, supplied); + await session.createLane("captured", ROOT_ID, supplied, BACKGROUND_CONTEXT); - expect((await session.getValue(storedValues.laneConfig("captured")))?.value).toBe(supplied); + expect((await session.getValue(storedValues.laneConfig("captured"), BACKGROUND_CONTEXT))?.value).toBe(supplied); expect(storage.getCommitAttempts()).toHaveLength(1); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); // Slice 2 defines classifiable session failures. Mapping them to AgentHarness @@ -148,29 +157,29 @@ describe("StorageBackedSession.createLane", () => { await commitSession(session, rootTransaction()); storage.clearCommitAttempts(); - await expect(session.createLane("", ROOT_ID, configuration)).rejects.toMatchObject({ + await expect(session.createLane("", ROOT_ID, configuration, BACKGROUND_CONTEXT)).rejects.toMatchObject({ name: "SessionInvalidLaneError", lane: "", }); expect(storage.getCommitAttempts()).toEqual([]); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("rejects existing configured lanes and fresh unconfigured main without writing", async () => { const storage = new InstrumentedStorage(new MemoryStorage({ now: () => NOW })); const session = new StorageBackedSession(metadata, storage); await commitSession(session, rootTransaction()); - await session.createLane("existing", ROOT_ID, configuration); + await session.createLane("existing", ROOT_ID, configuration, BACKGROUND_CONTEXT); storage.clearCommitAttempts(); for (const name of ["existing", "main"]) { - await expect(session.createLane(name, null, configuration)).rejects.toMatchObject({ + await expect(session.createLane(name, null, configuration, BACKGROUND_CONTEXT)).rejects.toMatchObject({ name: "SessionLaneExistsError", lane: name, }); } expect(storage.getCommitAttempts()).toEqual([]); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("rejects unknown non-null anchors without writing", async () => { @@ -179,12 +188,14 @@ describe("StorageBackedSession.createLane", () => { await commitSession(session, rootTransaction()); storage.clearCommitAttempts(); - await expect(session.createLane("missing-target", MISSING_ID, configuration)).rejects.toMatchObject({ + await expect( + session.createLane("missing-target", MISSING_ID, configuration, BACKGROUND_CONTEXT), + ).rejects.toMatchObject({ name: "SessionUnknownTargetError", targetId: MISSING_ID, }); expect(storage.getCommitAttempts()).toEqual([]); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("fails fast on every partial lane-register combination instead of repairing it", async () => { @@ -212,9 +223,11 @@ describe("StorageBackedSession.createLane", () => { ); storage.clearCommitAttempts(); - await expect(session.createLane("broken", null, configuration)).rejects.toBeInstanceOf(SessionInvariantError); + await expect(session.createLane("broken", null, configuration, BACKGROUND_CONTEXT)).rejects.toBeInstanceOf( + SessionInvariantError, + ); expect(storage.getCommitAttempts()).toEqual([]); - await session.close(); + await session.close(BACKGROUND_CONTEXT); } }); @@ -225,8 +238,8 @@ describe("StorageBackedSession.createLane", () => { storage.clearCommitAttempts(); const results = await Promise.allSettled([ - session.createLane("race", ROOT_ID, configuration), - session.createLane("race", null, configuration), + session.createLane("race", ROOT_ID, configuration, BACKGROUND_CONTEXT), + session.createLane("race", null, configuration, BACKGROUND_CONTEXT), ]); const fulfilled = results.filter((result) => result.status === "fulfilled"); @@ -235,28 +248,31 @@ describe("StorageBackedSession.createLane", () => { expect(rejected).toHaveLength(1); expect(rejected[0]).toMatchObject({ reason: { name: "SessionLaneExistsError", lane: "race" } }); expect(storage.getCommitAttempts()).toHaveLength(1); - const leaf = await session.view("race").getLeafId(); + const leaf = await session.view("race").getLeafId(BACKGROUND_CONTEXT); expect([ROOT_ID, null]).toContain(leaf); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("orders lane creation with appends submitted through the prospective view", async () => { const session = new StorageBackedSession(metadata, new MemoryStorage({ now: () => NOW })); await commitSession(session, rootTransaction()); - const creating = session.createLane("created-first", ROOT_ID, configuration); - const appended = session.view("created-first").appendCustomEntry("note"); + const creating = session.createLane("created-first", ROOT_ID, configuration, BACKGROUND_CONTEXT); + const appended = session.view("created-first").appendCustomEntry("note", undefined, BACKGROUND_CONTEXT); const [, entryId] = await Promise.all([creating, appended]); - expect(await session.view("created-first").getLeafId()).toBe(entryId); - expect(await session.getEntry(entryId)).toMatchObject({ parentId: ROOT_ID, customType: "note" }); + expect(await session.view("created-first").getLeafId(BACKGROUND_CONTEXT)).toBe(entryId); + expect(await session.getEntry(entryId, BACKGROUND_CONTEXT)).toMatchObject({ + parentId: ROOT_ID, + customType: "note", + }); - const rejectedAppend = session.view("append-first").appendCustomEntry("note"); - const laterCreation = session.createLane("append-first", ROOT_ID, configuration); + const rejectedAppend = session.view("append-first").appendCustomEntry("note", undefined, BACKGROUND_CONTEXT); + const laterCreation = session.createLane("append-first", ROOT_ID, configuration, BACKGROUND_CONTEXT); await expect(rejectedAppend).rejects.toBeInstanceOf(SessionInvariantError); await laterCreation; - expect(await session.view("append-first").getLeafId()).toBe(ROOT_ID); - expect(await session.findEntries({ customType: "note" })).toHaveLength(1); - await session.close(); + expect(await session.view("append-first").getLeafId(BACKGROUND_CONTEXT)).toBe(ROOT_ID); + expect(await session.findEntries({ customType: "note" }, BACKGROUND_CONTEXT)).toHaveLength(1); + await session.close(BACKGROUND_CONTEXT); }); it("propagates commit failure without publishing partial lane state", async () => { @@ -266,12 +282,12 @@ describe("StorageBackedSession.createLane", () => { const rejection = new Error("commit failed"); storage.rejection = rejection; - await expect(session.createLane("failed", ROOT_ID, configuration)).rejects.toBe(rejection); + await expect(session.createLane("failed", ROOT_ID, configuration, BACKGROUND_CONTEXT)).rejects.toBe(rejection); storage.rejection = undefined; - expect(await session.getValue(storedValues.laneConfig("failed"))).toBeUndefined(); - expect(await session.getValue(storedValues.laneLeaf("failed"))).toBeUndefined(); - expect(await session.getValue(storedValues.laneState("failed"))).toBeUndefined(); - await session.close(); + expect(await session.getValue(storedValues.laneConfig("failed"), BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await session.getValue(storedValues.laneLeaf("failed"), BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await session.getValue(storedValues.laneState("failed"), BACKGROUND_CONTEXT)).toBeUndefined(); + await session.close(BACKGROUND_CONTEXT); }); it("drains creation after its commit is admitted to storage and rejects creation after close", async () => { @@ -280,14 +296,16 @@ describe("StorageBackedSession.createLane", () => { await commitSession(session, rootTransaction()); storage.block = true; - const creation = session.createLane("admitted", ROOT_ID, configuration); + const creation = session.createLane("admitted", ROOT_ID, configuration, BACKGROUND_CONTEXT); await storage.admitted; - const close = session.close(); + const close = session.close(BACKGROUND_CONTEXT); storage.release(); await creation; await close; - await expect(session.createLane("late", null, configuration)).rejects.toThrow("Session is closed"); + await expect(session.createLane("late", null, configuration, BACKGROUND_CONTEXT)).rejects.toThrow( + "Session is closed", + ); }); it("rejects queued duplicate creation when close seals the lane mutation line", async () => { @@ -296,10 +314,10 @@ describe("StorageBackedSession.createLane", () => { await commitSession(session, rootTransaction()); storage.block = true; - const admitted = session.createLane("queued", ROOT_ID, configuration); + const admitted = session.createLane("queued", ROOT_ID, configuration, BACKGROUND_CONTEXT); await storage.admitted; - const queued = session.createLane("queued", null, configuration); - const close = session.close(); + const queued = session.createLane("queued", null, configuration, BACKGROUND_CONTEXT); + const close = session.close(BACKGROUND_CONTEXT); storage.release(); await admitted; diff --git a/packages/agent/test/harness/session-tree.test.ts b/packages/agent/test/harness/session-tree.test.ts index 1d2dfdad3..b87a45813 100644 --- a/packages/agent/test/harness/session-tree.test.ts +++ b/packages/agent/test/harness/session-tree.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import * as sessionWrites from "../../src/harness/session/commit.ts"; import { LaneMutationLine } from "../../src/harness/session/lane-mutations.ts"; import { MemoryStorage } from "../../src/harness/session/memory.ts"; @@ -60,7 +61,7 @@ const runState = { } satisfies RunState; function commitSession(session: Session, transaction: Write[]) { - return session.mutate("main", (mutator) => mutator.commit(transaction)); + return session.mutate("main", (mutator) => mutator.commit(transaction, BACKGROUND_CONTEXT), BACKGROUND_CONTEXT); } function laneWrites(leafId: string | null, state: LaneState = idleLaneState) { @@ -101,7 +102,9 @@ class RejectingCommitStorage extends MemoryStorage { rejection: Error | undefined; override commit(transaction: Parameters[0]) { - return this.rejection === undefined ? super.commit(transaction) : Promise.reject(this.rejection); + return this.rejection === undefined + ? super.commit(transaction, BACKGROUND_CONTEXT) + : Promise.reject(this.rejection); } } @@ -117,7 +120,7 @@ class BlockingCommitStorage extends MemoryStorage { this.releaseCommit = resolve; }); } - return super.commit(transaction); + return super.commit(transaction, BACKGROUND_CONTEXT); } release(): void { @@ -131,55 +134,55 @@ describe("StorageBackedSession SessionTree", () => { const session = await createTreeSession(); const other = session.view("other"); - expect(await session.getLeafId()).toBe(CUSTOM_ID); - expect(await other.getLeafId()).toBe(OTHER_ID); - expect(await session.getEntry(CHILD_ID)).toMatchObject({ id: CHILD_ID, parentId: ROOT_ID }); - expect(await session.getEntry("00000000-0000-7000-8000-000000000099")).toBeUndefined(); - expect(await session.getStats()).toMatchObject({ messageCount: 1 }); - expect(await other.getStats()).toEqual(await session.getStats()); + expect(await session.getLeafId(BACKGROUND_CONTEXT)).toBe(CUSTOM_ID); + expect(await other.getLeafId(BACKGROUND_CONTEXT)).toBe(OTHER_ID); + expect(await session.getEntry(CHILD_ID, BACKGROUND_CONTEXT)).toMatchObject({ id: CHILD_ID, parentId: ROOT_ID }); + expect(await session.getEntry("00000000-0000-7000-8000-000000000099", BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await session.getStats(BACKGROUND_CONTEXT)).toMatchObject({ messageCount: 1 }); + expect(await other.getStats(BACKGROUND_CONTEXT)).toEqual(await session.getStats(BACKGROUND_CONTEXT)); - await session.setName("shared"); - expect(await other.getName()).toBe("shared"); - await session.close(); + await session.setName("shared", BACKGROUND_CONTEXT); + expect(await other.getName(BACKGROUND_CONTEXT)).toBe("shared"); + await session.close(BACKGROUND_CONTEXT); }); it("rejects missing lanes", async () => { const session = await createTreeSession(); - await expect(session.view("missing").getLeafId()).rejects.toThrow("Unknown lane"); - await session.close(); + await expect(session.view("missing").getLeafId(BACKGROUND_CONTEXT)).rejects.toThrow("Unknown lane"); + await session.close(BACKGROUND_CONTEXT); }); it("reads, sets, and deletes global values while preserving JSON null", async () => { const session = await createTreeSession(); const other = session.view("other"); - expect(await session.getName()).toBeUndefined(); - expect(await session.getLabel(ROOT_ID)).toBeUndefined(); - expect(await session.getValue(applicationState)).toBeUndefined(); + expect(await session.getName(BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await session.getLabel(ROOT_ID, BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await session.getValue(applicationState, BACKGROUND_CONTEXT)).toBeUndefined(); - await session.setName("name"); - await other.setLabel(ROOT_ID, "root label"); - await session.setValue(applicationState, null); - expect(await other.getName()).toBe("name"); - expect(await session.getLabel(ROOT_ID)).toBe("root label"); - expect((await other.getValue(applicationState))?.value).toBeNull(); + await session.setName("name", BACKGROUND_CONTEXT); + await other.setLabel(ROOT_ID, "root label", BACKGROUND_CONTEXT); + await session.setValue(applicationState, null, BACKGROUND_CONTEXT); + expect(await other.getName(BACKGROUND_CONTEXT)).toBe("name"); + expect(await session.getLabel(ROOT_ID, BACKGROUND_CONTEXT)).toBe("root label"); + expect((await other.getValue(applicationState, BACKGROUND_CONTEXT))?.value).toBeNull(); - await other.setName(undefined); - await session.setLabel(ROOT_ID, undefined); - await other.deleteValue(applicationState); - expect(await session.getName()).toBeUndefined(); - expect(await other.getLabel(ROOT_ID)).toBeUndefined(); - expect(await session.getValue(applicationState)).toBeUndefined(); - await session.close(); + await other.setName(undefined, BACKGROUND_CONTEXT); + await session.setLabel(ROOT_ID, undefined, BACKGROUND_CONTEXT); + await other.deleteValue(applicationState, BACKGROUND_CONTEXT); + expect(await session.getName(BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await other.getLabel(ROOT_ID, BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await session.getValue(applicationState, BACKGROUND_CONTEXT)).toBeUndefined(); + await session.close(BACKGROUND_CONTEXT); }); it("passes application values directly to storage", async () => { const session = await createTreeSession(); const value = { nested: ["original"] }; - await session.setValue(objectApplicationState, value); - expect((await session.getValue(objectApplicationState))?.value).toBe(value); - await session.close(); + await session.setValue(objectApplicationState, value, BACKGROUND_CONTEXT); + expect((await session.getValue(objectApplicationState, BACKGROUND_CONTEXT))?.value).toBe(value); + await session.close(BACKGROUND_CONTEXT); }); it("exposes bound application values and lists through every lane view", async () => { @@ -192,11 +195,13 @@ describe("StorageBackedSession SessionTree", () => { const workspaceEvents = events("workspace"); const other = session.view("main"); - await session.setValue(scalar, { count: 1 }); - await other.appendList(workspaceEvents, "first"); - await session.appendList(workspaceEvents, "second"); - expect((await other.getValue(scalar))?.value).toEqual({ count: 1 }); - expect((await session.readList(workspaceEvents, { limit: 1 })).map(({ value }) => value)).toEqual(["first"]); + await session.setValue(scalar, { count: 1 }, BACKGROUND_CONTEXT); + await other.appendList(workspaceEvents, "first", BACKGROUND_CONTEXT); + await session.appendList(workspaceEvents, "second", BACKGROUND_CONTEXT); + expect((await other.getValue(scalar, BACKGROUND_CONTEXT))?.value).toEqual({ count: 1 }); + expect( + (await session.readList(workspaceEvents, { limit: 1 }, BACKGROUND_CONTEXT)).map(({ value }) => value), + ).toEqual(["first"]); expect(storage.getCommitAttempts()).toHaveLength(3); expect(storage.getCommitAttempts().map((writes) => writes.map(({ kind }) => kind))).toEqual([ ["value"], @@ -204,72 +209,77 @@ describe("StorageBackedSession SessionTree", () => { ["list"], ]); - await other.deleteList(workspaceEvents); - await session.deleteValue(scalar); - expect(await session.readList(workspaceEvents)).toEqual([]); - expect(await session.getValue(scalar)).toBeUndefined(); - await session.close(); + await other.deleteList(workspaceEvents, BACKGROUND_CONTEXT); + await session.deleteValue(scalar, BACKGROUND_CONTEXT); + expect(await session.readList(workspaceEvents, undefined, BACKGROUND_CONTEXT)).toEqual([]); + expect(await session.getValue(scalar, BACKGROUND_CONTEXT)).toBeUndefined(); + await session.close(BACKGROUND_CONTEXT); }); it("applies global query ordering, filters, exclusive cursors, and limits", async () => { const session = await createTreeSession(); - const all = await session.findEntries(); + const all = await session.findEntries(undefined, BACKGROUND_CONTEXT); const custom = all.find((entry) => entry.id === CUSTOM_ID); if (custom === undefined) throw new Error("Expected custom entry"); expect(all.map((entry) => entry.id)).toEqual([OTHER_ID, CUSTOM_ID, CHILD_ID, ROOT_ID]); - expect((await session.findEntries({ order: "asc", type: "custom" })).map((entry) => entry.id)).toEqual([ - ROOT_ID, - CUSTOM_ID, - OTHER_ID, - ]); - expect((await session.findEntries({ customType: "note", limit: 1 })).map((entry) => entry.id)).toEqual([ - CUSTOM_ID, - ]); - expect((await session.findEntries({ cursor: { seq: custom.seq } })).map((entry) => entry.id)).toEqual([ - CHILD_ID, - ROOT_ID, - ]); expect( - (await session.findEntries({ order: "asc", cursor: { seq: custom.seq } })).map((entry) => entry.id), + (await session.findEntries({ order: "asc", type: "custom" }, BACKGROUND_CONTEXT)).map((entry) => entry.id), + ).toEqual([ROOT_ID, CUSTOM_ID, OTHER_ID]); + expect( + (await session.findEntries({ customType: "note", limit: 1 }, BACKGROUND_CONTEXT)).map((entry) => entry.id), + ).toEqual([CUSTOM_ID]); + expect( + (await session.findEntries({ cursor: { seq: custom.seq } }, BACKGROUND_CONTEXT)).map((entry) => entry.id), + ).toEqual([CHILD_ID, ROOT_ID]); + expect( + (await session.findEntries({ order: "asc", cursor: { seq: custom.seq } }, BACKGROUND_CONTEXT)).map( + (entry) => entry.id, + ), ).toEqual([OTHER_ID]); - expect((await session.findEntry({ type: "message" }))?.id).toBe(CHILD_ID); - expect(await session.findEntry({ customType: "missing" })).toBeUndefined(); - await session.close(); + expect((await session.findEntry({ type: "message" }, BACKGROUND_CONTEXT))?.id).toBe(CHILD_ID); + expect(await session.findEntry({ customType: "missing" }, BACKGROUND_CONTEXT)).toBeUndefined(); + await session.close(BACKGROUND_CONTEXT); }); it("applies branch defaults, inclusive stops, filters, cursors, and lane leaves", async () => { const session = await createTreeSession(); - const path = await session.findEntriesOnBranch(); + const path = await session.findEntriesOnBranch(undefined, BACKGROUND_CONTEXT); const child = path.find((entry) => entry.id === CHILD_ID); if (child === undefined) throw new Error("Expected child entry"); expect(path.map((entry) => entry.id)).toEqual([CUSTOM_ID, CHILD_ID, ROOT_ID]); - expect((await session.findEntriesOnBranch({ stopAtId: CHILD_ID })).map((entry) => entry.id)).toEqual([ - CUSTOM_ID, - CHILD_ID, - ]); - expect((await session.view("other").findEntriesOnBranch()).map((entry) => entry.id)).toEqual([OTHER_ID, ROOT_ID]); expect( - (await session.findEntriesOnBranch({ stopAtType: "message", type: "custom" })).map((entry) => entry.id), - ).toEqual([CUSTOM_ID]); + (await session.findEntriesOnBranch({ stopAtId: CHILD_ID }, BACKGROUND_CONTEXT)).map((entry) => entry.id), + ).toEqual([CUSTOM_ID, CHILD_ID]); expect( - (await session.findEntriesOnBranch({ order: "oldestFirst", cursor: { seq: child.seq }, limit: 1 })).map( + (await session.view("other").findEntriesOnBranch(undefined, BACKGROUND_CONTEXT)).map((entry) => entry.id), + ).toEqual([OTHER_ID, ROOT_ID]); + expect( + (await session.findEntriesOnBranch({ stopAtType: "message", type: "custom" }, BACKGROUND_CONTEXT)).map( (entry) => entry.id, ), ).toEqual([CUSTOM_ID]); - expect((await session.findEntryOnBranch({ customType: "root" }))?.id).toBe(ROOT_ID); - expect(await session.findEntryOnBranch({ customType: "missing" })).toBeUndefined(); - await session.close(); + expect( + ( + await session.findEntriesOnBranch( + { order: "oldestFirst", cursor: { seq: child.seq }, limit: 1 }, + BACKGROUND_CONTEXT, + ) + ).map((entry) => entry.id), + ).toEqual([CUSTOM_ID]); + expect((await session.findEntryOnBranch({ customType: "root" }, BACKGROUND_CONTEXT))?.id).toBe(ROOT_ID); + expect(await session.findEntryOnBranch({ customType: "missing" }, BACKGROUND_CONTEXT)).toBeUndefined(); + await session.close(BACKGROUND_CONTEXT); }); it("returns an empty default branch for a lane at the root", async () => { const session = new StorageBackedSession(metadata, new MemoryStorage({ now: () => NOW })); await commitSession(session, [storedValues.setValue(storedValues.laneLeaf("main"), null)]); - expect(await session.findEntriesOnBranch()).toEqual([]); - expect(await session.findEntryOnBranch()).toBeUndefined(); - await session.close(); + expect(await session.findEntriesOnBranch(undefined, BACKGROUND_CONTEXT)).toEqual([]); + expect(await session.findEntryOnBranch(undefined, BACKGROUND_CONTEXT)).toBeUndefined(); + await session.close(BACKGROUND_CONTEXT); }); it("atomically appends messages and custom entries to the bound lane", async () => { @@ -282,9 +292,12 @@ describe("StorageBackedSession SessionTree", () => { ]); storage.clearCommitAttempts(); - const messageId = await session.appendMessage({ role: "user", content: "hello", timestamp: NOW }); - const customId = await session.appendCustomEntry("note", { nested: ["value"] }); - const withoutDataId = await session.view("other").appendCustomEntry("marker"); + const messageId = await session.appendMessage( + { role: "user", content: "hello", timestamp: NOW }, + BACKGROUND_CONTEXT, + ); + const customId = await session.appendCustomEntry("note", { nested: ["value"] }, BACKGROUND_CONTEXT); + const withoutDataId = await session.view("other").appendCustomEntry("marker", undefined, BACKGROUND_CONTEXT); expect(storage.getCommitAttempts()).toEqual([ [ @@ -316,10 +329,10 @@ describe("StorageBackedSession SessionTree", () => { storedValues.setValue(storedValues.laneLeaf("other"), withoutDataId), ], ]); - expect(await session.getLeafId()).toBe(customId); - expect(await session.view("other").getLeafId()).toBe(withoutDataId); - expect(await session.getEntry(messageId)).toMatchObject({ parentId: null, type: "message" }); - expect(await session.getEntry(customId)).toEqual( + expect(await session.getLeafId(BACKGROUND_CONTEXT)).toBe(customId); + expect(await session.view("other").getLeafId(BACKGROUND_CONTEXT)).toBe(withoutDataId); + expect(await session.getEntry(messageId, BACKGROUND_CONTEXT)).toMatchObject({ parentId: null, type: "message" }); + expect(await session.getEntry(customId, BACKGROUND_CONTEXT)).toEqual( expect.objectContaining({ parentId: messageId, type: "custom", @@ -327,10 +340,10 @@ describe("StorageBackedSession SessionTree", () => { data: { nested: ["value"] }, }), ); - const withoutData = await session.getEntry(withoutDataId); + const withoutData = await session.getEntry(withoutDataId, BACKGROUND_CONTEXT); expect(withoutData).toEqual(expect.objectContaining({ parentId: null, type: "custom", customType: "marker" })); expect(withoutData).not.toHaveProperty("data"); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("defers appends into an active run without moving the lane leaf", async () => { @@ -345,23 +358,23 @@ describe("StorageBackedSession SessionTree", () => { storage.clearCommitAttempts(); const [customId, messageId] = await Promise.all([ - session.appendCustomEntry("note", { deferred: true }), - session.appendMessage({ role: "user", content: "later", timestamp: NOW }), + session.appendCustomEntry("note", { deferred: true }, BACKGROUND_CONTEXT), + session.appendMessage({ role: "user", content: "later", timestamp: NOW }, BACKGROUND_CONTEXT), ]); - expect(await session.getLeafId()).toBe(ROOT_ID); - expect(await session.getEntries([customId, messageId])).toEqual(new Map()); - expect(await session.getValue(storedValues.pendingEntry(customId))).toMatchObject({ + expect(await session.getLeafId(BACKGROUND_CONTEXT)).toBe(ROOT_ID); + expect(await session.getEntries([customId, messageId], BACKGROUND_CONTEXT)).toEqual(new Map()); + expect(await session.getValue(storedValues.pendingEntry(customId), BACKGROUND_CONTEXT)).toMatchObject({ value: { type: "custom", customType: "note", payload: { deferred: true } }, }); - expect(await session.getValue(storedValues.pendingEntry(messageId))).toMatchObject({ + expect(await session.getValue(storedValues.pendingEntry(messageId), BACKGROUND_CONTEXT)).toMatchObject({ value: { type: "message", payload: { role: "user", content: "later", timestamp: NOW } }, }); - expect(await session.getValue(storedValues.operationState(OPERATION_ID))).toMatchObject({ + expect(await session.getValue(storedValues.operationState(OPERATION_ID), BACKGROUND_CONTEXT)).toMatchObject({ value: { inbox: { writes: [customId, messageId] } }, }); expect(storage.getCommitAttempts()).toHaveLength(2); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("rejects without writing while structural-operation waiting is not implemented", async () => { @@ -383,11 +396,11 @@ describe("StorageBackedSession SessionTree", () => { ]); storage.clearCommitAttempts(); - const append = session.appendCustomEntry("during-structural"); + const append = session.appendCustomEntry("during-structural", undefined, BACKGROUND_CONTEXT); await expect(append).rejects.toThrow("Cannot append while structural operation"); expect(storage.getCommitAttempts()).toEqual([]); - expect(await session.getLeafId()).toBe(ROOT_ID); - await session.close(); + expect(await session.getLeafId(BACKGROUND_CONTEXT)).toBe(ROOT_ID); + await session.close(BACKGROUND_CONTEXT); }); it("diverges lanes that append from the same shared leaf", async () => { @@ -398,18 +411,24 @@ describe("StorageBackedSession SessionTree", () => { thinkingLevel: "off" as const, activeToolNames: [], }; - const review = await session.createLane("review", ROOT_ID, configuration); + const review = await session.createLane("review", ROOT_ID, configuration, BACKGROUND_CONTEXT); const [mainId, reviewId] = await Promise.all([ - session.appendCustomEntry("main-child"), - review.appendCustomEntry("review-child"), + session.appendCustomEntry("main-child", undefined, BACKGROUND_CONTEXT), + review.appendCustomEntry("review-child", undefined, BACKGROUND_CONTEXT), ]); - expect((await session.findEntriesOnBranch()).map((entry) => entry.id)).toEqual([mainId, ROOT_ID]); - expect((await review.findEntriesOnBranch()).map((entry) => entry.id)).toEqual([reviewId, ROOT_ID]); - expect((await session.getEntry(mainId))?.parentId).toBe(ROOT_ID); - expect((await session.getEntry(reviewId))?.parentId).toBe(ROOT_ID); - await session.close(); + expect((await session.findEntriesOnBranch(undefined, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ + mainId, + ROOT_ID, + ]); + expect((await review.findEntriesOnBranch(undefined, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ + reviewId, + ROOT_ID, + ]); + expect((await session.getEntry(mainId, BACKGROUND_CONTEXT))?.parentId).toBe(ROOT_ID); + expect((await session.getEntry(reviewId, BACKGROUND_CONTEXT))?.parentId).toBe(ROOT_ID); + await session.close(BACKGROUND_CONTEXT); }); it("serializes concurrent appends into one linear lane branch", async () => { @@ -417,12 +436,16 @@ describe("StorageBackedSession SessionTree", () => { await commitSession(session, [sessionWrites.insertEntry(customEntry(ROOT_ID, null)), ...laneWrites(ROOT_ID)]); const [firstId, secondId] = await Promise.all([ - session.appendCustomEntry("first"), - session.appendCustomEntry("second"), + session.appendCustomEntry("first", undefined, BACKGROUND_CONTEXT), + session.appendCustomEntry("second", undefined, BACKGROUND_CONTEXT), ]); - expect((await session.findEntriesOnBranch()).map((entry) => entry.id)).toEqual([secondId, firstId, ROOT_ID]); - await session.close(); + expect((await session.findEntriesOnBranch(undefined, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ + secondId, + firstId, + ROOT_ID, + ]); + await session.close(BACKGROUND_CONTEXT); }); it("passes append payloads directly to storage", async () => { @@ -434,18 +457,18 @@ describe("StorageBackedSession SessionTree", () => { const data = { nested: ["original"] }; const runCountBeforeAppends = laneMutationLine.runCount; - const id = await session.appendCustomEntry("note", data); - const entry = await session.getEntry(id); + const id = await session.appendCustomEntry("note", data, BACKGROUND_CONTEXT); + const entry = await session.getEntry(id, BACKGROUND_CONTEXT); if (entry?.type !== "custom") throw new Error("Expected custom entry"); expect(entry.data).toBe(data); expect(storage.getCommitAttempts()).toHaveLength(1); const message = { role: "user" as const, content: "original", timestamp: NOW }; - const messageId = await session.appendMessage(message); - const messageEntry = await session.getEntry(messageId); + const messageId = await session.appendMessage(message, BACKGROUND_CONTEXT); + const messageEntry = await session.getEntry(messageId, BACKGROUND_CONTEXT); if (messageEntry?.type !== "message") throw new Error("Expected message entry"); expect(messageEntry.message).toBe(message); expect(laneMutationLine.runCount - runCountBeforeAppends).toBe(2); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("fails fast when active operation registers are inconsistent", async () => { @@ -454,11 +477,11 @@ describe("StorageBackedSession SessionTree", () => { await commitSession(session, [...laneWrites(null, { currentOperationId: OPERATION_ID, pendingNextRun: [] })]); storage.clearCommitAttempts(); - const append = session.appendCustomEntry("note"); + const append = session.appendCustomEntry("note", undefined, BACKGROUND_CONTEXT); await expect(append).rejects.toBeInstanceOf(SessionInvariantError); await expect(append).rejects.toThrow("missing op.meta"); expect(storage.getCommitAttempts()).toEqual([]); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("propagates append storage failures without moving the leaf", async () => { @@ -468,11 +491,11 @@ describe("StorageBackedSession SessionTree", () => { const rejection = new Error("commit failed"); storage.rejection = rejection; - await expect(session.appendCustomEntry("note")).rejects.toBe(rejection); + await expect(session.appendCustomEntry("note", undefined, BACKGROUND_CONTEXT)).rejects.toBe(rejection); storage.rejection = undefined; - expect(await session.getLeafId()).toBe(ROOT_ID); - expect((await session.findEntries()).map((entry) => entry.id)).toEqual([ROOT_ID]); - await session.close(); + expect(await session.getLeafId(BACKGROUND_CONTEXT)).toBe(ROOT_ID); + expect((await session.findEntries(undefined, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ROOT_ID]); + await session.close(BACKGROUND_CONTEXT); }); it("drains an append whose storage commit was admitted before close", async () => { @@ -481,9 +504,9 @@ describe("StorageBackedSession SessionTree", () => { await commitSession(session, [...laneWrites(null)]); storage.block = true; - const append = session.appendCustomEntry("admitted"); + const append = session.appendCustomEntry("admitted", undefined, BACKGROUND_CONTEXT); while (!storage.admitted) await Promise.resolve(); - const close = session.close(); + const close = session.close(BACKGROUND_CONTEXT); storage.release(); const id = await append; @@ -494,17 +517,17 @@ describe("StorageBackedSession SessionTree", () => { it("rejects all SessionTree operations after close, including existing views", async () => { const session = await createTreeSession(); const view = session.view("other"); - await session.close(); + await session.close(BACKGROUND_CONTEXT); const operations: Array<() => Promise> = [ - () => session.getLeafId(), - () => view.getEntry(ROOT_ID), - () => session.getStats(), - () => view.getName(), - () => session.setName("closed"), - () => view.findEntries(), - () => session.findEntriesOnBranch(), - () => view.appendCustomEntry("closed"), + () => session.getLeafId(BACKGROUND_CONTEXT), + () => view.getEntry(ROOT_ID, BACKGROUND_CONTEXT), + () => session.getStats(BACKGROUND_CONTEXT), + () => view.getName(BACKGROUND_CONTEXT), + () => session.setName("closed", BACKGROUND_CONTEXT), + () => view.findEntries(undefined, BACKGROUND_CONTEXT), + () => session.findEntriesOnBranch(undefined, BACKGROUND_CONTEXT), + () => view.appendCustomEntry("closed", undefined, BACKGROUND_CONTEXT), ]; for (const operation of operations) await expect(operation()).rejects.toThrow("Session is closed"); }); diff --git a/packages/agent/test/harness/storage-backed-session.test.ts b/packages/agent/test/harness/storage-backed-session.test.ts index 7933fba7c..27d821786 100644 --- a/packages/agent/test/harness/storage-backed-session.test.ts +++ b/packages/agent/test/harness/storage-backed-session.test.ts @@ -1,5 +1,6 @@ import type { AssistantMessage } from "@earendil-works/pi-ai"; import { describe, expect, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import type { CustomMessage } from "../../src/harness/messages.ts"; import * as sessionWrites from "../../src/harness/session/commit.ts"; import { MemoryStorage } from "../../src/harness/session/memory.ts"; @@ -25,7 +26,7 @@ const metadata = { } satisfies SessionMetadata; function commitSession(session: Session, transaction: Write[]) { - return session.mutate("main", (mutator) => mutator.commit(transaction)); + return session.mutate("main", (mutator) => mutator.commit(transaction, BACKGROUND_CONTEXT), BACKGROUND_CONTEXT); } describe("StorageBackedSession", () => { @@ -41,12 +42,14 @@ describe("StorageBackedSession", () => { const result = await commitSession(session, transaction); expect(storage.getCommitAttempts()[0]).toBe(transaction); - const entry = (await session.getEntries([ENTRY_ID])).get(ENTRY_ID); + const entry = (await session.getEntries([ENTRY_ID], BACKGROUND_CONTEXT)).get(ENTRY_ID); expect(entry).toMatchObject({ seq: result.seqs[0], timestamp: NOW }); if (entry?.type !== "custom") throw new Error("Expected custom entry"); expect(entry.data).toBe(data); - expect((await session.getValue(storedValues.value("test.value", "state")))?.value).toBe(data); - await session.close(); + expect( + (await session.getValue(storedValues.value("test.value", "state"), BACKGROUND_CONTEXT))?.value, + ).toBe(data); + await session.close(BACKGROUND_CONTEXT); }); it("composes bound values and lists atomically with entries and usage", async () => { @@ -55,31 +58,39 @@ describe("StorageBackedSession", () => { const scalar = storedValues.value("test.application.scalar"); const events = storedValues.list("test.application.events"); - const result = await session.mutate("main", (mutator) => - mutator.commit([ - sessionWrites.insertEntry({ id: ENTRY_ID, parentId: null, type: "custom", customType: "note" }), - storedValues.setValue(scalar, "state"), - storedValues.appendList(events, "event"), - sessionWrites.insertUsage({ - id: "usage", - adjustment: false, - usage: { - input: 1, - output: 1, - cacheRead: 0, - cacheWrite: 0, - totalTokens: 2, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - }, - }), - ]), + const result = await session.mutate( + "main", + (mutator) => + mutator.commit( + [ + sessionWrites.insertEntry({ id: ENTRY_ID, parentId: null, type: "custom", customType: "note" }), + storedValues.setValue(scalar, "state"), + storedValues.appendList(events, "event"), + sessionWrites.insertUsage({ + id: "usage", + adjustment: false, + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); expect(result.seqs).toHaveLength(4); - expect((await session.getValue(scalar))?.seq).toBe(result.seqs[1]); - expect(await session.readList(events)).toEqual([{ seq: result.seqs[2], value: "event" }]); + expect((await session.getValue(scalar, BACKGROUND_CONTEXT))?.seq).toBe(result.seqs[1]); + expect(await session.readList(events, undefined, BACKGROUND_CONTEXT)).toEqual([ + { seq: result.seqs[2], value: "event" }, + ]); expect(storage.getCommitAttempts()).toHaveLength(1); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("rejects pending assistant entries at the durable session write boundary", async () => { @@ -109,8 +120,8 @@ describe("StorageBackedSession", () => { ]), ).rejects.toThrow("Cannot persist a pending assistant message"); expect(storage.getCommitAttempts()).toEqual([]); - expect(await session.getEntries([ENTRY_ID])).toEqual(new Map()); - await session.close(); + expect(await session.getEntries([ENTRY_ID], BACKGROUND_CONTEXT)).toEqual(new Map()); + await session.close(BACKGROUND_CONTEXT); }); it("trusts typed custom messages without repository schema registration", async () => { @@ -127,12 +138,12 @@ describe("StorageBackedSession", () => { const result = await commitSession(session, [sessionWrites.insertEntry(entry)]); - expect((await session.getEntries([ENTRY_ID])).get(ENTRY_ID)).toEqual({ + expect((await session.getEntries([ENTRY_ID], BACKGROUND_CONTEXT)).get(ENTRY_ID)).toEqual({ ...entry, seq: result.firstSeq, timestamp: result.timestamp, }); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("serializes mutations, permits one commit attempt, and invalidates the mutator", async () => { @@ -140,20 +151,24 @@ describe("StorageBackedSession", () => { const session = new StorageBackedSession(metadata, storage); let captured: SessionMutator | undefined; - await session.mutate("review", async (mutator) => { - captured = mutator; - expect(mutator.lane).toBe("review"); - expect(await mutator.getValue(storedValues.sessionName)).toBeUndefined(); - await mutator.commit([storedValues.setValue(storedValues.sessionName, "committed")]); - await expect(mutator.commit([])).rejects.toThrow("commit already attempted"); - }); + await session.mutate( + "review", + async (mutator) => { + captured = mutator; + expect(mutator.lane).toBe("review"); + expect(await mutator.getValue(storedValues.sessionName, BACKGROUND_CONTEXT)).toBeUndefined(); + await mutator.commit([storedValues.setValue(storedValues.sessionName, "committed")], BACKGROUND_CONTEXT); + await expect(mutator.commit([], BACKGROUND_CONTEXT)).rejects.toThrow("commit already attempted"); + }, + BACKGROUND_CONTEXT, + ); expect(storage.getCommitAttempts()).toHaveLength(1); - expect(await session.getName()).toBe("committed"); + expect(await session.getName(BACKGROUND_CONTEXT)).toBe("committed"); const invalidated = captured; if (invalidated === undefined) throw new Error("Expected captured mutator"); - expect(() => invalidated.getEntries([])).toThrow("outside its mutation callback"); - await session.close(); + expect(() => invalidated.getEntries([], BACKGROUND_CONTEXT)).toThrow("outside its mutation callback"); + await session.close(BACKGROUND_CONTEXT); }); it("consumes the commit guard when the first commit fails", async () => { @@ -164,13 +179,17 @@ describe("StorageBackedSession", () => { ] satisfies Write[]; await expect( - session.mutate("main", async (mutator) => { - await expect(mutator.commit(transaction)).rejects.toThrow("Missing parent entry"); - await expect(mutator.commit([])).rejects.toThrow("commit already attempted"); - }), + session.mutate( + "main", + async (mutator) => { + await expect(mutator.commit(transaction, BACKGROUND_CONTEXT)).rejects.toThrow("Missing parent entry"); + await expect(mutator.commit([], BACKGROUND_CONTEXT)).rejects.toThrow("commit already attempted"); + }, + BACKGROUND_CONTEXT, + ), ).resolves.toBeUndefined(); expect(storage.getCommitAttempts()).toHaveLength(1); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("mints distinct follower ids with the leader timestamp", async () => { @@ -182,7 +201,7 @@ describe("StorageBackedSession", () => { expect([leader, ...followers].map(decodeTimestamp)).toEqual([leaderTimestamp, leaderTimestamp, leaderTimestamp]); expect(new Set([leader, ...followers])).toHaveLength(3); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("exposes metadata directly and the shared UUIDv7 id generator", async () => { @@ -193,17 +212,19 @@ describe("StorageBackedSession", () => { expect(session.idGenerator.next()).toMatch( /^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/, ); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); it("closes idempotently and rejects operations not admitted before close", async () => { const storage = new MemoryStorage({ now: () => NOW }); const session = new StorageBackedSession(metadata, storage); - await Promise.all([session.close(), session.close()]); - await expect(session.mutate("main", () => undefined)).rejects.toThrow("Session is closed"); - await expect(session.getEntries([])).rejects.toThrow("Session is closed"); - await expect(session.getValue(storedValues.sessionName)).rejects.toThrow("Session is closed"); - await expect(session.scanValues(storedValues.sessionName)).rejects.toThrow("Session is closed"); + await Promise.all([session.close(BACKGROUND_CONTEXT), session.close(BACKGROUND_CONTEXT)]); + await expect(session.mutate("main", () => undefined, BACKGROUND_CONTEXT)).rejects.toThrow("Session is closed"); + await expect(session.getEntries([], BACKGROUND_CONTEXT)).rejects.toThrow("Session is closed"); + await expect(session.getValue(storedValues.sessionName, BACKGROUND_CONTEXT)).rejects.toThrow("Session is closed"); + await expect(session.scanValues(storedValues.sessionName, BACKGROUND_CONTEXT)).rejects.toThrow( + "Session is closed", + ); }); }); diff --git a/packages/agent/test/harness/types.test.ts b/packages/agent/test/harness/types.test.ts index da691c643..4d4c029f6 100644 --- a/packages/agent/test/harness/types.test.ts +++ b/packages/agent/test/harness/types.test.ts @@ -13,6 +13,7 @@ import type { CancelQueuedResult, CheckpointPhase, CompactionState, + Context, Control, CustomEntry, Deferred, @@ -344,18 +345,30 @@ it("covers storage, session, repository, search, and identity signatures", () => expectTypeOf[0]["start"]>().toEqualTypeOf(); expectTypeOf().toEqualTypeOf(); expectTypeOf().toEqualTypeOf< - (transactionToCommit: Write[]) => Promise<{ firstSeq: number; seqs: number[]; timestamp: number }> + ( + transactionToCommit: Write[], + context: Context, + ) => Promise<{ firstSeq: number; seqs: number[]; timestamp: number }> >(); expectTypeOf().toEqualTypeOf< - (lane: string, mutation: (mutator: SessionMutator) => T | Promise) => Promise + ( + lane: string, + mutation: (mutator: SessionMutator, context: Context) => T | Promise, + context: Context, + ) => Promise >(); expectTypeOf().toEqualTypeOf< - (transactionToCommit: Write[]) => Promise<{ firstSeq: number; seqs: number[]; timestamp: number }> + ( + transactionToCommit: Write[], + context: Context, + ) => Promise<{ firstSeq: number; seqs: number[]; timestamp: number }> >(); expectTypeOf().toEqualTypeOf< - (name: string, at: string | null, laneConfiguration: LaneConfiguration) => Promise + (name: string, at: string | null, laneConfiguration: LaneConfiguration, context: Context) => Promise + >(); + expectTypeOf().toEqualTypeOf< + (options: SessionCreateOptions, context: Context) => Promise >(); - expectTypeOf().toEqualTypeOf<(options: SessionCreateOptions) => Promise>(); expectTypeOf().toEqualTypeOf< (query: SearchQuery) => Promise >(); @@ -435,7 +448,7 @@ it("covers Part 5 results, events, hooks, snapshots, tools, and stream options", boolean | { window?: "15m" | "1h" | "24h" } | undefined >(); expectTypeOf().toEqualTypeOf< - (entry: CustomEntry) => AgentMessage[] | undefined | Promise + (entry: CustomEntry, context: Context) => AgentMessage[] | undefined | Promise >(); expectTypeOf>().toEqualTypeOf>(); diff --git a/packages/agent/test/harness/values.test.ts b/packages/agent/test/harness/values.test.ts index 075955730..f4cc8ff6b 100644 --- a/packages/agent/test/harness/values.test.ts +++ b/packages/agent/test/harness/values.test.ts @@ -1,5 +1,6 @@ import type { AssistantMessageFrame } from "@earendil-works/pi-ai"; import { describe, expect, expectTypeOf, it } from "vitest"; +import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; import { MemoryStorage } from "../../src/harness/session/memory.ts"; import type { DurableStructuralPreparation, @@ -90,22 +91,22 @@ describe("bound value addresses", () => { const storage = new MemoryStorage({ now: () => 1 }); const events = list<{ name: string }>("app.events"); const reader: SessionReader = storage; - const storageElements = await storage.readList(events); - const readerElements = await reader.readList(events); + const storageElements = await storage.readList(events, undefined, BACKGROUND_CONTEXT); + const readerElements = await reader.readList(events, undefined, BACKGROUND_CONTEXT); expectTypeOf(storageElements).toEqualTypeOf[]>(); expectTypeOf(readerElements).toEqualTypeOf[]>(); - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); }); it("uses separately constructed equal addresses for one durable location", async () => { const storage = new MemoryStorage({ now: () => 1 }); const first = value<{ ready: boolean }>("app.state", "workspace"); const second = value<{ ready: boolean }>("app.state", "workspace"); - await storage.commit([setValue(first, { ready: true })]); - const stored = await storage.getValue(second); + await storage.commit([setValue(first, { ready: true })], BACKGROUND_CONTEXT); + const stored = await storage.getValue(second, BACKGROUND_CONTEXT); expect(stored).toEqual({ address: first, value: { ready: true }, seq: 1 }); expectTypeOf(stored).toEqualTypeOf | undefined>(); - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); }); }); diff --git a/packages/coding-agent/src/experimental/server.ts b/packages/coding-agent/src/experimental/server.ts index 17b137502..44be4ec11 100644 --- a/packages/coding-agent/src/experimental/server.ts +++ b/packages/coding-agent/src/experimental/server.ts @@ -3,7 +3,7 @@ import { chmod, lstat, mkdir, readFile, writeFile } from "node:fs/promises"; import { homedir } from "node:os"; import { isAbsolute, join, resolve } from "node:path"; import { fileURLToPath } from "node:url"; -import { type JsonlSessionMetadata, JsonlSessionRepo } from "@earendil-works/pi-agent-core"; +import { type JsonlSessionMetadata, JsonlSessionRepo, TODO_CONTEXT } from "@earendil-works/pi-agent-core"; import { NodeExecutionEnv } from "@earendil-works/pi-agent-core/node"; import { PiClient, PiDisconnectedError, PiServerError } from "@earendil-works/pi-client"; import { createUnixTransportFactory, type UnixServerRoute } from "@earendil-works/pi-client/unix"; @@ -332,16 +332,18 @@ async function startServerBackend( const host: PiServerHost = { sessions: { list: async () => { - const sessions = new Map((await repo.list()).map((metadata) => [metadata.path, metadata])); + const sessions = new Map( + (await repo.list(undefined, TODO_CONTEXT)).map((metadata) => [metadata.path, metadata]), + ); for (const metadata of workers.trackedSessions) sessions.set(metadata.path, metadata); return [...sessions.values()]; }, create: async (createOptions) => { - const session = await repo.create(createOptions); + const session = await repo.create(createOptions, TODO_CONTEXT); try { return session.metadata; } finally { - await session.close(); + await session.close(TODO_CONTEXT); } }, }, @@ -349,7 +351,7 @@ async function startServerBackend( }; const socketPath = options.path; const closeCatalog = async (): Promise => { - const cleanup = await Promise.allSettled([repo.close(), executionEnv.cleanup()]); + const cleanup = await Promise.allSettled([repo.close(TODO_CONTEXT), executionEnv.cleanup()]); const errors = cleanup.flatMap((result) => (result.status === "rejected" ? [result.reason] : [])); if (errors.length === 1) throw errors[0]; if (errors.length > 1) throw new AggregateError(errors, "Experimental session catalog cleanup failed"); diff --git a/packages/coding-agent/src/experimental/session-worker.ts b/packages/coding-agent/src/experimental/session-worker.ts index 27219354e..598344d37 100644 --- a/packages/coding-agent/src/experimental/session-worker.ts +++ b/packages/coding-agent/src/experimental/session-worker.ts @@ -11,6 +11,7 @@ import { type JsonlSessionMetadata, JsonlSessionRepo, type Session, + TODO_CONTEXT, } from "@earendil-works/pi-agent-core"; import { NodeExecutionEnv } from "@earendil-works/pi-agent-core/node"; import { @@ -469,13 +470,13 @@ async function closeResources(resources: { }): Promise { const errors: unknown[] = []; try { - if (resources.harness) await resources.harness.close(); - else await resources.session?.close(); + if (resources.harness) await resources.harness.close(TODO_CONTEXT); + else await resources.session?.close(TODO_CONTEXT); } catch (error) { errors.push(error); } try { - await resources.repo.close(); + await resources.repo.close(TODO_CONTEXT); } catch (error) { errors.push(error); } @@ -517,7 +518,7 @@ async function run(options: SessionWorkerOptions, createHarness: CreateSessionWo let session: Session | undefined; let harness: AgentHarnessInstance | undefined; try { - session = await repo.open(metadata); + session = await repo.open(metadata, TODO_CONTEXT); harness = await createHarness(session, options, executionEnv); } catch (error) { try { @@ -597,11 +598,13 @@ async function run(options: SessionWorkerOptions, createHarness: CreateSessionWo prompt: async (_scope: WorkerOperationScope, prompt) => { const args = toHarnessPromptArguments(prompt); const result = - typeof args[0] === "string" ? await harness.prompt(args[0], args[1]) : await harness.prompt(args[0]); + typeof args[0] === "string" + ? await harness.prompt(args[0], args[1], TODO_CONTEXT) + : await harness.prompt(args[0], TODO_CONTEXT); return toWireRunResult(result); }, watch: async (scope: WorkerOperationScope, ..._args: never[]) => { - const handle = await harness.watch(); + const handle = await harness.watch(TODO_CONTEXT); const watchId = randomUUID(); laneWatches.set(watchId, { scope, handle }); return { watchId, snapshot: toWireLaneSnapshot(handle.snapshot) }; @@ -790,42 +793,48 @@ async function createCodingAgentHarness( const tools = [createReadTool(), createWriteTool(), createBashTool()]; const activeToolNames = tools.map((tool) => tool.name); const harness = ( - await AgentHarness.create({ - session, - models: modelRuntime, - model: resolved.model, - thinkingLevel: resolved.thinkingLevel, - tools, - activeToolNames, - toolContext: { env: executionEnv }, - resources: {}, - }) + await AgentHarness.create( + { + session, + models: modelRuntime, + model: resolved.model, + thinkingLevel: resolved.thinkingLevel, + tools, + activeToolNames, + toolContext: { env: executionEnv }, + resources: {}, + }, + TODO_CONTEXT, + ) ).harness; try { - const currentActiveToolNames = await harness.getActiveTools(); + const currentActiveToolNames = await harness.getActiveTools(TODO_CONTEXT); if ( currentActiveToolNames.length !== activeToolNames.length || currentActiveToolNames.some((name, index) => name !== activeToolNames[index]) ) { - await harness.setActiveTools(activeToolNames); + await harness.setActiveTools(activeToolNames, TODO_CONTEXT); } if (options.model !== undefined) { - const currentModel = await harness.getModel(); + const currentModel = await harness.getModel(TODO_CONTEXT); if ( !currentModel || currentModel.provider !== resolved.model.provider || currentModel.id !== resolved.model.id ) { - await harness.setModel(resolved.model); + await harness.setModel(resolved.model, TODO_CONTEXT); } - if (resolved.thinkingLevel !== undefined && (await harness.getThinkingLevel()) !== resolved.thinkingLevel) { - await harness.setThinkingLevel(resolved.thinkingLevel); + if ( + resolved.thinkingLevel !== undefined && + (await harness.getThinkingLevel(TODO_CONTEXT)) !== resolved.thinkingLevel + ) { + await harness.setThinkingLevel(resolved.thinkingLevel, TODO_CONTEXT); } } return harness; } catch (error) { try { - await harness.close(); + await harness.close(TODO_CONTEXT); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "Session worker model selection and cleanup failed"); } diff --git a/packages/coding-agent/test/experimental-session-support.ts b/packages/coding-agent/test/experimental-session-support.ts index a80298666..315983cb1 100644 --- a/packages/coding-agent/test/experimental-session-support.ts +++ b/packages/coding-agent/test/experimental-session-support.ts @@ -1,6 +1,12 @@ import { mkdir, writeFile } from "node:fs/promises"; import { join } from "node:path"; -import { type Entry, type JsonlSessionMetadata, JsonlSessionRepo, laneConfig } from "@earendil-works/pi-agent-core"; +import { + BACKGROUND_CONTEXT, + type Entry, + type JsonlSessionMetadata, + JsonlSessionRepo, + laneConfig, +} from "@earendil-works/pi-agent-core"; import { NodeExecutionEnv } from "@earendil-works/pi-agent-core/node"; export async function createExperimentalSessions( @@ -13,13 +19,13 @@ export async function createExperimentalSessions( const metadata: JsonlSessionMetadata[] = []; try { for (const id of ids) { - const session = await repo.create({ id, cwd }); + const session = await repo.create({ id, cwd }, BACKGROUND_CONTEXT); metadata.push(session.metadata); - await session.close(); + await session.close(BACKGROUND_CONTEXT); } return metadata; } finally { - await repo.close(); + await repo.close(BACKGROUND_CONTEXT); await fileSystem.cleanup(); } } @@ -43,17 +49,17 @@ export async function readExperimentalSessionState( const repo = new JsonlSessionRepo({ fileSystem, sessionsRoot }); let session: Awaited> | undefined; try { - const matches = (await repo.list()).filter((metadata) => metadata.id === sessionId); + const matches = (await repo.list(undefined, BACKGROUND_CONTEXT)).filter((metadata) => metadata.id === sessionId); if (matches.length !== 1) throw new Error(`Expected one Session ${sessionId}, found ${matches.length}`); - session = await repo.open(matches[0]!); + session = await repo.open(matches[0]!, BACKGROUND_CONTEXT); const [branch, configuration] = await Promise.all([ - session.findEntriesOnBranch({ order: "oldestFirst" }), - session.getValue(laneConfig("main")), + session.findEntriesOnBranch({ order: "oldestFirst" }, BACKGROUND_CONTEXT), + session.getValue(laneConfig("main"), BACKGROUND_CONTEXT), ]); return { branch, model: configuration?.value.model, activeTools: configuration?.value.activeToolNames ?? [] }; } finally { - await session?.close(); - await repo.close(); + await session?.close(BACKGROUND_CONTEXT); + await repo.close(BACKGROUND_CONTEXT); await fileSystem.cleanup(); } } diff --git a/packages/coding-agent/test/fixtures/faux-session-worker.ts b/packages/coding-agent/test/fixtures/faux-session-worker.ts index fe96e3012..470b87e38 100644 --- a/packages/coding-agent/test/fixtures/faux-session-worker.ts +++ b/packages/coding-agent/test/fixtures/faux-session-worker.ts @@ -1,6 +1,6 @@ import { resolve } from "node:path"; import { fileURLToPath } from "node:url"; -import { AgentHarness } from "@earendil-works/pi-agent-core"; +import { AgentHarness, BACKGROUND_CONTEXT } from "@earendil-works/pi-agent-core"; import { createModels, fauxAssistantMessage, fauxProvider } from "@earendil-works/pi-ai"; import { consumeInternalProcessRole } from "../../src/experimental/process.ts"; import { runSessionWorkerWithHarness } from "../../src/experimental/session-worker.ts"; @@ -17,13 +17,16 @@ if (process.argv[1] !== undefined && resolve(process.argv[1]) === fileURLToPath( const models = createModels(); models.setProvider(faux.provider); return ( - await AgentHarness.create({ - session, - models, - model: faux.getModel(), - tools: [], - resources: {}, - }) + await AgentHarness.create( + { + session, + models, + model: faux.getModel(), + tools: [], + resources: {}, + }, + BACKGROUND_CONTEXT, + ) ).harness; }).catch((error: unknown) => { console.error(error); diff --git a/packages/server/src/testing/host.ts b/packages/server/src/testing/host.ts index afc3c3708..a40804238 100644 --- a/packages/server/src/testing/host.ts +++ b/packages/server/src/testing/host.ts @@ -1,5 +1,5 @@ import type { Session, SessionMetadata } from "@earendil-works/pi-agent-core"; -import { MemorySessionRepo } from "@earendil-works/pi-agent-core"; +import { BACKGROUND_CONTEXT, MemorySessionRepo } from "@earendil-works/pi-agent-core"; import type { LaneEvent, LaneSnapshot, PromptArguments, RunResult } from "@earendil-works/pi-protocol"; import type { HostedHarnessHandle, HostedHarnessWatch, PiServerHost } from "../types.ts"; @@ -162,13 +162,13 @@ export class TestHarness { this.failClose = undefined; throw error; } - await this.session.close(); + await this.session.close(BACKGROUND_CONTEXT); this.closed.resolve(undefined); this.#termination.resolve(undefined); } async terminate(error: Error): Promise { - await this.session.close(); + await this.session.close(BACKGROUND_CONTEXT); this.#termination.resolve(error); } @@ -204,14 +204,14 @@ export class TestServerHost implements PiServerHost { delay.entered.resolve(undefined); await delay.release.promise; } - return this.repo.list(); + return this.repo.list(undefined, BACKGROUND_CONTEXT); }, create: async ({ id }) => { - const session = await this.repo.create({ id }); + const session = await this.repo.create({ id }, BACKGROUND_CONTEXT); try { return session.metadata; } finally { - await session.close(); + await session.close(BACKGROUND_CONTEXT); } }, }; @@ -226,7 +226,7 @@ export class TestServerHost implements PiServerHost { gate.entered.resolve(undefined); await gate.release.promise; } - const session = await this.repo.open(metadata); + const session = await this.repo.open(metadata, BACKGROUND_CONTEXT); try { if (this.nextCreateHarnessError) { const error = this.nextCreateHarnessError; @@ -243,15 +243,15 @@ export class TestServerHost implements PiServerHost { this.harnesses.set(metadata.id, harnesses); return harness; } catch (error) { - await session.close(); + await session.close(BACKGROUND_CONTEXT); throw error; } } async seed(id = "session-1", parentSessionId?: string): Promise { - const session = await this.repo.create({ id, parentSessionId }); + const session = await this.repo.create({ id, parentSessionId }, BACKGROUND_CONTEXT); const metadata = session.metadata; - await session.close(); + await session.close(BACKGROUND_CONTEXT); return metadata; } diff --git a/packages/session-backends/sqlite-node/src/sqlite/repo.ts b/packages/session-backends/sqlite-node/src/sqlite/repo.ts index 4b0eceb66..8eba001e9 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/repo.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/repo.ts @@ -1,6 +1,6 @@ import { access, mkdir, open as openFile, readdir, rm } from "node:fs/promises"; import { basename, join } from "node:path"; -import type { Entry, ForkOptions, SessionCreateOptions, StoredValue } from "@earendil-works/pi-agent-core"; +import type { Context, Entry, ForkOptions, SessionCreateOptions, StoredValue } from "@earendil-works/pi-agent-core"; import { createForkSnapshot, laneLeaf, StorageBackedSession } from "@earendil-works/pi-agent-core"; import { uuidv7 } from "@earendil-works/pi-ai"; import { applyInitialSchema } from "./migrations.ts"; @@ -151,7 +151,8 @@ export class SqliteSessionRepo { this.now = options.now ?? Date.now; } - async create(options: SqliteSessionCreateOptions = {}): Promise { + async create(options: SqliteSessionCreateOptions | undefined, _context: Context): Promise { + options ??= {}; await mkdir(this.directory, { recursive: true }); const createdAt = this.now(); const id = options.id ?? uuidv7(createdAt); @@ -202,7 +203,7 @@ export class SqliteSessionRepo { } } - async open(metadata: SqliteSessionMetadata): Promise { + async open(metadata: SqliteSessionMetadata, _context: Context): Promise { this.reserveId(metadata.id); let db: SqliteDatabase | undefined; let lease: WriterLeaseRow | undefined; @@ -242,7 +243,7 @@ export class SqliteSessionRepo { } } - async list(): Promise { + async list(_options: undefined, _context: Context): Promise { await mkdir(this.directory, { recursive: true }); const names = await readdir(this.directory); const sessions: SqliteSessionMetadata[] = []; @@ -266,7 +267,7 @@ export class SqliteSessionRepo { return sessions.sort((left, right) => right.createdAt - left.createdAt); } - async delete(metadata: SqliteSessionMetadata): Promise { + async delete(metadata: SqliteSessionMetadata, _context: Context): Promise { if (this.pendingIds.has(metadata.id)) throw new Error(`Session is open: ${metadata.id}`); await access(metadata.path); const db = await this.databaseFactory.open(metadata.path); @@ -282,12 +283,12 @@ export class SqliteSessionRepo { await removeSessionFiles(metadata.path, { force: false }); } - async fork(source: SqliteSessionMetadata, options: ForkOptions): Promise { + async fork(source: SqliteSessionMetadata, options: ForkOptions, context: Context): Promise { const createdAt = this.now(); const id = options.id ?? uuidv7(createdAt); this.reserveId(id); const sourceStorage = this.openStorages.get(source.id); - const activeSourceSnapshot = sourceStorage?.snapshot(options); + const activeSourceSnapshot = sourceStorage?.snapshot(options, context); void activeSourceSnapshot?.catch(() => undefined); await mkdir(this.directory, { recursive: true }); const path = sessionPath(this.directory, id); diff --git a/packages/session-backends/sqlite-node/src/sqlite/session.ts b/packages/session-backends/sqlite-node/src/sqlite/session.ts index cf9f523a1..286fcf989 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/session.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/session.ts @@ -1,6 +1,7 @@ import type { AgentMessage, BranchScan, + Context, Entry, EntryQuery, JsonValue, @@ -57,132 +58,146 @@ export class SqliteOpenSession implements Session { this.renewalTimer.unref?.(); } - mutate(lane: string, mutation: (mutator: SessionMutator) => T | Promise): Promise { - return this.admit(() => this.session.mutate(lane, mutation)); + mutate( + lane: string, + mutation: (mutator: SessionMutator, context: Context) => T | Promise, + context: Context, + ): Promise { + return this.admit(() => this.session.mutate(lane, mutation, context)); } - getEntries(ids: string[]): Promise> { - return this.admit(() => this.session.getEntries(ids)); + getEntries(ids: string[], context: Context): Promise> { + return this.admit(() => this.session.getEntries(ids, context)); } - getValue(address: Value): Promise | undefined> { - return this.admit(() => this.session.getValue(address)); + getValue(address: Value, context: Context): Promise | undefined> { + return this.admit(() => this.session.getValue(address, context)); } - scanValues(prefix: Value): Promise[]> { - return this.admit(() => this.session.scanValues(prefix)); + scanValues(prefix: Value, context: Context): Promise[]> { + return this.admit(() => this.session.scanValues(prefix, context)); } - readList(address: ValueList, options?: ListReadOptions): Promise[]> { - return this.admit(() => this.session.readList(address, options)); + readList( + address: ValueList, + options: ListReadOptions | undefined, + context: Context, + ): Promise[]> { + return this.admit(() => this.session.readList(address, options, context)); } view(lane: string): SessionTree { const view = this.session.view(lane); return { - getLeafId: () => this.admit(() => view.getLeafId()), - getEntry: (id) => this.admit(() => view.getEntry(id)), - getStats: () => this.admit(() => view.getStats()), - getValue: (address) => this.admit(() => view.getValue(address)), - scanValues: (prefix) => this.admit(() => view.scanValues(prefix)), - readList: (address, options) => this.admit(() => view.readList(address, options)), - setValue: (address, next) => this.admit(() => view.setValue(address, next)), - deleteValue: (address) => this.admit(() => view.deleteValue(address)), - appendList: (address, element) => this.admit(() => view.appendList(address, element)), - deleteList: (address) => this.admit(() => view.deleteList(address)), - getName: () => this.admit(() => view.getName()), - setName: (name) => this.admit(() => view.setName(name)), - getLabel: (targetId) => this.admit(() => view.getLabel(targetId)), - setLabel: (targetId, label) => this.admit(() => view.setLabel(targetId, label)), - findEntries: (query) => this.admit(() => view.findEntries(query)), - findEntry: (query) => this.admit(() => view.findEntry(query)), - findEntriesOnBranch: (query) => this.admit(() => view.findEntriesOnBranch(query)), - findEntryOnBranch: (query) => this.admit(() => view.findEntryOnBranch(query)), - appendMessage: (message) => this.admit(() => view.appendMessage(message)), - appendCustomEntry: (customType, data) => this.admit(() => view.appendCustomEntry(customType, data)), + getLeafId: (context) => this.admit(() => view.getLeafId(context)), + getEntry: (id, context) => this.admit(() => view.getEntry(id, context)), + getStats: (context) => this.admit(() => view.getStats(context)), + getValue: (address, context) => this.admit(() => view.getValue(address, context)), + scanValues: (prefix, context) => this.admit(() => view.scanValues(prefix, context)), + readList: (address, options, context) => this.admit(() => view.readList(address, options, context)), + setValue: (address, next, context) => this.admit(() => view.setValue(address, next, context)), + deleteValue: (address, context) => this.admit(() => view.deleteValue(address, context)), + appendList: (address, element, context) => this.admit(() => view.appendList(address, element, context)), + deleteList: (address, context) => this.admit(() => view.deleteList(address, context)), + getName: (context) => this.admit(() => view.getName(context)), + setName: (name, context) => this.admit(() => view.setName(name, context)), + getLabel: (targetId, context) => this.admit(() => view.getLabel(targetId, context)), + setLabel: (targetId, label, context) => this.admit(() => view.setLabel(targetId, label, context)), + findEntries: (query, context) => this.admit(() => view.findEntries(query, context)), + findEntry: (query, context) => this.admit(() => view.findEntry(query, context)), + findEntriesOnBranch: (query, context) => this.admit(() => view.findEntriesOnBranch(query, context)), + findEntryOnBranch: (query, context) => this.admit(() => view.findEntryOnBranch(query, context)), + appendMessage: (message, context) => this.admit(() => view.appendMessage(message, context)), + appendCustomEntry: (customType, data, context) => + this.admit(() => view.appendCustomEntry(customType, data, context)), }; } - createLane(name: string, at: string | null, configuration: LaneConfiguration): Promise { + createLane( + name: string, + at: string | null, + configuration: LaneConfiguration, + context: Context, + ): Promise { return this.admit(async () => { - await this.session.createLane(name, at, configuration); + await this.session.createLane(name, at, configuration, context); return this.view(name); }); } - getLeafId(): Promise { - return this.admit(() => this.session.getLeafId()); + getLeafId(context: Context): Promise { + return this.admit(() => this.session.getLeafId(context)); } - getEntry(id: string): Promise { - return this.admit(() => this.session.getEntry(id)); + getEntry(id: string, context: Context): Promise { + return this.admit(() => this.session.getEntry(id, context)); } - getStats(): Promise { - return this.admit(() => this.session.getStats()); + getStats(context: Context): Promise { + return this.admit(() => this.session.getStats(context)); } - setValue(address: Value, next: NoInfer): Promise { - return this.admit(() => this.session.setValue(address, next)); + setValue(address: Value, next: NoInfer, context: Context): Promise { + return this.admit(() => this.session.setValue(address, next, context)); } - deleteValue(address: Value): Promise { - return this.admit(() => this.session.deleteValue(address)); + deleteValue(address: Value, context: Context): Promise { + return this.admit(() => this.session.deleteValue(address, context)); } - appendList(address: ValueList, element: NoInfer): Promise { - return this.admit(() => this.session.appendList(address, element)); + appendList(address: ValueList, element: NoInfer, context: Context): Promise { + return this.admit(() => this.session.appendList(address, element, context)); } - deleteList(address: ValueList): Promise { - return this.admit(() => this.session.deleteList(address)); + deleteList(address: ValueList, context: Context): Promise { + return this.admit(() => this.session.deleteList(address, context)); } - getName(): Promise { - return this.admit(() => this.session.getName()); + getName(context: Context): Promise { + return this.admit(() => this.session.getName(context)); } - setName(name: string | undefined): Promise { - return this.admit(() => this.session.setName(name)); + setName(name: string | undefined, context: Context): Promise { + return this.admit(() => this.session.setName(name, context)); } - getLabel(targetId: string): Promise { - return this.admit(() => this.session.getLabel(targetId)); + getLabel(targetId: string, context: Context): Promise { + return this.admit(() => this.session.getLabel(targetId, context)); } - setLabel(targetId: string, label: string | undefined): Promise { - return this.admit(() => this.session.setLabel(targetId, label)); + setLabel(targetId: string, label: string | undefined, context: Context): Promise { + return this.admit(() => this.session.setLabel(targetId, label, context)); } - findEntries(query?: EntryQuery): Promise { - return this.admit(() => this.session.findEntries(query)); + findEntries(query: EntryQuery | undefined, context: Context): Promise { + return this.admit(() => this.session.findEntries(query, context)); } - findEntry(query?: EntryQuery): Promise { - return this.admit(() => this.session.findEntry(query)); + findEntry(query: EntryQuery | undefined, context: Context): Promise { + return this.admit(() => this.session.findEntry(query, context)); } - findEntriesOnBranch(query?: BranchScan): Promise { - return this.admit(() => this.session.findEntriesOnBranch(query)); + findEntriesOnBranch(query: BranchScan | undefined, context: Context): Promise { + return this.admit(() => this.session.findEntriesOnBranch(query, context)); } - findEntryOnBranch(query?: BranchScan): Promise { - return this.admit(() => this.session.findEntryOnBranch(query)); + findEntryOnBranch(query: BranchScan | undefined, context: Context): Promise { + return this.admit(() => this.session.findEntryOnBranch(query, context)); } - appendMessage(message: AgentMessage): Promise { - return this.admit(() => this.session.appendMessage(message)); + appendMessage(message: AgentMessage, context: Context): Promise { + return this.admit(() => this.session.appendMessage(message, context)); } - appendCustomEntry(customType: string, data?: JsonValue): Promise { - return this.admit(() => this.session.appendCustomEntry(customType, data)); + appendCustomEntry(customType: string, data: JsonValue | undefined, context: Context): Promise { + return this.admit(() => this.session.appendCustomEntry(customType, data, context)); } - close(): Promise { + close(context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; this.state = "closing"; this.closePromise = Promise.allSettled([...this.admitted]) - .then(() => this.session.close()) + .then(() => this.session.close(context)) .finally(() => { clearInterval(this.renewalTimer); try { diff --git a/packages/session-backends/sqlite-node/src/sqlite/storage.ts b/packages/session-backends/sqlite-node/src/sqlite/storage.ts index 6ebed1d57..740520979 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/storage.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/storage.ts @@ -1,5 +1,6 @@ import type { CommitResult, + Context, Entry, EntryScan, EntryStructure, @@ -63,7 +64,7 @@ export class SqliteStorage implements Storage { this.usageWriter = new UsageLedgerRowWriter(db); } - async commit(writes: Write[]): Promise { + async commit(writes: Write[], _context: Context): Promise { if (this.state !== "open") throw new Error("SqliteStorage is closed"); const result = this.commitQueue.then(() => this.applyCommit(writes)); this.commitQueue = result.then( @@ -73,7 +74,7 @@ export class SqliteStorage implements Storage { return result; } - getEntries(ids: string[]): Promise> { + getEntries(ids: string[], _context: Context): Promise> { if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); const rowsById = new Map(readEntryRows(this.db, ids).map((row) => [row.id, row])); const entries = new Map(); @@ -84,47 +85,52 @@ export class SqliteStorage implements Storage { return Promise.resolve(entries); } - getValue(address: Value): Promise | undefined> { + getValue(address: Value, _context: Context): Promise | undefined> { if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); return Promise.resolve(readScalarValueRow(this.db, address)); } - scanValues(prefix: Value): Promise[]> { + scanValues(prefix: Value, _context: Context): Promise[]> { if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); return Promise.resolve(scanScalarValueRows(this.db, prefix)); } - async readList(address: ValueList, options?: ListReadOptions): Promise[]> { + async readList( + address: ValueList, + options: ListReadOptions | undefined, + _context: Context, + ): Promise[]> { if (this.state !== "open") throw new Error("SqliteStorage is closed"); return readListValueRows(this.db, address, options); } - scanBranch(query: StorageBranchScan): Promise { + scanBranch(query: StorageBranchScan, _context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); return Promise.resolve(scanBranchEntries(this.db, query)); } - scanBranchStructure(query: StorageBranchScan): Promise { + scanBranchStructure(query: StorageBranchScan, _context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); return Promise.resolve(scanBranchEntryStructures(this.db, query)); } - scanEntries(query: EntryScan): Promise { + scanEntries(query: EntryScan, _context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); return Promise.resolve(scanEntryRows(this.db, query).map(decodeEntryRow)); } - scanUsage(query: UsageScan): Promise { + scanUsage(query: UsageScan, _context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); return Promise.resolve(scanUsageLedgerRows(this.db, query).map(decodeUsageLedgerRow)); } - getStats(): Promise { + getStats(_context: Context): Promise { if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); return Promise.resolve(readSessionStats(this.db)); } - snapshot(options: ForkOptions = {}): Promise { + snapshot(options: ForkOptions | undefined, _context: Context): Promise { + options ??= {}; if (this.state !== "open") return Promise.reject(new Error("SqliteStorage is closed")); const result = this.commitQueue.then(() => this.readSnapshot(options)); this.commitQueue = result.then( @@ -195,7 +201,7 @@ export class SqliteStorage implements Storage { }); } - close(): Promise { + close(_context: Context): Promise { if (this.closePromise !== undefined) return this.closePromise; this.state = "closing"; this.closePromise = this.commitQueue.then(() => { diff --git a/packages/session-backends/sqlite-node/test/repo.test.ts b/packages/session-backends/sqlite-node/test/repo.test.ts index caa638a6b..24c657060 100644 --- a/packages/session-backends/sqlite-node/test/repo.test.ts +++ b/packages/session-backends/sqlite-node/test/repo.test.ts @@ -3,6 +3,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import * as storedValues from "@earendil-works/pi-agent-core"; import * as sessionWrites from "@earendil-works/pi-agent-core"; +import { BACKGROUND_CONTEXT } from "@earendil-works/pi-agent-core"; import { describe, expect, it } from "vitest"; import type { SqliteDatabase } from "../src/index.ts"; import { createNodeSqliteFactory, SqliteSessionRepo, sql } from "../src/index.ts"; @@ -34,7 +35,7 @@ describe("SqliteSessionRepo", () => { now: () => 1_700_000_000_000, }); - const session = await repo.create({ id: "session" }); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); const metadata = session.metadata; expect(metadata).toMatchObject({ id: "session", @@ -73,7 +74,7 @@ describe("SqliteSessionRepo", () => { ]); expect(sql`SELECT COUNT(*) AS count FROM list_values`.get(db)).toEqual({ count: 0 }); }); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); }); @@ -84,14 +85,14 @@ describe("SqliteSessionRepo", () => { databaseFactory: createNodeSqliteFactory(), now: () => 1, }); - const session = await repo.create({ id: "session" }); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); const { metadata } = session; - await expect(repo.create({ id: "session" })).rejects.toThrow(); + await expect(repo.create({ id: "session" }, BACKGROUND_CONTEXT)).rejects.toThrow(); await withDb(metadata.path, (db) => { expect(sql`SELECT COUNT(*) AS count FROM session`.get<{ count: number }>(db)).toEqual({ count: 1 }); }); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); }); @@ -102,18 +103,20 @@ describe("SqliteSessionRepo", () => { databaseFactory: createNodeSqliteFactory(), now: () => 1, }); - const session = await repo.create({ id: "session" }); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); const { metadata } = session; let leaseBeforeList: unknown[] = []; await withDb(metadata.path, (db) => { leaseBeforeList = sql`SELECT owner_id, fence FROM writer_lease`.all(db); }); - await expect(repo.list()).resolves.toMatchObject([{ id: "session", path: metadata.path }]); + await expect(repo.list(undefined, BACKGROUND_CONTEXT)).resolves.toMatchObject([ + { id: "session", path: metadata.path }, + ]); await withDb(metadata.path, (db) => { expect(sql`SELECT owner_id, fence FROM writer_lease`.all(db)).toEqual(leaseBeforeList); }); - await session.close(); + await session.close(BACKGROUND_CONTEXT); }); }); @@ -124,14 +127,14 @@ describe("SqliteSessionRepo", () => { databaseFactory: createNodeSqliteFactory(), now: () => 1, }); - const session = await repo.create({ id: "session" }); - await session.close(); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); + await session.close(BACKGROUND_CONTEXT); await writeFile(join(directory, "corrupt.sqlite"), "not a sqlite database"); await withDb(session.metadata.path, (db) => { sql`UPDATE session SET storage_version = ${999}`.run(db); }); - expect(await repo.list()).toEqual([]); + expect(await repo.list(undefined, BACKGROUND_CONTEXT)).toEqual([]); }); }); @@ -145,7 +148,7 @@ describe("SqliteSessionRepo", () => { const path = join(directory, "session.sqlite"); await writeFile(path, "not a sqlite database"); - await expect(repo.create({ id: "session" })).rejects.toThrow(); + await expect(repo.create({ id: "session" }, BACKGROUND_CONTEXT)).rejects.toThrow(); await expect(access(path)).resolves.toBeUndefined(); }); }); @@ -157,22 +160,24 @@ describe("SqliteSessionRepo", () => { databaseFactory: createNodeSqliteFactory(), now: () => 1, }); - const session = await repo.create({ id: "session" }); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); const { metadata } = session; - await session.close(); + await session.close(BACKGROUND_CONTEXT); await withDb(metadata.path, (db) => { sql`INSERT INTO writer_lease (owner_id, fence, expires_at_ms) VALUES (${"external"}, ${1}, ${1_000})`.run( db, ); }); - await expect(repo.delete(metadata)).rejects.toThrow("already claimed"); + await expect(repo.delete(metadata, BACKGROUND_CONTEXT)).rejects.toThrow("already claimed"); await withDb(metadata.path, (db) => { sql`DELETE FROM writer_lease`.run(db); }); - await expect(repo.delete({ ...metadata, path: join(directory, "missing.sqlite") })).rejects.toThrow(); - await repo.delete(metadata); + await expect( + repo.delete({ ...metadata, path: join(directory, "missing.sqlite") }, BACKGROUND_CONTEXT), + ).rejects.toThrow(); + await repo.delete(metadata, BACKGROUND_CONTEXT); await expect(access(metadata.path)).rejects.toThrow(); }); }); @@ -184,20 +189,20 @@ describe("SqliteSessionRepo", () => { databaseFactory: createNodeSqliteFactory(), now: () => 1, }); - const created = await repo.create({ id: "session" }); + const created = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); const { metadata } = created; - await created.close(); + await created.close(BACKGROUND_CONTEXT); - const opened = await repo.open(metadata); + const opened = await repo.open(metadata, BACKGROUND_CONTEXT); expect(opened.metadata).toMatchObject({ id: "session", path: metadata.path }); - await opened.close(); + await opened.close(BACKGROUND_CONTEXT); await withDb(metadata.path, (db) => { sql`INSERT INTO writer_lease (owner_id, fence, expires_at_ms) VALUES (${"external"}, ${1}, ${1_000})`.run( db, ); }); - await expect(repo.open(metadata)).rejects.toThrow("already claimed"); + await expect(repo.open(metadata, BACKGROUND_CONTEXT)).rejects.toThrow("already claimed"); }); }); @@ -208,33 +213,47 @@ describe("SqliteSessionRepo", () => { databaseFactory: createNodeSqliteFactory(), now: () => 1_700_000_000_000, }); - const source = await repo.create({ id: "source" }); + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); const applicationValue = storedValues.value("test.application.value"); const applicationList = storedValues.list("test.application.list"); - const firstCommit = source.mutate("main", (mutator) => - mutator.commit([ - sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "root" }), - storedValues.setValue(storedValues.laneLeaf("main"), "root"), - storedValues.setValue(applicationValue, "excluded"), - storedValues.appendList(applicationList, "excluded"), - ]), + const firstCommit = source.mutate( + "main", + (mutator) => + mutator.commit( + [ + sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "root" }), + storedValues.setValue(storedValues.laneLeaf("main"), "root"), + storedValues.setValue(applicationValue, "excluded"), + storedValues.appendList(applicationList, "excluded"), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - const forkPromise = repo.fork(source.metadata, { id: "fork" }); - const secondCommit = source.mutate("main", (mutator) => - mutator.commit([ - sessionWrites.insertEntry({ id: "child", parentId: "root", type: "custom", customType: "child" }), - storedValues.setValue(storedValues.laneLeaf("main"), "child"), - ]), + const forkPromise = repo.fork(source.metadata, { id: "fork" }, BACKGROUND_CONTEXT); + const secondCommit = source.mutate( + "main", + (mutator) => + mutator.commit( + [ + sessionWrites.insertEntry({ id: "child", parentId: "root", type: "custom", customType: "child" }), + storedValues.setValue(storedValues.laneLeaf("main"), "child"), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); const [, fork] = await Promise.all([firstCommit, forkPromise]); await secondCommit; - expect(await fork.getLeafId()).toBe("root"); - expect((await fork.findEntries({ order: "asc" })).map((entry) => entry.id)).toEqual(["root"]); - expect(await fork.getValue(applicationValue)).toBeUndefined(); - expect(await fork.readList(applicationList)).toEqual([]); - await Promise.all([source.close(), fork.close()]); + expect(await fork.getLeafId(BACKGROUND_CONTEXT)).toBe("root"); + expect((await fork.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ + "root", + ]); + expect(await fork.getValue(applicationValue, BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await fork.readList(applicationList, undefined, BACKGROUND_CONTEXT)).toEqual([]); + await Promise.all([source.close(BACKGROUND_CONTEXT), fork.close(BACKGROUND_CONTEXT)]); }); }); @@ -245,41 +264,62 @@ describe("SqliteSessionRepo", () => { databaseFactory: createNodeSqliteFactory(), now: () => 1_700_000_000_000, }); - const source = await repo.create({ id: "source" }); + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); const configuration = { model: { provider: "test", modelId: "model" }, thinkingLevel: "off" as const, activeToolNames: [], }; - await source.mutate("main", (mutator) => - mutator.commit([ - sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "root" }), - sessionWrites.insertEntry({ id: "child", parentId: "root", type: "custom", customType: "child" }), - sessionWrites.insertEntry({ id: "sibling", parentId: "root", type: "custom", customType: "sibling" }), - storedValues.setValue(storedValues.laneLeaf("main"), "child"), - storedValues.setValue(storedValues.laneLeaf("review"), "sibling"), - storedValues.setValue(storedValues.laneConfig("review"), configuration), - storedValues.setValue(storedValues.laneState("review"), { - currentOperationId: null, - pendingNextRun: [], - }), - ]), + await source.mutate( + "main", + (mutator) => + mutator.commit( + [ + sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "root" }), + sessionWrites.insertEntry({ id: "child", parentId: "root", type: "custom", customType: "child" }), + sessionWrites.insertEntry({ + id: "sibling", + parentId: "root", + type: "custom", + customType: "sibling", + }), + storedValues.setValue(storedValues.laneLeaf("main"), "child"), + storedValues.setValue(storedValues.laneLeaf("review"), "sibling"), + storedValues.setValue(storedValues.laneConfig("review"), configuration), + storedValues.setValue(storedValues.laneState("review"), { + currentOperationId: null, + pendingNextRun: [], + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - const fork = await repo.fork(source.metadata, { id: "fork", scope: "tree" }); + const fork = await repo.fork(source.metadata, { id: "fork", scope: "tree" }, BACKGROUND_CONTEXT); - expect((await fork.findEntries({ order: "asc" })).map(({ id }) => id)).toEqual(["root", "child", "sibling"]); - expect(await fork.getLeafId()).toBe("child"); - expect(await fork.view("review").getLeafId()).toBe("sibling"); - expect((await fork.getValue(storedValues.laneState("review")))?.value).toEqual({ + expect((await fork.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id)).toEqual([ + "root", + "child", + "sibling", + ]); + expect(await fork.getLeafId(BACKGROUND_CONTEXT)).toBe("child"); + expect(await fork.view("review").getLeafId(BACKGROUND_CONTEXT)).toBe("sibling"); + expect((await fork.getValue(storedValues.laneState("review"), BACKGROUND_CONTEXT))?.value).toEqual({ currentOperationId: null, pendingNextRun: [], }); - const branchFork = await repo.fork(source.metadata, { id: "branch", entryId: "root" }); - expect((await branchFork.findEntries({ order: "asc" })).map(({ id }) => id)).toEqual(["root"]); - expect(await branchFork.getLeafId()).toBe("root"); - await Promise.all([source.close(), fork.close(), branchFork.close()]); + const branchFork = await repo.fork(source.metadata, { id: "branch", entryId: "root" }, BACKGROUND_CONTEXT); + expect((await branchFork.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id)).toEqual([ + "root", + ]); + expect(await branchFork.getLeafId(BACKGROUND_CONTEXT)).toBe("root"); + await Promise.all([ + source.close(BACKGROUND_CONTEXT), + fork.close(BACKGROUND_CONTEXT), + branchFork.close(BACKGROUND_CONTEXT), + ]); }); }); @@ -290,69 +330,85 @@ describe("SqliteSessionRepo", () => { databaseFactory: createNodeSqliteFactory(), now: () => 1_700_000_000_000, }); - const source = await repo.create({ id: "source" }); + const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); const applicationValue = storedValues.value("test.application.value"); const applicationList = storedValues.list("test.application.list"); - await source.mutate("main", (mutator) => - mutator.commit([ - sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "root" }), - sessionWrites.insertEntry({ - id: "child", - parentId: "root", - type: "message", - message: { role: "user", content: "child", timestamp: 1 }, - }), - sessionWrites.insertEntry({ id: "sibling", parentId: "root", type: "custom", customType: "sibling" }), - storedValues.setValue(storedValues.laneLeaf("main"), "child"), - storedValues.setValue(storedValues.sessionName, "source name"), - storedValues.setValue(storedValues.entryLabel("root"), "root label"), - storedValues.setValue(storedValues.entryLabel("sibling"), "sibling label"), - storedValues.setValue(applicationValue, "excluded"), - storedValues.appendList(applicationList, "excluded"), - storedValues.setValue(storedValues.laneLastResult("main"), { - operationId: "previous", - kind: "navigation", - leafId: "child", - oldLeafId: "root", - outcome: "completed", - }), - storedValues.setValue(storedValues.pendingEntry("pending"), { - type: "custom", - customType: "pending", - }), - storedValues.setValue(storedValues.operationToolMemo("operation", "invocation", "memo"), true), - sessionWrites.insertUsage({ - id: "usage", - adjustment: false, - usage: { - input: 1, - output: 1, - cacheRead: 0, - cacheWrite: 0, - totalTokens: 2, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - }, - }), - ]), + await source.mutate( + "main", + (mutator) => + mutator.commit( + [ + sessionWrites.insertEntry({ id: "root", parentId: null, type: "custom", customType: "root" }), + sessionWrites.insertEntry({ + id: "child", + parentId: "root", + type: "message", + message: { role: "user", content: "child", timestamp: 1 }, + }), + sessionWrites.insertEntry({ + id: "sibling", + parentId: "root", + type: "custom", + customType: "sibling", + }), + storedValues.setValue(storedValues.laneLeaf("main"), "child"), + storedValues.setValue(storedValues.sessionName, "source name"), + storedValues.setValue(storedValues.entryLabel("root"), "root label"), + storedValues.setValue(storedValues.entryLabel("sibling"), "sibling label"), + storedValues.setValue(applicationValue, "excluded"), + storedValues.appendList(applicationList, "excluded"), + storedValues.setValue(storedValues.laneLastResult("main"), { + operationId: "previous", + kind: "navigation", + leafId: "child", + oldLeafId: "root", + outcome: "completed", + }), + storedValues.setValue(storedValues.pendingEntry("pending"), { + type: "custom", + customType: "pending", + }), + storedValues.setValue(storedValues.operationToolMemo("operation", "invocation", "memo"), true), + sessionWrites.insertUsage({ + id: "usage", + adjustment: false, + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }), + ], + BACKGROUND_CONTEXT, + ), + BACKGROUND_CONTEXT, ); - await source.close(); - const fork = await repo.fork(source.metadata, { id: "fork", entryId: "child" }); + await source.close(BACKGROUND_CONTEXT); + const fork = await repo.fork(source.metadata, { id: "fork", entryId: "child" }, BACKGROUND_CONTEXT); - expect((await fork.findEntries({ order: "asc" })).map((entry) => entry.id)).toEqual(["root", "child"]); - expect(await fork.getLeafId()).toBe("child"); - expect(await fork.getName()).toBe("source name"); - expect(await fork.getLabel("root")).toBe("root label"); - expect(await fork.getLabel("sibling")).toBeUndefined(); - expect(await fork.getValue(applicationValue)).toBeUndefined(); - expect(await fork.readList(applicationList)).toEqual([]); - expect(await fork.getValue(storedValues.laneLastResult("main"))).toBeUndefined(); - expect(await fork.getValue(storedValues.pendingEntry("pending"))).toBeUndefined(); - expect(await fork.getValue(storedValues.operationToolMemo("operation", "invocation", "memo"))).toBeUndefined(); + expect((await fork.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ + "root", + "child", + ]); + expect(await fork.getLeafId(BACKGROUND_CONTEXT)).toBe("child"); + expect(await fork.getName(BACKGROUND_CONTEXT)).toBe("source name"); + expect(await fork.getLabel("root", BACKGROUND_CONTEXT)).toBe("root label"); + expect(await fork.getLabel("sibling", BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await fork.getValue(applicationValue, BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await fork.readList(applicationList, undefined, BACKGROUND_CONTEXT)).toEqual([]); + expect(await fork.getValue(storedValues.laneLastResult("main"), BACKGROUND_CONTEXT)).toBeUndefined(); + expect(await fork.getValue(storedValues.pendingEntry("pending"), BACKGROUND_CONTEXT)).toBeUndefined(); + expect( + await fork.getValue(storedValues.operationToolMemo("operation", "invocation", "memo"), BACKGROUND_CONTEXT), + ).toBeUndefined(); await withDb(fork.metadata.path, (db) => { expect(sql`SELECT COUNT(*) AS count FROM usage_ledger`.get<{ count: number }>(db)).toEqual({ count: 0 }); }); - expect(await fork.getStats()).toEqual({ + expect(await fork.getStats(BACKGROUND_CONTEXT)).toEqual({ messageCount: 1, usage: { input: 0, @@ -363,7 +419,7 @@ describe("SqliteSessionRepo", () => { cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, }); - await Promise.all([source.close(), fork.close()]); + await Promise.all([source.close(BACKGROUND_CONTEXT), fork.close(BACKGROUND_CONTEXT)]); }); }); }); diff --git a/packages/session-backends/sqlite-node/test/storage-conformance.test.ts b/packages/session-backends/sqlite-node/test/storage-conformance.test.ts index 0d65b67ec..2d9e132cb 100644 --- a/packages/session-backends/sqlite-node/test/storage-conformance.test.ts +++ b/packages/session-backends/sqlite-node/test/storage-conformance.test.ts @@ -1,3 +1,4 @@ +import { BACKGROUND_CONTEXT } from "@earendil-works/pi-agent-core"; import { type ConformanceCase, createStorageConformance, @@ -49,7 +50,7 @@ registerConformance( storage, async [Symbol.asyncDispose]() { try { - await storage.close(); + await storage.close(BACKGROUND_CONTEXT); } finally { db.close(); } diff --git a/packages/session-backends/sqlite-node/test/storage.test.ts b/packages/session-backends/sqlite-node/test/storage.test.ts index feca3dfa7..7e2928b79 100644 --- a/packages/session-backends/sqlite-node/test/storage.test.ts +++ b/packages/session-backends/sqlite-node/test/storage.test.ts @@ -1,6 +1,6 @@ import * as storedValues from "@earendil-works/pi-agent-core"; import * as sessionWrites from "@earendil-works/pi-agent-core"; -import { prepareStorageCommit } from "@earendil-works/pi-agent-core"; +import { BACKGROUND_CONTEXT, prepareStorageCommit } from "@earendil-works/pi-agent-core"; import { describe, expect, it } from "vitest"; import { createNodeSqliteFactory, type SqliteDatabase, SqliteStorage, sql } from "../src/index.ts"; import { applyInitialSchema } from "../src/sqlite/migrations.ts"; @@ -59,24 +59,30 @@ describe("SqliteStorage", () => { insertCommitSessionRow(db); expect( - await storage.commit([ - sessionWrites.insertEntry({ - id: "root", - parentId: null, - type: "message", - message: { role: "user", content: "root", timestamp: 10 }, - }), - ]), + await storage.commit( + [ + sessionWrites.insertEntry({ + id: "root", + parentId: null, + type: "message", + message: { role: "user", content: "root", timestamp: 10 }, + }), + ], + BACKGROUND_CONTEXT, + ), ).toEqual({ firstSeq: 1, seqs: [1], timestamp: 1_700_000_000_000 }); expect( - await storage.commit([ - sessionWrites.insertEntry({ - id: "child", - parentId: "root", - type: "message", - message: { role: "user", content: "child", timestamp: 11 }, - }), - ]), + await storage.commit( + [ + sessionWrites.insertEntry({ + id: "child", + parentId: "root", + type: "message", + message: { role: "user", content: "child", timestamp: 11 }, + }), + ], + BACKGROUND_CONTEXT, + ), ).toEqual({ firstSeq: 2, seqs: [2], timestamp: 1_700_000_000_000 }); expect(sql`SELECT branch_id, tip_entry_id, tip_seq FROM branch_meta`.all(db)).toEqual([ @@ -88,33 +94,39 @@ describe("SqliteStorage", () => { { branch_id: "root", entry_id: "root", entry_seq: 1, entry_type: "message" }, { branch_id: "root", entry_id: "child", entry_seq: 2, entry_type: "message" }, ]); - expect((await storage.scanBranch({ start: "child" })).map((entry) => entry.id)).toEqual(["child", "root"]); + expect((await storage.scanBranch({ start: "child" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ + "child", + "root", + ]); }); }); it("commits divergent branch entries by materializing a new segment", async () => { await withStorage(async (storage, db) => { insertCommitSessionRow(db); - await storage.commit([ - sessionWrites.insertEntry({ - id: "root", - parentId: null, - type: "message", - message: { role: "user", content: "root", timestamp: 10 }, - }), - sessionWrites.insertEntry({ - id: "left", - parentId: "root", - type: "message", - message: { role: "user", content: "left", timestamp: 11 }, - }), - sessionWrites.insertEntry({ - id: "right", - parentId: "root", - type: "message", - message: { role: "user", content: "right", timestamp: 12 }, - }), - ]); + await storage.commit( + [ + sessionWrites.insertEntry({ + id: "root", + parentId: null, + type: "message", + message: { role: "user", content: "root", timestamp: 10 }, + }), + sessionWrites.insertEntry({ + id: "left", + parentId: "root", + type: "message", + message: { role: "user", content: "left", timestamp: 11 }, + }), + sessionWrites.insertEntry({ + id: "right", + parentId: "root", + type: "message", + message: { role: "user", content: "right", timestamp: 12 }, + }), + ], + BACKGROUND_CONTEXT, + ); expect( sql`SELECT branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq FROM branch_meta ORDER BY branch_id`.all( @@ -134,49 +146,58 @@ describe("SqliteStorage", () => { { branch_id: "root", entry_id: "root", entry_seq: 1, entry_type: "message" }, { branch_id: "root", entry_id: "left", entry_seq: 2, entry_type: "message" }, ]); - expect((await storage.scanBranch({ start: "right" })).map((entry) => entry.id)).toEqual(["right", "root"]); - expect((await storage.scanBranch({ start: "left" })).map((entry) => entry.id)).toEqual(["left", "root"]); + expect((await storage.scanBranch({ start: "right" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ + "right", + "root", + ]); + expect((await storage.scanBranch({ start: "left" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ + "left", + "root", + ]); }); }); it("bases divergent branch segments at the newest compaction", async () => { await withStorage(async (storage, db) => { insertCommitSessionRow(db); - await storage.commit([ - sessionWrites.insertEntry({ - id: "root", - parentId: null, - type: "message", - message: { role: "user", content: "root", timestamp: 10 }, - }), - sessionWrites.insertEntry({ - id: "compact", - parentId: "root", - type: "compaction", - summary: "summary", - retainedTail: [], - tokensBefore: 1, - fromHook: false, - }), - sessionWrites.insertEntry({ - id: "left", - parentId: "compact", - type: "message", - message: { role: "user", content: "left", timestamp: 11 }, - }), - sessionWrites.insertEntry({ - id: "leaf", - parentId: "left", - type: "message", - message: { role: "user", content: "leaf", timestamp: 12 }, - }), - sessionWrites.insertEntry({ - id: "right", - parentId: "left", - type: "message", - message: { role: "user", content: "right", timestamp: 13 }, - }), - ]); + await storage.commit( + [ + sessionWrites.insertEntry({ + id: "root", + parentId: null, + type: "message", + message: { role: "user", content: "root", timestamp: 10 }, + }), + sessionWrites.insertEntry({ + id: "compact", + parentId: "root", + type: "compaction", + summary: "summary", + retainedTail: [], + tokensBefore: 1, + fromHook: false, + }), + sessionWrites.insertEntry({ + id: "left", + parentId: "compact", + type: "message", + message: { role: "user", content: "left", timestamp: 11 }, + }), + sessionWrites.insertEntry({ + id: "leaf", + parentId: "left", + type: "message", + message: { role: "user", content: "leaf", timestamp: 12 }, + }), + sessionWrites.insertEntry({ + id: "right", + parentId: "left", + type: "message", + message: { role: "user", content: "right", timestamp: 13 }, + }), + ], + BACKGROUND_CONTEXT, + ); expect( sql`SELECT branch_id, tip_entry_id, tip_seq, base_branch_id, base_seq FROM branch_meta WHERE branch_id = ${"right"}`.get( @@ -197,7 +218,7 @@ describe("SqliteStorage", () => { { branch_id: "right", entry_id: "left", entry_seq: 3, entry_type: "message" }, { branch_id: "right", entry_id: "right", entry_seq: 5, entry_type: "message" }, ]); - expect((await storage.scanBranch({ start: "right" })).map((entry) => entry.id)).toEqual([ + expect((await storage.scanBranch({ start: "right" }, BACKGROUND_CONTEXT)).map((entry) => entry.id)).toEqual([ "right", "left", "compact", @@ -215,7 +236,7 @@ describe("SqliteStorage", () => { db, ); - const entries = await storage.getEntries(["second", "missing", "first"]); + const entries = await storage.getEntries(["second", "missing", "first"], BACKGROUND_CONTEXT); expect([...entries.keys()]).toEqual(["second", "first"]); expect(entries.get("second")).toMatchObject({ id: "second", type: "message", parentId: "first" }); @@ -232,10 +253,14 @@ describe("SqliteStorage", () => { (${"three"}, ${"two"}, ${3}, ${"custom"}, ${"note"}, ${12}, ${JSON.stringify({ data: 3 })})`.run(db); expect( - (await storage.scanEntries({ order: "desc", type: "custom", fromSeq: 2 })).map((entry) => entry.id), + (await storage.scanEntries({ order: "desc", type: "custom", fromSeq: 2 }, BACKGROUND_CONTEXT)).map( + (entry) => entry.id, + ), ).toEqual(["three"]); expect( - (await storage.scanEntries({ order: "asc", customType: "note", limit: 2 })).map((entry) => entry.id), + (await storage.scanEntries({ order: "asc", customType: "note", limit: 2 }, BACKGROUND_CONTEXT)).map( + (entry) => entry.id, + ), ).toEqual(["one", "three"]); }); }); @@ -264,11 +289,16 @@ describe("SqliteStorage", () => { (${"new"}, ${"leaf"}, ${5}, ${"message"})`.run(db); expect( - (await storage.scanBranch({ start: "leaf", stopAtType: "compaction", limit: 2 })).map((entry) => entry.id), + (await storage.scanBranch({ start: "leaf", stopAtType: "compaction", limit: 2 }, BACKGROUND_CONTEXT)).map( + (entry) => entry.id, + ), ).toEqual(["leaf", "custom"]); expect( ( - await storage.scanBranch({ start: "leaf", order: "oldestFirst", type: "message", cursor: { seq: 1 } }) + await storage.scanBranch( + { start: "leaf", order: "oldestFirst", type: "message", cursor: { seq: 1 } }, + BACKGROUND_CONTEXT, + ) ).map((entry) => entry.id), ).toEqual(["old", "leaf"]); }); @@ -294,7 +324,9 @@ describe("SqliteStorage", () => { (${"zzz-base"}, ${"base-tip"}, ${2}, ${"message"})`.run(db); expect( - (await storage.scanBranch({ start: "base-tip", order: "oldestFirst" })).map((entry) => entry.id), + (await storage.scanBranch({ start: "base-tip", order: "oldestFirst" }, BACKGROUND_CONTEXT)).map( + (entry) => entry.id, + ), ).toEqual(["root", "base-tip"]); }); }); @@ -319,12 +351,17 @@ describe("SqliteStorage", () => { (${"new"}, ${"leaf"}, ${4}, ${"custom"})`.run(db); expect( - (await storage.scanBranch({ start: "leaf", stopAtType: "compaction" })).map((entry) => entry.id), - ).toEqual(["leaf", "after", "compact"]); - expect( - (await storage.scanBranch({ start: "leaf", stopAtType: "compaction", type: "message" })).map( + (await storage.scanBranch({ start: "leaf", stopAtType: "compaction" }, BACKGROUND_CONTEXT)).map( (entry) => entry.id, ), + ).toEqual(["leaf", "after", "compact"]); + expect( + ( + await storage.scanBranch( + { start: "leaf", stopAtType: "compaction", type: "message" }, + BACKGROUND_CONTEXT, + ) + ).map((entry) => entry.id), ).toEqual(["after"]); }); }); @@ -361,16 +398,18 @@ describe("SqliteStorage", () => { (${"main"}, ${"root"}, ${1}, ${"message"}), (${"main"}, ${"custom"}, ${2}, ${"custom"})`.run(db); - expect(await storage.scanBranchStructure({ start: "custom", customType: "note" })).toEqual([ - { - id: "custom", - parentId: "root", - seq: 2, - timestamp: 11, - type: "custom", - customType: "note", - }, - ]); + expect(await storage.scanBranchStructure({ start: "custom", customType: "note" }, BACKGROUND_CONTEXT)).toEqual( + [ + { + id: "custom", + parentId: "root", + seq: 2, + timestamp: 11, + type: "custom", + customType: "note", + }, + ], + ); }); }); @@ -409,10 +448,12 @@ describe("SqliteStorage", () => { (${"u2"}, ${2}, ${null}, ${1}, ${JSON.stringify(usage)}, ${JSON.stringify({ reason: "adjust" })}), (${"u3"}, ${3}, ${"e3"}, ${0}, ${JSON.stringify(usage)}, ${null})`.run(db); - expect(await storage.scanUsage({ fromSeq: 2, order: "asc", limit: 1 })).toEqual([ + expect(await storage.scanUsage({ fromSeq: 2, order: "asc", limit: 1 }, BACKGROUND_CONTEXT)).toEqual([ { id: "u2", seq: 2, usage, adjustment: true, details: { reason: "adjust" } }, ]); - expect((await storage.scanUsage({ toSeq: 2, order: "desc" })).map((row) => row.id)).toEqual(["u2", "u1"]); + expect( + (await storage.scanUsage({ toSeq: 2, order: "desc" }, BACKGROUND_CONTEXT)).map((row) => row.id), + ).toEqual(["u2", "u1"]); }); }); @@ -479,7 +520,7 @@ describe("SqliteStorage", () => { (created_at, parent_session_id, storage_version, metadata, message_count, usage_payload, next_seq) VALUES (${1}, ${null}, ${1}, ${null}, ${2}, ${JSON.stringify(usage)}, ${3})`.run(db); - expect(await storage.getStats()).toEqual({ messageCount: 2, usage }); + expect(await storage.getStats(BACKGROUND_CONTEXT)).toEqual({ messageCount: 2, usage }); }); }); @@ -489,12 +530,14 @@ describe("SqliteStorage", () => { sql`INSERT INTO scalar_values (namespace, key, seq, value) VALUES (${address.namespace}, ${address.key}, ${1}, ${JSON.stringify({ ready: true })})`.run(db); - expect(await storage.getValue(address)).toEqual({ + expect(await storage.getValue(address, BACKGROUND_CONTEXT)).toEqual({ address, seq: 1, value: { ready: true }, }); - expect(await storage.getValue(storedValues.value("test.value", "missing"))).toBeUndefined(); + expect( + await storage.getValue(storedValues.value("test.value", "missing"), BACKGROUND_CONTEXT), + ).toBeUndefined(); }); }); @@ -509,7 +552,9 @@ describe("SqliteStorage", () => { (${"test.value"}, ${"other"}, ${5}, ${JSON.stringify(5)}), (${"other.value"}, ${""}, ${6}, ${JSON.stringify("name")})`.run(db); - expect(await storage.scanValues(storedValues.value("test.value", "app:"))).toEqual([ + expect( + await storage.scanValues(storedValues.value("test.value", "app:"), BACKGROUND_CONTEXT), + ).toEqual([ { address: storedValues.value("test.value", "app:one"), seq: 1, value: 1 }, { address: storedValues.value("test.value", "app:two"), seq: 2, value: 2 }, { address: storedValues.value("test.value", "app:\ufffftail"), seq: 3, value: 3 },