mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(durable): keep each conversation's context in memory; Durable Object SQLite adapter
runtime.context() keeps each conversation's last read range and derived view. A later read by any of the conversation's tasks with the same head marker scans and derives only the entries committed since, so a generation no longer rereads and re-derives the whole transcript for each model request. Busy conversations keep the range; idle ones for settings.contextRetentionMs (default ten minutes), checked on task changes and by one unreferenced timer where the runtime supports it. Kept entries are frozen; each read returns its own view arrays. Add openDurableObjectSqliteStorage(ctx.storage) in @earendil-works/pi-durable/storage/sqlite/cloudflare: SQLite storage on a SQLite-backed Cloudflare Durable Object, without per-row copies.
This commit is contained in:
Mario Zechner
parent
fe11328b0d
commit
da866ada17
12 changed files with 927 additions and 76 deletions
|
|
@ -15,10 +15,12 @@
|
|||
- Task records carry `startedAt` and `endedAt`, wall-clock times the Session stamps at the first change to `running` and at the change to `terminal`, with `HarnessOptions.now` or the new `createSession(storage, { now })` option. `startedAt` survives waits and reopen. `inspect()` shows them on each live task's record. Records written by earlier versions have neither ([#10549](https://github.com/earendil-works/pi/issues/10549)).
|
||||
- `pi.tool-result` messages carry `durationMs`: how long `execute()` took in that attempt, measured with a monotonic clock. Calls that did not execute, and interrupted or aborted calls, have none ([#10549](https://github.com/earendil-works/pi/issues/10549)).
|
||||
- Assistant messages carry the `durationMs` pi-ai measures for each response ([#10549](https://github.com/earendil-works/pi/issues/10549)).
|
||||
- `openDurableObjectSqliteStorage(ctx.storage)` in `@earendil-works/pi-durable/storage/sqlite/cloudflare`: SQLite storage on a SQLite-backed Cloudflare Durable Object.
|
||||
- `settings.contextRetentionMs` (default ten minutes): how long an idle conversation keeps its last context read in memory; `0` drops it once idle.
|
||||
|
||||
### Changed
|
||||
|
||||
- Within one task invocation, `runtime.context()` reuses the range its previous read scanned and reads only entries committed since then, unless the head marker changed. A generation's request no longer rereads the whole transcript after its prepare, which roughly halves context reads per turn. The range is dropped when the invocation ends or sleeps.
|
||||
- `runtime.context()` keeps each conversation's last read range and derived view in memory: a later read by any of its tasks with the same head marker scans and derives only the entries committed since. A generation no longer rereads and re-derives the whole transcript for each model request. Busy conversations keep it; idle ones for `settings.contextRetentionMs`.
|
||||
|
||||
### Fixed
|
||||
|
||||
|
|
|
|||
|
|
@ -224,6 +224,7 @@ const harness = await Harness.open(storage, {
|
|||
compaction: { reserveTokens: 16384 },
|
||||
progress: { partialIntervalMs: 100, outputIntervalMs: 100 },
|
||||
toolExecution: "parallel",
|
||||
contextRetentionMs: 600_000, // idle conversations keep their context in memory this long (checked lazily where timers cannot be unreferenced)
|
||||
get followUpMode() {
|
||||
return userSettings.followUpMode;
|
||||
},
|
||||
|
|
@ -539,6 +540,7 @@ const usage = await harness.usage(context); // { models: { "openai/gpt-6-sol": U
|
|||
| Memory | `MemoryStorage` from the package root | Nothing is persisted. |
|
||||
| SQLite | `openNodeSqliteStorage(file)` from `@earendil-works/pi-durable/storage/sqlite/node` | One database file. WAL mode with `synchronous = NORMAL`: commits survive process crashes; the newest may be lost on power or host failure. |
|
||||
| JSONL | `openNodeJsonlStorage(directory, context)` from `@earendil-works/pi-durable/storage/jsonl/node` | Append-only files in one directory. Pass `{ fsync: true }` to flush before each commit marker. |
|
||||
| Cloudflare Durable Object | `openDurableObjectSqliteStorage(ctx.storage)` from `@earendil-works/pi-durable/storage/sqlite/cloudflare` | The SQLite storage of one SQLite-backed Durable Object; commits are its storage transactions. |
|
||||
|
||||
One process owns a storage at a time; there is no cross-process locking. The portable SQLite and JSONL cores (`/storage/sqlite`, `/storage/jsonl`) run without Node APIs, for example on Bun or in Cloudflare Durable Objects, given an asynchronous `SqliteDatabase` facade or a `FileSystem` from `@earendil-works/pi-durable/env`.
|
||||
|
||||
|
|
|
|||
|
|
@ -412,6 +412,8 @@ type HarnessSettings = {
|
|||
readonly toolExecution?: ToolExecutionMode;
|
||||
readonly steeringMode?: QueueMode;
|
||||
readonly followUpMode?: QueueMode;
|
||||
/** How long an idle conversation keeps its last context read in memory; busy ones always keep it. */
|
||||
readonly contextRetentionMs?: number;
|
||||
};
|
||||
|
||||
/** Resolved: every field over its built-in default, object fields merged. */
|
||||
|
|
@ -423,6 +425,7 @@ type Settings = {
|
|||
readonly toolExecution: ToolExecutionMode; // "parallel"
|
||||
readonly steeringMode: QueueMode; // "one-at-a-time"
|
||||
readonly followUpMode: QueueMode; // "one-at-a-time"
|
||||
readonly contextRetentionMs: number; // 600000
|
||||
};
|
||||
|
||||
/** Stored choices of one conversation; names, not objects. Unset fields follow the host. */
|
||||
|
|
@ -1806,8 +1809,15 @@ committed documents, for example to supply `PromptInput.read` (section 7.4).
|
|||
`getTask()` and `entry()` read committed records with one lookup each; `waitForTask()` waits for a terminal receipt,
|
||||
for example a child task created by a tool.
|
||||
`context()` captures its bounds on the Session line and derives the view from
|
||||
immutable entries off the line, like `Conversation.context()`. Like every runtime
|
||||
operation, these reject after the invocation ends.
|
||||
immutable entries off the line, like `Conversation.context()`. The scheduler keeps
|
||||
each conversation's last range and view in memory: a later read by any of its tasks
|
||||
with the same head marker scans and derives only the entries after that range's tail.
|
||||
A conversation keeps them while busy and for `contextRetentionMs` once idle.
|
||||
Expiry is checked on task changes and by one unreferenced timer, which never keeps
|
||||
the process alive; where timers cannot be unreferenced, as in Cloudflare Workers,
|
||||
only task changes check it. Kept entries are frozen and each read returns its own
|
||||
view arrays. Like every runtime operation, these
|
||||
reject after the invocation ends.
|
||||
|
||||
Reservation durably changes `pending`, or `waiting` once it may resume, to
|
||||
`running`. One invocation runs phase
|
||||
|
|
|
|||
|
|
@ -54,6 +54,12 @@
|
|||
"import": "./dist/storage/sqlite/index.js",
|
||||
"default": "./dist/storage/sqlite/index.js"
|
||||
},
|
||||
"./storage/sqlite/cloudflare": {
|
||||
"source": "./src/storage/sqlite/cloudflare.ts",
|
||||
"types": "./dist/storage/sqlite/cloudflare.d.ts",
|
||||
"import": "./dist/storage/sqlite/cloudflare.js",
|
||||
"default": "./dist/storage/sqlite/cloudflare.js"
|
||||
},
|
||||
"./storage/sqlite/node": {
|
||||
"source": "./src/storage/sqlite/node.ts",
|
||||
"types": "./dist/storage/sqlite/node.d.ts",
|
||||
|
|
|
|||
|
|
@ -49,6 +49,9 @@ export const AgentDoc = defineDoc<AgentState>({
|
|||
checkpointWhen: () => true,
|
||||
});
|
||||
|
||||
/** Default `settings.contextRetentionMs`: ten minutes. */
|
||||
const DEFAULT_CONTEXT_RETENTION_MS = 600_000;
|
||||
|
||||
/** Resolve the host settings: every field over its built-in default, object fields merged. */
|
||||
export function resolveSettings(settings: HarnessSettings | undefined): Settings {
|
||||
const extensions = settings?.extensions;
|
||||
|
|
@ -65,6 +68,7 @@ export function resolveSettings(settings: HarnessSettings | undefined): Settings
|
|||
toolExecution: settings?.toolExecution ?? "parallel",
|
||||
steeringMode: settings?.steeringMode ?? "one-at-a-time",
|
||||
followUpMode: settings?.followUpMode ?? "one-at-a-time",
|
||||
contextRetentionMs: settings?.contextRetentionMs ?? DEFAULT_CONTEXT_RETENTION_MS,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -40,18 +40,29 @@ export async function captureContextBounds(
|
|||
}
|
||||
|
||||
/**
|
||||
* Visible entries one context read scanned: from the head marker's head, or transcript start, through `bounds.tail`,
|
||||
* oldest first. Entries at or below the tail never change, so a later read with the same head marker extends it.
|
||||
* One context read: the visible entries it scanned, from the head marker's head, or transcript start, through
|
||||
* `bounds.tail`, oldest first, and the view derived from them. Entries at or below the tail never change, so a later
|
||||
* read with the same head marker extends both.
|
||||
*/
|
||||
export type ContextRange = {
|
||||
readonly conversationId: ConversationId;
|
||||
readonly bounds: ContextBounds;
|
||||
readonly entries: readonly EntryRecord[];
|
||||
readonly view: ContextView;
|
||||
/** Targets of the edits in `entries`. */
|
||||
readonly edited: ReadonlySet<EntryId>;
|
||||
/**
|
||||
* `view.messages` from the contributed messages before the last assistant message. Tool results are ordered within
|
||||
* the messages up to the next assistant message, so later entries cannot change these.
|
||||
*/
|
||||
readonly settled: readonly Message[];
|
||||
/** Contributed messages from the last assistant message on, before tool result ordering. */
|
||||
readonly open: readonly Message[];
|
||||
};
|
||||
|
||||
/**
|
||||
* `readContext()` that reuses `previous`: with the same conversation and head marker, only entries after its tail are
|
||||
* scanned. Returns the view and the range to pass to the next read.
|
||||
* `readContext()` that reuses `previous`, an earlier range of the same conversation: with the same head marker, only
|
||||
* entries after its tail are scanned. Returns the view and the range to pass to the next read. A range outlives the
|
||||
* read, so its entries are frozen like `MemoryStorage` records, and the returned view has its own arrays.
|
||||
*/
|
||||
export async function readContextFrom(
|
||||
session: SessionImpl,
|
||||
|
|
@ -63,21 +74,32 @@ export async function readContextFrom(
|
|||
): Promise<{ readonly view: ContextView; readonly range: ContextRange | undefined }> {
|
||||
const bounds = await session.readOnLine(() => captureContextBounds(storage, conversationId, context, at));
|
||||
if (bounds === undefined) return { view: emptyView(), range: undefined };
|
||||
let entries: readonly EntryRecord[];
|
||||
if (
|
||||
previous === undefined ||
|
||||
previous.conversationId !== conversationId ||
|
||||
previous.bounds.head?.id !== bounds.head?.id
|
||||
) {
|
||||
entries = await scanRange(storage, conversationId, rangeQuery(bounds), context);
|
||||
} else if (bounds.tail <= previous.bounds.tail) {
|
||||
entries = previous.entries.filter((entry) => entry.id <= bounds.tail);
|
||||
freezeJson(bounds.head);
|
||||
let range: ContextRange;
|
||||
if (previous === undefined || previous.bounds.head?.id !== bounds.head?.id) {
|
||||
range = deriveRange(bounds, freezeJson(await scanRange(storage, conversationId, rangeQuery(bounds), context)));
|
||||
} else if (bounds.tail === previous.bounds.tail) {
|
||||
range = previous;
|
||||
} else if (bounds.tail < previous.bounds.tail) {
|
||||
range = deriveRange(
|
||||
bounds,
|
||||
previous.entries.filter((entry) => entry.id <= bounds.tail),
|
||||
);
|
||||
} else {
|
||||
const minEntryId = idFromNumber<EntryId>(previous.bounds.tail + 1);
|
||||
const added = await scanRange(storage, conversationId, { minEntryId, maxEntryId: bounds.tail }, context);
|
||||
entries = [...previous.entries, ...added];
|
||||
range = extendRange(previous, bounds, freezeJson(added));
|
||||
}
|
||||
return { view: deriveView(bounds.head, entries), range: { conversationId, bounds, entries } };
|
||||
const { view } = range;
|
||||
return {
|
||||
view: {
|
||||
...view,
|
||||
entries: [...view.entries],
|
||||
contributions: [...view.contributions],
|
||||
messages: [...view.messages],
|
||||
},
|
||||
range,
|
||||
};
|
||||
}
|
||||
|
||||
/** Committed context of one conversation: bounds captured on the Session line, entries derived off it. */
|
||||
|
|
@ -105,29 +127,44 @@ export async function deriveContext(
|
|||
context: Context,
|
||||
): Promise<ContextView> {
|
||||
if (bounds === undefined) return emptyView();
|
||||
return deriveView(bounds.head, await scanRange(storage, conversationId, rangeQuery(bounds), context));
|
||||
return deriveRange(bounds, await scanRange(storage, conversationId, rangeQuery(bounds), context)).view;
|
||||
}
|
||||
|
||||
function emptyView(): ContextView {
|
||||
return { head: undefined, entries: [], contributions: [], messages: [] };
|
||||
}
|
||||
|
||||
/** Derive the context view from the scanned range of `head`. */
|
||||
function deriveView(head: ContextBounds["head"], range: readonly EntryRecord[]): ContextView {
|
||||
/** Derive the context view from the entries scanned within `bounds`. */
|
||||
function deriveRange(bounds: ContextBounds, entries: readonly EntryRecord[]): ContextRange {
|
||||
const edits = new Map<EntryId, ContextEdit>();
|
||||
// Edits of every entry in the range count, including older head markers that `selectActive()` drops.
|
||||
for (const entry of range) for (const edit of entry.edits ?? []) edits.set(edit.target, edit);
|
||||
for (const entry of entries) for (const edit of entry.edits ?? []) edits.set(edit.target, edit);
|
||||
const active = selectActive(bounds.head, entries);
|
||||
const contributions = active.map((entry) => contribute(entry, edits.get(entry.id)));
|
||||
const { settled, open } = settle([], contributions.flat());
|
||||
const messages = leadWithSystem([...settled, ...orderToolResults(open)]);
|
||||
const view = { head: bounds.head, entries: active, contributions, messages };
|
||||
return { bounds, entries, view, edited: new Set(edits.keys()), settled, open };
|
||||
}
|
||||
|
||||
const entries = selectActive(head, range);
|
||||
const contributions = entries.map((entry): Message[] => {
|
||||
const edit = edits.get(entry.id);
|
||||
if (edit?.action === "omit") return [];
|
||||
const contributed = edit?.action === "replace" ? edit.messages : (entry.model ?? []);
|
||||
return contributed.filter(
|
||||
(message) => message.role !== "assistant" || !EXCLUDED_STOP_REASONS.has(message.stopReason),
|
||||
);
|
||||
});
|
||||
return { head, entries, contributions, messages: leadWithSystem(orderToolResults(contributions.flat())) };
|
||||
/**
|
||||
* `previous` extended by the entries `added` after its tail under the same head marker: only their contributions and
|
||||
* the open messages are derived. An added edit can change an earlier entry, so it derives the whole range again.
|
||||
*/
|
||||
function extendRange(previous: ContextRange, bounds: ContextBounds, added: readonly EntryRecord[]): ContextRange {
|
||||
const entries = [...previous.entries, ...added];
|
||||
if (added.some((entry) => entry.edits !== undefined || entry.head !== undefined || previous.edited.has(entry.id))) {
|
||||
return deriveRange(bounds, entries);
|
||||
}
|
||||
const contributions = added.map((entry) => contribute(entry, undefined));
|
||||
const { settled, open } = settle(previous.settled, [...previous.open, ...contributions.flat()]);
|
||||
const view = {
|
||||
head: bounds.head,
|
||||
entries: [...previous.view.entries, ...added],
|
||||
contributions: [...previous.view.contributions, ...contributions],
|
||||
messages: leadWithSystem([...settled, ...orderToolResults(open)]),
|
||||
};
|
||||
return { bounds, entries, view, edited: previous.edited, settled, open };
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -142,6 +179,40 @@ function leadWithSystem(messages: Message[]): Message[] {
|
|||
return [messages[index]!, ...messages.slice(0, index), ...messages.slice(index + 1)];
|
||||
}
|
||||
|
||||
/** One active entry's model messages after its edit and excluded stop reasons, before tool result ordering. */
|
||||
function contribute(entry: EntryRecord, edit: ContextEdit | undefined): readonly Message[] {
|
||||
if (edit?.action === "omit") return Object.freeze([]);
|
||||
const contributed = edit?.action === "replace" ? edit.messages : (entry.model ?? []);
|
||||
return Object.freeze(
|
||||
contributed.filter((message) => message.role !== "assistant" || !EXCLUDED_STOP_REASONS.has(message.stopReason)),
|
||||
);
|
||||
}
|
||||
|
||||
/** Objects `freezeJson()` froze with everything in them. `Object.isFrozen()` holds for shallow freezes too. */
|
||||
const deeplyFrozen = new WeakSet<object>();
|
||||
|
||||
/** Freeze a JSON value and everything in it. */
|
||||
function freezeJson<T>(value: T): T {
|
||||
if (value === null || typeof value !== "object" || deeplyFrozen.has(value)) return value;
|
||||
for (const child of Object.values(value)) freezeJson(child);
|
||||
Object.freeze(value);
|
||||
deeplyFrozen.add(value);
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* Move the ordered messages before the last assistant message of `open` to `settled`. `orderToolResults()` of a
|
||||
* sequence equals the concatenation over its parts when each later part starts with an assistant message.
|
||||
*/
|
||||
function settle(
|
||||
settled: readonly Message[],
|
||||
open: readonly Message[],
|
||||
): { readonly settled: readonly Message[]; readonly open: readonly Message[] } {
|
||||
const last = open.findLastIndex((message) => message.role === "assistant");
|
||||
if (last <= 0) return { settled, open };
|
||||
return { settled: [...settled, ...orderToolResults(open.slice(0, last))], open: open.slice(last) };
|
||||
}
|
||||
|
||||
/** The raw active entries within captured bounds, without deriving model context. */
|
||||
export async function activeEntries(
|
||||
storage: Storage,
|
||||
|
|
@ -210,7 +281,7 @@ export function orderToolResults(messages: readonly Message[]): Message[] {
|
|||
}
|
||||
|
||||
function missingResult(call: ToolCall, timestamp: number): ToolResultMessage {
|
||||
return {
|
||||
return freezeJson({
|
||||
role: "toolResult",
|
||||
toolCallId: call.id,
|
||||
toolName: call.name,
|
||||
|
|
@ -218,5 +289,5 @@ function missingResult(call: ToolCall, timestamp: number): ToolResultMessage {
|
|||
isError: true,
|
||||
details: { reason: "missing_result" },
|
||||
timestamp,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -75,11 +75,6 @@ type Invocation = {
|
|||
readonly context: Context;
|
||||
/** Watches acquired through the runtime; stopped at invocation end. */
|
||||
readonly watches: Set<DocumentWatch<JsonObject>>;
|
||||
/**
|
||||
* Range of the invocation's last context read, so the next read, such as a generation's request after its prepare,
|
||||
* scans only newer entries. Derived and never persisted; dropped at invocation end and before sleeping.
|
||||
*/
|
||||
contextRange: ContextRange | undefined;
|
||||
ended: boolean;
|
||||
readonly done: Promise<void>;
|
||||
readonly finish: () => void;
|
||||
|
|
@ -191,6 +186,14 @@ export class TaskScheduler {
|
|||
readonly #conversation: TaskSchedulerOptions["conversation"];
|
||||
readonly #context: Context;
|
||||
readonly #live = new Map<TaskId, AnyTaskRecord>();
|
||||
/**
|
||||
* Context range last read through a task runtime, per conversation: a later read, by any of its tasks, scans only
|
||||
* newer entries. `idleSince` is the Harness time it was first seen idle. Derived and never persisted; dropped at the
|
||||
* first idle check after `settings.contextRetentionMs` of idleness, and at close.
|
||||
*/
|
||||
readonly #contexts = new Map<ConversationId, { range: ContextRange; idleSince: number | undefined }>();
|
||||
/** Timer for the earliest expiry of an idle context, only where timers can be unreferenced (see `#scheduleExpiry`). */
|
||||
#expiry: { readonly at: number; readonly timer: ReturnType<typeof setTimeout> } | undefined;
|
||||
readonly #invocations = new Map<TaskId, Invocation>();
|
||||
readonly #taskWaiters = new Waiters<TaskId, SettledTask<JsonValue>>();
|
||||
/** Idle waiters by conversation; `undefined` waits for the whole Harness. */
|
||||
|
|
@ -419,14 +422,62 @@ export class TaskScheduler {
|
|||
// Also retries, with the next commit of any kind, a cascade whose commit failed.
|
||||
if (this.#cascadePending) this.#scheduleReconcile();
|
||||
if (!changed) return;
|
||||
this.#resolveIdleWaiters();
|
||||
this.#settleIdle();
|
||||
this.#kick();
|
||||
}
|
||||
|
||||
#resolveIdleWaiters(): void {
|
||||
/** Resolve idle waiters, and drop each kept context whose conversation has been idle for the retention period. */
|
||||
#settleIdle(): void {
|
||||
for (const conversationId of this.#idleWaiters.keys()) {
|
||||
if (this.#idle(conversationId)) this.#idleWaiters.resolve(conversationId);
|
||||
}
|
||||
const now = this.#now();
|
||||
const retention = this.#settings().contextRetentionMs;
|
||||
for (const [conversationId, kept] of this.#contexts) {
|
||||
if (kept.idleSince !== undefined && now - kept.idleSince >= retention) {
|
||||
this.#contexts.delete(conversationId);
|
||||
} else if (!this.#idle(conversationId)) {
|
||||
kept.idleSince = undefined;
|
||||
} else if (retention > 0) {
|
||||
kept.idleSince ??= now;
|
||||
} else {
|
||||
this.#contexts.delete(conversationId);
|
||||
}
|
||||
}
|
||||
this.#scheduleExpiry();
|
||||
}
|
||||
|
||||
/**
|
||||
* Run `#settleIdle()` when the earliest idle context expires. The timer is unreferenced, so it never keeps the process
|
||||
* alive. Where timers cannot be unreferenced, as in Cloudflare Workers, none is kept: a pending timer could keep a
|
||||
* Durable Object from being evicted, and eviction frees the contexts. There, task changes alone check expiry.
|
||||
*/
|
||||
#scheduleExpiry(): void {
|
||||
const retention = this.#settings().contextRetentionMs;
|
||||
let at: number | undefined;
|
||||
for (const kept of this.#contexts.values()) {
|
||||
if (kept.idleSince !== undefined && (at === undefined || kept.idleSince + retention < at)) {
|
||||
at = kept.idleSince + retention;
|
||||
}
|
||||
}
|
||||
if (this.#expiry !== undefined && this.#expiry.at === at) return;
|
||||
if (this.#expiry !== undefined) clearTimeout(this.#expiry.timer);
|
||||
this.#expiry = undefined;
|
||||
if (at === undefined || this.#closing) return;
|
||||
const timer = setTimeout(
|
||||
() => {
|
||||
this.#expiry = undefined;
|
||||
this.#settleIdle();
|
||||
},
|
||||
Math.min(Math.max(0, at - this.#now()), MAX_TIMER_DELAY),
|
||||
);
|
||||
const unref = (timer as { unref?: unknown }).unref;
|
||||
if (typeof unref !== "function") {
|
||||
clearTimeout(timer);
|
||||
return;
|
||||
}
|
||||
unref.call(timer);
|
||||
this.#expiry = { at, timer };
|
||||
}
|
||||
|
||||
// ─── Ownership ───────────────────────────────────────────────────────────
|
||||
|
|
@ -481,7 +532,7 @@ export class TaskScheduler {
|
|||
for (const id of checks) this.#failFastChecks.add(id);
|
||||
if (!this.#closing) this.#report(error);
|
||||
}
|
||||
this.#resolveIdleWaiters();
|
||||
this.#settleIdle();
|
||||
}
|
||||
|
||||
/** Whether any of `ids` holds or ended with an outcome other than `completed`. */
|
||||
|
|
@ -674,6 +725,9 @@ export class TaskScheduler {
|
|||
const error = closedError();
|
||||
this.#taskWaiters.rejectAll(error);
|
||||
this.#idleWaiters.rejectAll(error);
|
||||
this.#contexts.clear();
|
||||
if (this.#expiry !== undefined) clearTimeout(this.#expiry.timer);
|
||||
this.#expiry = undefined;
|
||||
for (const invocation of this.#invocations.values()) invocation.controller.abort();
|
||||
}
|
||||
|
||||
|
|
@ -830,7 +884,6 @@ export class TaskScheduler {
|
|||
controller,
|
||||
context: withAbortSignal(controller.signal, this.#context),
|
||||
watches: new Set(),
|
||||
contextRange: undefined,
|
||||
ended: false,
|
||||
done,
|
||||
finish: () => finish(),
|
||||
|
|
@ -1034,7 +1087,6 @@ export class TaskScheduler {
|
|||
#end(invocation: Invocation): void {
|
||||
if (invocation.ended) return;
|
||||
invocation.ended = true;
|
||||
invocation.contextRange = undefined;
|
||||
if (this.#invocations.get(invocation.taskId) === invocation) this.#invocations.delete(invocation.taskId);
|
||||
for (const watch of invocation.watches) void watch.stop();
|
||||
// Pending waits bound to the invocation, such as a tool's waitForTask(), reject with it.
|
||||
|
|
@ -1182,9 +1234,24 @@ export class TaskScheduler {
|
|||
conversationId,
|
||||
context,
|
||||
options?.at,
|
||||
invocation.contextRange,
|
||||
this.#contexts.get(conversationId)?.range,
|
||||
);
|
||||
if (!invocation.ended) invocation.contextRange = range;
|
||||
// Keep it unless the invocation ended or a concurrent read already kept a newer range.
|
||||
const kept = this.#contexts.get(conversationId);
|
||||
if (
|
||||
range !== undefined &&
|
||||
!this.#closing &&
|
||||
!invocation.ended &&
|
||||
(kept === undefined || kept.range.bounds.tail <= range.bounds.tail)
|
||||
) {
|
||||
// A read of another, idle conversation starts or continues its retention period.
|
||||
if (!this.#idle(conversationId)) {
|
||||
this.#contexts.set(conversationId, { range, idleSince: undefined });
|
||||
} else if (this.#settings().contextRetentionMs > 0) {
|
||||
this.#contexts.set(conversationId, { range, idleSince: kept?.idleSince ?? this.#now() });
|
||||
this.#scheduleExpiry();
|
||||
}
|
||||
}
|
||||
return view;
|
||||
}),
|
||||
now: () => {
|
||||
|
|
@ -1232,8 +1299,6 @@ export class TaskScheduler {
|
|||
/** Wait until the Harness clock reaches `until`, rechecking it after every timer. */
|
||||
async #sleep(invocation: Invocation, until: number, context: Context): Promise<void> {
|
||||
if (invocation.ended) throw endedError(invocation);
|
||||
// A sleeping task holds no context range; a retry or poll can wait for a long time.
|
||||
invocation.contextRange = undefined;
|
||||
const signals = [invocation.controller.signal];
|
||||
if (context.abortSignal !== undefined) signals.push(context.abortSignal);
|
||||
const signal = AbortSignal.any(signals);
|
||||
|
|
|
|||
|
|
@ -421,6 +421,11 @@ export type HarnessSettings = {
|
|||
readonly toolExecution?: ToolExecutionMode;
|
||||
readonly steeringMode?: QueueMode;
|
||||
readonly followUpMode?: QueueMode;
|
||||
/**
|
||||
* How long an idle conversation keeps its last context read in memory, so its next run reads only newer entries.
|
||||
* Busy conversations always keep it; `0` drops it once the conversation is idle.
|
||||
*/
|
||||
readonly contextRetentionMs?: number;
|
||||
};
|
||||
|
||||
/** Resolved settings: every field over its built-in default, object fields merged. */
|
||||
|
|
@ -434,6 +439,7 @@ export type Settings = {
|
|||
readonly toolExecution: ToolExecutionMode;
|
||||
readonly steeringMode: QueueMode;
|
||||
readonly followUpMode: QueueMode;
|
||||
readonly contextRetentionMs: number;
|
||||
};
|
||||
|
||||
/** What `HarnessOptions.env` builds an environment for. */
|
||||
|
|
|
|||
139
packages/durable/src/storage/sqlite/cloudflare.ts
Normal file
139
packages/durable/src/storage/sqlite/cloudflare.ts
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
import type { SqliteDatabase, SqliteExecutor, SqliteValue } from "./database.ts";
|
||||
import { SqliteStorage } from "./storage.ts";
|
||||
|
||||
/** Values a Durable Object SQL binding accepts and returns. */
|
||||
type DurableObjectSqlValue = ArrayBuffer | string | number | null;
|
||||
|
||||
/**
|
||||
* The parts of a SQLite-backed Durable Object's `ctx.storage` this adapter uses, typed structurally so the package
|
||||
* needs no Workers type dependency.
|
||||
*/
|
||||
export type DurableObjectSqliteStorage = {
|
||||
readonly sql: {
|
||||
exec(query: string, ...bindings: DurableObjectSqlValue[]): { toArray(): Record<string, DurableObjectSqlValue>[] };
|
||||
};
|
||||
transaction<T>(closure: () => Promise<T>): Promise<T>;
|
||||
};
|
||||
|
||||
/** Runs operations one at a time in call order, so a transaction excludes every other operation. */
|
||||
class SerialQueue {
|
||||
private tail: Promise<unknown> = Promise.resolve();
|
||||
|
||||
run<T>(operation: () => T | Promise<T>): Promise<T> {
|
||||
const result = this.tail.then(operation);
|
||||
this.tail = result.catch(() => {});
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
/** Durable Object SQL binds numbers as doubles: a `bigint` outside the safe integer range would lose precision. */
|
||||
function bindValue(value: SqliteValue): DurableObjectSqlValue {
|
||||
if (value instanceof Uint8Array) return value.slice().buffer as ArrayBuffer;
|
||||
if (typeof value !== "bigint") return value;
|
||||
if (value < BigInt(Number.MIN_SAFE_INTEGER) || value > BigInt(Number.MAX_SAFE_INTEGER)) {
|
||||
throw new RangeError(`Durable Object SQL cannot bind ${value} without losing precision`);
|
||||
}
|
||||
return Number(value);
|
||||
}
|
||||
|
||||
const bind = (params: readonly SqliteValue[]): DurableObjectSqlValue[] => params.map(bindValue);
|
||||
|
||||
/** Rows as returned, with `BLOB` columns converted to `Uint8Array` in place; no other copy is made. */
|
||||
function rows<T>(rows: Record<string, DurableObjectSqlValue>[]): T[] {
|
||||
for (const row of rows) {
|
||||
for (const key in row) {
|
||||
const value = row[key];
|
||||
if (value instanceof ArrayBuffer) (row as Record<string, unknown>)[key] = new Uint8Array(value);
|
||||
}
|
||||
}
|
||||
return rows as T[];
|
||||
}
|
||||
|
||||
/** Executes SQL through `ctx.storage.sql`, which has no separate prepare step. */
|
||||
class DurableObjectSqliteExecutor implements SqliteExecutor {
|
||||
protected readonly storage: DurableObjectSqliteStorage;
|
||||
|
||||
constructor(storage: DurableObjectSqliteStorage) {
|
||||
this.storage = storage;
|
||||
}
|
||||
|
||||
async exec(sql: string): Promise<void> {
|
||||
this.check();
|
||||
this.storage.sql.exec(sql);
|
||||
}
|
||||
|
||||
async run(sql: string, ...params: SqliteValue[]): Promise<void> {
|
||||
this.check();
|
||||
this.storage.sql.exec(sql, ...bind(params));
|
||||
}
|
||||
|
||||
async get<T extends object>(sql: string, ...params: SqliteValue[]): Promise<T | undefined> {
|
||||
this.check();
|
||||
return rows<T>(this.storage.sql.exec(sql, ...bind(params)).toArray())[0];
|
||||
}
|
||||
|
||||
async all<T extends object>(sql: string, ...params: SqliteValue[]): Promise<T[]> {
|
||||
this.check();
|
||||
return rows<T>(this.storage.sql.exec(sql, ...bind(params)).toArray());
|
||||
}
|
||||
|
||||
/** Throws when this handle may no longer run SQL. */
|
||||
protected check(): void {}
|
||||
}
|
||||
|
||||
/** The handle a transaction callback receives; invalid once the callback settles. */
|
||||
class DurableObjectSqliteTransaction extends DurableObjectSqliteExecutor {
|
||||
active = true;
|
||||
|
||||
protected override check(): void {
|
||||
if (!this.active) throw new Error("SQLite transaction handle is no longer active");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* `SqliteDatabase` adapter for the SQLite storage of a Cloudflare Durable Object. Integers are JavaScript numbers;
|
||||
* binding a `bigint` outside the safe integer range throws.
|
||||
*/
|
||||
export class DurableObjectSqliteDatabase extends DurableObjectSqliteExecutor implements SqliteDatabase {
|
||||
private readonly queue = new SerialQueue();
|
||||
|
||||
override exec(sql: string): Promise<void> {
|
||||
return this.queue.run(() => super.exec(sql));
|
||||
}
|
||||
|
||||
override run(sql: string, ...params: SqliteValue[]): Promise<void> {
|
||||
return this.queue.run(() => super.run(sql, ...params));
|
||||
}
|
||||
|
||||
override get<T extends object>(sql: string, ...params: SqliteValue[]): Promise<T | undefined> {
|
||||
return this.queue.run(() => super.get<T>(sql, ...params));
|
||||
}
|
||||
|
||||
override all<T extends object>(sql: string, ...params: SqliteValue[]): Promise<T[]> {
|
||||
return this.queue.run(() => super.all<T>(sql, ...params));
|
||||
}
|
||||
|
||||
/** `ctx.storage.transaction()` commits when the callback resolves and rolls back when it rejects. */
|
||||
transaction<T>(callback: (transaction: SqliteExecutor) => Promise<T>): Promise<T> {
|
||||
return this.queue.run(() =>
|
||||
this.storage.transaction(async () => {
|
||||
const transaction = new DurableObjectSqliteTransaction(this.storage);
|
||||
try {
|
||||
return await callback(transaction);
|
||||
} finally {
|
||||
transaction.active = false;
|
||||
}
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
/** Waits for queued work. The Durable Object owns its storage, so nothing else closes. */
|
||||
close(): Promise<void> {
|
||||
return this.queue.run(() => {});
|
||||
}
|
||||
}
|
||||
|
||||
/** Open durable storage on a SQLite-backed Durable Object's `ctx.storage`. */
|
||||
export function openDurableObjectSqliteStorage(storage: DurableObjectSqliteStorage): Promise<SqliteStorage> {
|
||||
return SqliteStorage.open(new DurableObjectSqliteDatabase(storage));
|
||||
}
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import type { Message } from "@earendil-works/pi-ai";
|
||||
import type { Context } from "@earendil-works/chord";
|
||||
import { createModels, type Message } from "@earendil-works/pi-ai";
|
||||
import {
|
||||
type ContextView,
|
||||
createRegistry,
|
||||
|
|
@ -6,10 +7,12 @@ import {
|
|||
type EntryDraft,
|
||||
type EntryId,
|
||||
type EntryRecord,
|
||||
Harness,
|
||||
type HarnessSettings,
|
||||
MemoryStorage,
|
||||
type Storage,
|
||||
} from "@earendil-works/pi-durable";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { addTask, assistant, describeMessage, openHarness, system, toolResult, user } from "./harness-support.ts";
|
||||
import { context } from "./session-support.ts";
|
||||
|
||||
|
|
@ -22,10 +25,52 @@ async function setup() {
|
|||
return { harness, root, append, message };
|
||||
}
|
||||
|
||||
/** A Harness over MemoryStorage that counts scanned entry rows, with a 21-entry root transcript. */
|
||||
async function countingSetup(options: { readonly settings?: HarnessSettings; readonly now?: () => number } = {}) {
|
||||
const scanned = { rows: 0 };
|
||||
/** When set, the next range scan waits for it; bounds probes (one row, on the Session line) do not. */
|
||||
const hold: { scan?: Promise<void> } = {};
|
||||
const memory = new MemoryStorage();
|
||||
const storage = new Proxy<Storage>(memory, {
|
||||
get(target, key) {
|
||||
const value: unknown = Reflect.get(target, key, target);
|
||||
if (key === "scanEntries") {
|
||||
return async (...args: Parameters<Storage["scanEntries"]>) => {
|
||||
const held = args[1] > 1 ? hold.scan : undefined;
|
||||
if (held !== undefined) hold.scan = undefined;
|
||||
await held;
|
||||
const page = await target.scanEntries(...args);
|
||||
scanned.rows += page.items.length;
|
||||
return page;
|
||||
};
|
||||
}
|
||||
return typeof value === "function" ? value.bind(target) : value;
|
||||
},
|
||||
});
|
||||
const registry = createRegistry();
|
||||
const harness = await Harness.open(storage, { models: createModels(), registry, ...options }, context);
|
||||
const root = await harness.root(context);
|
||||
const first = await root.commit(
|
||||
(tx) => tx.appendEntry(root.id, { kind: "message", model: [user("first")] }),
|
||||
context,
|
||||
);
|
||||
for (let index = 0; index < 20; index++) {
|
||||
await root.commit(
|
||||
(tx) => tx.appendEntry(root.id, { kind: "message", model: [assistant(`old ${index}`)] }),
|
||||
context,
|
||||
);
|
||||
}
|
||||
return { harness, registry, root, first, scanned, hold };
|
||||
}
|
||||
|
||||
function ids(entries: readonly EntryRecord[]): EntryId[] {
|
||||
return entries.map((entry) => entry.id);
|
||||
}
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
describe("conversation context", () => {
|
||||
it("returns the whole transcript without a head and excludes model-less entries from messages", async () => {
|
||||
const { root, append, message } = await setup();
|
||||
|
|
@ -206,29 +251,8 @@ describe("conversation context", () => {
|
|||
await expect(other.context(context, { at: tail.id })).rejects.toThrow("is not visible");
|
||||
});
|
||||
|
||||
it("extends a task invocation's context read with only newer entries", async () => {
|
||||
let scanned = 0;
|
||||
const memory = new MemoryStorage();
|
||||
const storage = new Proxy<Storage>(memory, {
|
||||
get(target, key) {
|
||||
const value: unknown = Reflect.get(target, key, target);
|
||||
if (key === "scanEntries") {
|
||||
return async (...args: Parameters<Storage["scanEntries"]>) => {
|
||||
const page = await target.scanEntries(...args);
|
||||
scanned += page.items.length;
|
||||
return page;
|
||||
};
|
||||
}
|
||||
return typeof value === "function" ? value.bind(target) : value;
|
||||
},
|
||||
});
|
||||
const registry = createRegistry();
|
||||
const { harness } = await openHarness(storage, [], { registry });
|
||||
const root = await harness.root(context);
|
||||
const append = (draft: EntryDraft): Promise<EntryRecord> =>
|
||||
root.commit((tx) => tx.appendEntry(root.id, draft), context);
|
||||
const first = await append({ kind: "message", model: [user("first")] });
|
||||
for (let index = 0; index < 20; index++) await append({ kind: "message", model: [assistant(`old ${index}`)] });
|
||||
it("extends a task's context read with only newer entries", async () => {
|
||||
const { harness, registry, root, first, scanned } = await countingSetup();
|
||||
const Reads = defineTask<Record<string, never>, { phase: "run" }, null>({
|
||||
name: "test.context-reads",
|
||||
version: 1,
|
||||
|
|
@ -241,9 +265,9 @@ describe("conversation context", () => {
|
|||
return undefined;
|
||||
}, taskContext);
|
||||
const read = async (at?: EntryId): Promise<{ view: ContextView; rows: number }> => {
|
||||
const before = scanned;
|
||||
const before = scanned.rows;
|
||||
const view = await runtime.context(root.id, taskContext, { at });
|
||||
return { view, rows: scanned - before };
|
||||
return { view, rows: scanned.rows - before };
|
||||
};
|
||||
const initial = await read();
|
||||
expect(initial.view).toEqual(await root.context(taskContext));
|
||||
|
|
@ -294,4 +318,395 @@ describe("conversation context", () => {
|
|||
expect(settled.state.outcome).toEqual({ status: "completed", result: null });
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("keeps a conversation's context read across its tasks and for the retention period once idle", async () => {
|
||||
let clock = 1_000;
|
||||
const { harness, registry, root, scanned } = await countingSetup({ now: () => clock });
|
||||
const rows: Record<string, number> = {};
|
||||
const read = async (name: string, load: () => Promise<ContextView>, taskContext: Context) => {
|
||||
const before = scanned.rows;
|
||||
const view = await load();
|
||||
rows[name] = scanned.rows - before;
|
||||
expect(view).toEqual(await root.context(taskContext));
|
||||
};
|
||||
const Child = defineTask<Record<string, never>, { phase: "run" }, null>({
|
||||
name: "test.context-child",
|
||||
version: 1,
|
||||
initial: () => ({ phase: "run" }),
|
||||
phases: {
|
||||
run: async (_task, runtime, taskContext) => {
|
||||
await runtime.commit(async (tx) => {
|
||||
await tx.appendEntry(root.id, { kind: "message", model: [user("from child")] });
|
||||
return undefined;
|
||||
}, taskContext);
|
||||
await read("child", () => runtime.context(root.id, taskContext), taskContext);
|
||||
await runtime.commit(() => DONE, taskContext);
|
||||
},
|
||||
},
|
||||
abort: async (_task, runtime, taskContext) => {
|
||||
await runtime.commit(() => ABORTED, taskContext);
|
||||
},
|
||||
});
|
||||
const Parent = defineTask<Record<string, never>, { phase: "start" } | { phase: "after" }, null>({
|
||||
name: "test.context-parent",
|
||||
version: 1,
|
||||
initial: () => ({ phase: "start" }),
|
||||
phases: {
|
||||
start: async (_task, runtime, taskContext) => {
|
||||
await read("parent", () => runtime.context(root.id, taskContext), taskContext);
|
||||
await runtime.commit(async (tx) => {
|
||||
const child = await tx.createTask(
|
||||
Child,
|
||||
{},
|
||||
{ ownership: { kind: "task", taskId: runtime.taskId }, conversationId: root.id },
|
||||
);
|
||||
return { status: "waiting", checkpoint: { phase: "after" }, on: [child], policy: "allSettled" };
|
||||
}, taskContext);
|
||||
},
|
||||
after: async (_task, runtime, taskContext) => {
|
||||
await read("after", () => runtime.context(root.id, taskContext), taskContext);
|
||||
await runtime.commit(() => DONE, taskContext);
|
||||
},
|
||||
},
|
||||
abort: async (_task, runtime, taskContext) => {
|
||||
await runtime.commit(() => ABORTED, taskContext);
|
||||
},
|
||||
});
|
||||
const Probe = probeTask(read);
|
||||
for (const task of [Child, Parent, Probe]) addTask(registry, task);
|
||||
harness.resume();
|
||||
const options = { ownership: { kind: "conversation" }, conversationId: root.id } as const;
|
||||
const parent = await harness.commit((tx) => tx.createTask(Parent, {}, options), context);
|
||||
expect((await harness.waitForTask(parent, context)).state.outcome).toEqual(DONE.outcome);
|
||||
await root.waitForIdle(context);
|
||||
const probe = async (name: string) => {
|
||||
const id = await harness.commit((tx) => tx.createTask(Probe, { name }, options), context);
|
||||
expect((await harness.waitForTask(id, context)).state.outcome).toEqual(DONE.outcome);
|
||||
await root.waitForIdle(context);
|
||||
};
|
||||
await probe("idle");
|
||||
clock += 600_000;
|
||||
await probe("expired");
|
||||
// Bounds probes read one row each. The child scans its own new entry; the waiting parent's later invocation and
|
||||
// the probe within the retention period scan nothing new. After it, the probe reads the whole transcript again.
|
||||
expect(rows).toEqual({ parent: 22, child: 2, after: 1, idle: 1, expired: 23 });
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("derives the same view from extended reads as from whole reads", async () => {
|
||||
const { harness, registry, root } = await countingSetup();
|
||||
// Deterministic pseudo-random transcript (32-bit LCG, high bits): calls and results in any order, missing and
|
||||
// stray results, excluded stop reasons, system entries, notes, edits of earlier entries, and head markers.
|
||||
let seed = 7;
|
||||
const random = (n: number) => {
|
||||
seed = (Math.imul(seed, 1_664_525) + 1_013_904_223) >>> 0;
|
||||
return Math.floor((seed / 2 ** 32) * n);
|
||||
};
|
||||
const seen = new Map<string, number>();
|
||||
const count = (category: string) => seen.set(category, (seen.get(category) ?? 0) + 1);
|
||||
const Steps = defineTask<Record<string, never>, { phase: "run" }, null>({
|
||||
name: "test.context-steps",
|
||||
version: 1,
|
||||
initial: () => ({ phase: "run" }),
|
||||
phases: {
|
||||
run: async (_task, runtime, taskContext) => {
|
||||
const written: EntryRecord[] = [];
|
||||
const calls: string[] = [];
|
||||
let next = 0;
|
||||
const draft = (): EntryDraft => {
|
||||
const pick = random(14);
|
||||
if (pick < 3) {
|
||||
count("user");
|
||||
return { kind: "message", model: [user(`u${next++}`)] };
|
||||
}
|
||||
if (pick < 5) {
|
||||
const ids = Array.from({ length: random(3) }, () => `c${next++}`);
|
||||
calls.push(...ids);
|
||||
const stopReason = (["aborted", "error", "deferred"] as const)[random(9)];
|
||||
count(stopReason ?? "assistant");
|
||||
return {
|
||||
kind: "message",
|
||||
model: [assistant(`a${next++}`, { calls: ids, ...(stopReason ? { stopReason } : {}) })],
|
||||
};
|
||||
}
|
||||
if (pick < 8) {
|
||||
const matched = calls.length > 0 && random(5) > 0;
|
||||
count(matched ? "result" : "stray result");
|
||||
const id = matched ? calls.splice(random(calls.length), 1)[0]! : `stray${next++}`;
|
||||
return { kind: "message", model: [toolResult(id)] };
|
||||
}
|
||||
if (pick < 9) {
|
||||
count("system");
|
||||
return { kind: "message", model: [system({ s: `v${next++}` })] };
|
||||
}
|
||||
if (pick < 10) {
|
||||
count("note");
|
||||
return { kind: "note", data: { n: next++ } };
|
||||
}
|
||||
const target = written[random(written.length)];
|
||||
if (pick < 12 && target !== undefined) {
|
||||
const omit = random(2) === 0;
|
||||
count(omit ? "omit" : "replace");
|
||||
const action = omit
|
||||
? { action: "omit" as const }
|
||||
: { action: "replace" as const, messages: [user(`r${next++}`)] };
|
||||
return { kind: "edit", edits: [{ target: target.id, ...action }] };
|
||||
}
|
||||
if (pick < 13 && target !== undefined) {
|
||||
count("head");
|
||||
return { kind: "summary", head: target.id, model: [user(`h${next++}`)] };
|
||||
}
|
||||
count("user");
|
||||
return { kind: "message", model: [user(`u${next++}`)] };
|
||||
};
|
||||
const append = (conversationId: typeof root.id) =>
|
||||
runtime.commit(async (tx) => {
|
||||
const added = 1 + random(3);
|
||||
for (let index = 0; index < added; index++) {
|
||||
written.push(await tx.appendEntry(conversationId, draft()));
|
||||
}
|
||||
return undefined;
|
||||
}, taskContext);
|
||||
for (let step = 0; step < 150; step++) {
|
||||
await append(root.id);
|
||||
expect(await runtime.context(root.id, taskContext)).toEqual(await root.context(taskContext));
|
||||
if (step % 10 === 9) {
|
||||
// Overlapping reads, one with an earlier cutoff, share the kept range.
|
||||
const at = written[random(written.length)]!.id;
|
||||
const [whole, cut] = await Promise.all([
|
||||
runtime.context(root.id, taskContext),
|
||||
runtime.context(root.id, taskContext, { at }),
|
||||
]);
|
||||
expect(whole).toEqual(await root.context(taskContext));
|
||||
expect(cut).toEqual(await root.context(taskContext, { at }));
|
||||
}
|
||||
}
|
||||
// A fork sees its parent's entries through the fork point and extends with its own.
|
||||
const fork = await root.fork(
|
||||
written[random(written.length)]!.id,
|
||||
{ ownership: { kind: "ownerless" } },
|
||||
taskContext,
|
||||
);
|
||||
for (let step = 0; step < 20; step++) {
|
||||
expect(await runtime.context(fork.id, taskContext)).toEqual(await fork.context(taskContext));
|
||||
await append(fork.id);
|
||||
}
|
||||
await runtime.commit(() => DONE, taskContext);
|
||||
},
|
||||
},
|
||||
abort: async (_task, runtime, taskContext) => {
|
||||
await runtime.commit(() => ABORTED, taskContext);
|
||||
},
|
||||
});
|
||||
addTask(registry, Steps);
|
||||
harness.resume();
|
||||
const id = await harness.commit(
|
||||
(tx) => tx.createTask(Steps, {}, { ownership: { kind: "conversation" }, conversationId: root.id }),
|
||||
context,
|
||||
);
|
||||
expect((await harness.waitForTask(id, context)).state.outcome).toEqual(DONE.outcome);
|
||||
const categories = [
|
||||
"user",
|
||||
"assistant",
|
||||
"aborted",
|
||||
"error",
|
||||
"deferred",
|
||||
"result",
|
||||
"stray result",
|
||||
"system",
|
||||
"note",
|
||||
];
|
||||
for (const category of [...categories, "omit", "replace", "head"]) expect(seen.get(category)).toBeGreaterThan(0);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("keeps returned views and kept ranges apart", async () => {
|
||||
const { harness, registry, root } = await countingSetup();
|
||||
const Mutate = defineTask<Record<string, never>, { phase: "run" }, null>({
|
||||
name: "test.context-mutate",
|
||||
version: 1,
|
||||
initial: () => ({ phase: "run" }),
|
||||
phases: {
|
||||
run: async (_task, runtime, taskContext) => {
|
||||
const view = await runtime.context(root.id, taskContext);
|
||||
(view.messages as Message[]).length = 0;
|
||||
(view.entries as EntryRecord[]).pop();
|
||||
// Kept entries are frozen like MemoryStorage records.
|
||||
expect(() => {
|
||||
(view.contributions[0]![0] as { content: unknown }).content = "changed";
|
||||
}).toThrow(TypeError);
|
||||
expect(await runtime.context(root.id, taskContext)).toEqual(await root.context(taskContext));
|
||||
await runtime.commit(() => DONE, taskContext);
|
||||
},
|
||||
},
|
||||
abort: async (_task, runtime, taskContext) => {
|
||||
await runtime.commit(() => ABORTED, taskContext);
|
||||
},
|
||||
});
|
||||
addTask(registry, Mutate);
|
||||
harness.resume();
|
||||
const id = await harness.commit(
|
||||
(tx) => tx.createTask(Mutate, {}, { ownership: { kind: "conversation" }, conversationId: root.id }),
|
||||
context,
|
||||
);
|
||||
expect((await harness.waitForTask(id, context)).state.outcome).toEqual(DONE.outcome);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("freezes kept entries deeply even when storage freezes them shallowly", async () => {
|
||||
const memory = new MemoryStorage();
|
||||
// Fresh copies with only the record itself frozen; their messages stay mutable.
|
||||
const storage = new Proxy<Storage>(memory, {
|
||||
get(target, key) {
|
||||
const value: unknown = Reflect.get(target, key, target);
|
||||
if (key === "scanEntries") {
|
||||
return async (...args: Parameters<Storage["scanEntries"]>) => {
|
||||
const page = await target.scanEntries(...args);
|
||||
return { ...page, items: page.items.map((item) => Object.freeze(structuredClone(item))) };
|
||||
};
|
||||
}
|
||||
return typeof value === "function" ? value.bind(target) : value;
|
||||
},
|
||||
});
|
||||
const registry = createRegistry();
|
||||
const harness = await Harness.open(storage, { models: createModels(), registry }, context);
|
||||
const root = await harness.root(context);
|
||||
await root.commit((tx) => tx.appendEntry(root.id, { kind: "message", model: [user("hello")] }), context);
|
||||
const Mutate = defineTask<Record<string, never>, { phase: "run" }, null>({
|
||||
name: "test.context-shallow",
|
||||
version: 1,
|
||||
initial: () => ({ phase: "run" }),
|
||||
phases: {
|
||||
run: async (_task, runtime, taskContext) => {
|
||||
const view = await runtime.context(root.id, taskContext);
|
||||
expect(() => {
|
||||
(view.messages[0] as { content: unknown }).content = "changed";
|
||||
}).toThrow(TypeError);
|
||||
await runtime.commit(() => DONE, taskContext);
|
||||
},
|
||||
},
|
||||
abort: async (_task, runtime, taskContext) => {
|
||||
await runtime.commit(() => ABORTED, taskContext);
|
||||
},
|
||||
});
|
||||
addTask(registry, Mutate);
|
||||
harness.resume();
|
||||
const id = await harness.commit(
|
||||
(tx) => tx.createTask(Mutate, {}, { ownership: { kind: "conversation" }, conversationId: root.id }),
|
||||
context,
|
||||
);
|
||||
expect((await harness.waitForTask(id, context)).state.outcome).toEqual(DONE.outcome);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("does not keep a read that finishes after its task ended", async () => {
|
||||
const { harness, registry, root, scanned, hold } = await countingSetup({ settings: { contextRetentionMs: 0 } });
|
||||
const scans = Promise.withResolvers<void>();
|
||||
const probeStart = Promise.withResolvers<void>();
|
||||
let late: Promise<unknown> | undefined;
|
||||
const Late = defineTask<Record<string, never>, { phase: "run" }, null>({
|
||||
name: "test.context-late",
|
||||
version: 1,
|
||||
initial: () => ({ phase: "run" }),
|
||||
phases: {
|
||||
run: async (_task, runtime, taskContext) => {
|
||||
hold.scan = scans.promise;
|
||||
late = runtime.context(root.id, taskContext).catch((error: unknown) => error);
|
||||
await runtime.commit(() => DONE, taskContext);
|
||||
},
|
||||
},
|
||||
abort: async (_task, runtime, taskContext) => {
|
||||
await runtime.commit(() => ABORTED, taskContext);
|
||||
},
|
||||
});
|
||||
const rows: Record<string, number> = {};
|
||||
const Probe = probeTask(async (name, load) => {
|
||||
await probeStart.promise;
|
||||
const before = scanned.rows;
|
||||
await load();
|
||||
rows[name] = scanned.rows - before;
|
||||
});
|
||||
for (const task of [Late, Probe]) addTask(registry, task);
|
||||
harness.resume();
|
||||
const options = { ownership: { kind: "conversation" }, conversationId: root.id } as const;
|
||||
const lateId = await harness.commit((tx) => tx.createTask(Late, {}, options), context);
|
||||
expect((await harness.waitForTask(lateId, context)).state.outcome).toEqual(DONE.outcome);
|
||||
// The conversation is busy again when the late read finishes; its range must not be kept.
|
||||
const probeId = await harness.commit((tx) => tx.createTask(Probe, { name: "probe" }, options), context);
|
||||
scans.resolve();
|
||||
await late;
|
||||
probeStart.resolve();
|
||||
expect((await harness.waitForTask(probeId, context)).state.outcome).toEqual(DONE.outcome);
|
||||
expect(rows).toEqual({ probe: 22 });
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("drops an idle conversation's context read on a timer when nothing else runs", async () => {
|
||||
vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] });
|
||||
const { harness, registry, root, scanned } = await countingSetup({ settings: { contextRetentionMs: 1_000 } });
|
||||
const rows: Record<string, number> = {};
|
||||
const Probe = probeTask(async (name, load) => {
|
||||
const before = scanned.rows;
|
||||
await load();
|
||||
rows[name] = scanned.rows - before;
|
||||
});
|
||||
addTask(registry, Probe);
|
||||
harness.resume();
|
||||
const options = { ownership: { kind: "conversation" }, conversationId: root.id } as const;
|
||||
const probe = async (name: string) => {
|
||||
const id = await harness.commit((tx) => tx.createTask(Probe, { name }, options), context);
|
||||
expect((await harness.waitForTask(id, context)).state.outcome).toEqual(DONE.outcome);
|
||||
await root.waitForIdle(context);
|
||||
};
|
||||
const idle = vi.getTimerCount();
|
||||
await probe("first");
|
||||
expect(vi.getTimerCount()).toBe(idle + 1);
|
||||
vi.advanceTimersByTime(1_000);
|
||||
expect(vi.getTimerCount()).toBe(idle);
|
||||
await probe("second");
|
||||
expect(rows).toEqual({ first: 22, second: 22 });
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("drops a conversation's context read once idle with zero retention", async () => {
|
||||
const { harness, registry, root, scanned } = await countingSetup({ settings: { contextRetentionMs: 0 } });
|
||||
const rows: Record<string, number> = {};
|
||||
const Probe = probeTask(async (name, load) => {
|
||||
const before = scanned.rows;
|
||||
await load();
|
||||
rows[name] = scanned.rows - before;
|
||||
});
|
||||
addTask(registry, Probe);
|
||||
harness.resume();
|
||||
const options = { ownership: { kind: "conversation" }, conversationId: root.id } as const;
|
||||
for (const name of ["first", "second"]) {
|
||||
const id = await harness.commit((tx) => tx.createTask(Probe, { name }, options), context);
|
||||
expect((await harness.waitForTask(id, context)).state.outcome).toEqual(DONE.outcome);
|
||||
await root.waitForIdle(context);
|
||||
}
|
||||
expect(rows).toEqual({ first: 22, second: 22 });
|
||||
await harness.close(context);
|
||||
});
|
||||
});
|
||||
|
||||
const DONE = { status: "terminal", outcome: { status: "completed", result: null } } as const;
|
||||
|
||||
const ABORTED = { status: "terminal", outcome: { status: "aborted" } } as const;
|
||||
|
||||
/** A task that reads its conversation's context once, named by its input. */
|
||||
function probeTask(read: (name: string, load: () => Promise<ContextView>, taskContext: Context) => Promise<void>) {
|
||||
return defineTask<{ name: string }, { phase: "run" }, null>({
|
||||
name: "test.context-probe",
|
||||
version: 1,
|
||||
initial: () => ({ phase: "run" }),
|
||||
phases: {
|
||||
run: async (task, runtime, taskContext) => {
|
||||
await read(task.input.name, () => runtime.context(runtime.conversationId, taskContext), taskContext);
|
||||
await runtime.commit(() => DONE, taskContext);
|
||||
},
|
||||
},
|
||||
abort: async (_task, runtime, taskContext) => {
|
||||
await runtime.commit(() => ABORTED, taskContext);
|
||||
},
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -493,6 +493,7 @@ describe("generation", () => {
|
|||
toolExecution: "parallel",
|
||||
steeringMode: "one-at-a-time",
|
||||
followUpMode: "one-at-a-time",
|
||||
contextRetentionMs: 600_000,
|
||||
});
|
||||
expect(resolveSettings({ retry: { enabled: false }, compaction: { backgroundTokens: 0 } })).toMatchObject({
|
||||
retry: { enabled: false, maxRetries: 3, baseDelayMs: 2000 },
|
||||
|
|
|
|||
130
packages/durable/test/sqlite-cloudflare.test.ts
Normal file
130
packages/durable/test/sqlite-cloudflare.test.ts
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
import { DatabaseSync } from "node:sqlite";
|
||||
import { registerStorageConformance } from "@earendil-works/pi-durable/testing";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
DurableObjectSqliteDatabase,
|
||||
type DurableObjectSqliteStorage,
|
||||
openDurableObjectSqliteStorage,
|
||||
} from "../src/storage/sqlite/cloudflare.ts";
|
||||
|
||||
type SqlValue = ArrayBuffer | string | number | null;
|
||||
|
||||
/**
|
||||
* A stand-in for a SQLite-backed Durable Object's `ctx.storage` over `node:sqlite`, with its value types: `BLOB`
|
||||
* bindings and results are `ArrayBuffer`s, and `exec()` without bindings may run several statements.
|
||||
*/
|
||||
function durableObjectStorage(database = new DatabaseSync(":memory:")): DurableObjectSqliteStorage {
|
||||
const toNode = (value: SqlValue) => (value instanceof ArrayBuffer ? new Uint8Array(value) : value);
|
||||
const fromNode = (value: unknown): SqlValue =>
|
||||
value instanceof Uint8Array ? (value.slice().buffer as ArrayBuffer) : (value as SqlValue);
|
||||
return {
|
||||
sql: {
|
||||
exec(query, ...bindings) {
|
||||
if (bindings.length === 0 && query.trim().replace(/;\s*$/, "").includes(";")) {
|
||||
database.exec(query);
|
||||
return { toArray: () => [] };
|
||||
}
|
||||
const statement = database.prepare(query);
|
||||
if (statement.columns().length === 0) {
|
||||
statement.run(...bindings.map(toNode));
|
||||
return { toArray: () => [] };
|
||||
}
|
||||
const rows = statement.all(...bindings.map(toNode)) as Record<string, unknown>[];
|
||||
return {
|
||||
toArray: () =>
|
||||
rows.map((row) =>
|
||||
Object.fromEntries(Object.entries(row).map(([key, value]) => [key, fromNode(value)])),
|
||||
),
|
||||
};
|
||||
},
|
||||
},
|
||||
async transaction(closure) {
|
||||
database.exec("BEGIN");
|
||||
try {
|
||||
const result = await closure();
|
||||
database.exec("COMMIT");
|
||||
return result;
|
||||
} catch (error) {
|
||||
database.exec("ROLLBACK");
|
||||
throw error;
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
registerStorageConformance({ describe, expect, it }, "SqliteStorage on a Durable Object", async (use) =>
|
||||
use(await openDurableObjectSqliteStorage(durableObjectStorage())),
|
||||
);
|
||||
|
||||
describe("Durable Object SQLite adapter", () => {
|
||||
it("binds and returns BLOB values as Uint8Array", async () => {
|
||||
const database = new DurableObjectSqliteDatabase(durableObjectStorage());
|
||||
await database.exec("CREATE TABLE blobs (id INTEGER PRIMARY KEY, data BLOB, n INTEGER)");
|
||||
await database.run("INSERT INTO blobs (id, data, n) VALUES (?, ?, ?)", 1, new Uint8Array([1, 2, 3]), 7n);
|
||||
expect(await database.get("SELECT data, n FROM blobs WHERE id = ?", 1)).toEqual({
|
||||
data: new Uint8Array([1, 2, 3]),
|
||||
n: 7,
|
||||
});
|
||||
});
|
||||
|
||||
it("rejects a bigint binding outside the safe integer range", async () => {
|
||||
const database = new DurableObjectSqliteDatabase(durableObjectStorage());
|
||||
await database.exec("CREATE TABLE numbers (n INTEGER)");
|
||||
await expect(database.run("INSERT INTO numbers (n) VALUES (?)", 9_007_199_254_740_993n)).rejects.toThrow(
|
||||
RangeError,
|
||||
);
|
||||
expect(await database.all("SELECT n FROM numbers")).toEqual([]);
|
||||
});
|
||||
|
||||
it("closes after an active transaction settles", async () => {
|
||||
const database = new DurableObjectSqliteDatabase(durableObjectStorage());
|
||||
await database.exec("CREATE TABLE items (id INTEGER PRIMARY KEY)");
|
||||
const order: string[] = [];
|
||||
const { promise: gate, resolve } = Promise.withResolvers<void>();
|
||||
const transaction = database.transaction(async (handle) => {
|
||||
await handle.run("INSERT INTO items (id) VALUES (?)", 1);
|
||||
await gate;
|
||||
order.push("transaction");
|
||||
});
|
||||
const closed = database.close().then(() => order.push("close"));
|
||||
resolve();
|
||||
await Promise.all([transaction, closed]);
|
||||
expect(order).toEqual(["transaction", "close"]);
|
||||
});
|
||||
|
||||
it("rolls back a rejected transaction and rejects its handle afterwards", async () => {
|
||||
const database = new DurableObjectSqliteDatabase(durableObjectStorage());
|
||||
await database.exec("CREATE TABLE items (id INTEGER PRIMARY KEY)");
|
||||
let handle: Parameters<Parameters<DurableObjectSqliteDatabase["transaction"]>[0]>[0] | undefined;
|
||||
await expect(
|
||||
database.transaction(async (transaction) => {
|
||||
handle = transaction;
|
||||
await transaction.run("INSERT INTO items (id) VALUES (?)", 1);
|
||||
throw new Error("abort");
|
||||
}),
|
||||
).rejects.toThrow("abort");
|
||||
expect(await database.all("SELECT id FROM items")).toEqual([]);
|
||||
await expect(handle!.run("INSERT INTO items (id) VALUES (?)", 2)).rejects.toThrow("no longer active");
|
||||
});
|
||||
|
||||
it("queues operations behind an active transaction", async () => {
|
||||
const database = new DurableObjectSqliteDatabase(durableObjectStorage());
|
||||
await database.exec("CREATE TABLE items (id INTEGER PRIMARY KEY)");
|
||||
const order: string[] = [];
|
||||
const { promise: gate, resolve } = Promise.withResolvers<void>();
|
||||
const transaction = database.transaction(async (handle) => {
|
||||
await handle.run("INSERT INTO items (id) VALUES (?)", 1);
|
||||
await gate;
|
||||
order.push("transaction");
|
||||
});
|
||||
const read = database.all("SELECT id FROM items").then((rows) => {
|
||||
order.push("read");
|
||||
return rows;
|
||||
});
|
||||
resolve();
|
||||
await transaction;
|
||||
expect(await read).toEqual([{ id: 1 }]);
|
||||
expect(order).toEqual(["transaction", "read"]);
|
||||
await database.close();
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue