Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
8.0 kB · 280 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281import { mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";import { join } from "node:path";import type { AgentMessage } from "@earendil-works/pi-agent-core";
export const SNAPSHOT_TARGET_CHUNK_BYTES = 512 * 1024;export const SNAPSHOT_MEMORY_CACHE_BYTES = 4 * 1024 * 1024;
interface SnapshotTranscriptChunk { buffer?: Buffer; path?: string;}
export interface SnapshotTranscriptCacheOptions { activeSessionId: string; snapshotId: string; messages?: readonly AgentMessage[]; cacheRoot: string; targetChunkBytes?: number; memoryCacheBytes?: number;}
export type SnapshotTranscriptChunkSource = (Iterable<Buffer> | AsyncIterable<Buffer>) & { markFailed?(error: Error): void; dispose?(): void;};
export function createSnapshotTranscriptChunks(options: { activeSessionId: string; snapshotId: string; messages: readonly AgentMessage[]; targetChunkBytes?: number; signal?: AbortSignal;}): Iterable<Buffer> { const messages = [...options.messages]; const targetChunkBytes = options.targetChunkBytes ?? SNAPSHOT_TARGET_CHUNK_BYTES; return { *[Symbol.iterator](): Iterator<Buffer> { options.signal?.throwIfAborted(); let serializedMessages: string[] = []; let serializedBytes = 0; let index = 0; const flush = (): Buffer | undefined => { if (serializedMessages.length === 0) { return undefined; } const prefix = `{"type":"session_snapshot_chunk","activeSessionId":${JSON.stringify(options.activeSessionId)},` + `"snapshotId":${JSON.stringify(options.snapshotId)},"index":${index},"messages":[`; const line = Buffer.from(`${prefix}${serializedMessages.join(",")}]}\n`); serializedMessages = []; serializedBytes = 0; index++; return line; };
for (const message of messages) { options.signal?.throwIfAborted(); const serialized = JSON.stringify(message); const bytes = Buffer.byteLength(serialized) + (serializedMessages.length > 0 ? 1 : 0); if (serializedMessages.length > 0 && serializedBytes + bytes > targetChunkBytes) { const chunk = flush(); if (chunk) yield chunk; } serializedMessages.push(serialized); serializedBytes += bytes; } options.signal?.throwIfAborted(); const chunk = flush(); if (chunk) yield chunk; }, };}
export class SnapshotTranscriptCache { private readonly chunks: SnapshotTranscriptChunk[] = []; private cacheDirectory?: string; private totalBytes = 0; private completed = false; private readers = 0; private disposeRequested = false; private disposed = false; private failure?: Error; private readonly chunkWaiters = new Map< number, Array<{ resolve: (buffer: Buffer | undefined) => void; reject: (error: Error) => void }> >(); readonly targetChunkBytes: number; readonly snapshotId: string; readonly activeSessionId: string;
constructor(private readonly options: SnapshotTranscriptCacheOptions) { this.targetChunkBytes = options.targetChunkBytes ?? SNAPSHOT_TARGET_CHUNK_BYTES; this.snapshotId = options.snapshotId; this.activeSessionId = options.activeSessionId; if (options.messages) { for (const chunk of createSnapshotTranscriptChunks({ ...options, messages: options.messages })) { this.storeChunk(chunk); } this.completed = true; } }
get chunkCount(): number { return this.chunks.length; }
get complete(): boolean { return this.completed && !this.failure && !this.disposed; }
get bytes(): number { return this.totalBytes; }
get fileBacked(): boolean { return this.cacheDirectory !== undefined; }
readChunk(index: number): Buffer { const chunk = this.chunks[index]; if (!chunk) { throw new Error(`Unknown snapshot transcript chunk: ${index}`); } if (chunk.buffer) { return chunk.buffer; } if (!chunk.path) { throw new Error(`Snapshot transcript chunk ${index} has no backing storage`); } return readFileSync(chunk.path); }
*[Symbol.iterator](): Iterator<Buffer> { for (let index = 0; index < this.chunkCount; index++) { yield this.readChunk(index); } }
appendEncodedChunk(buffer: Buffer): void { if (this.completed || this.failure || this.disposed) { throw new Error(`Snapshot transcript ${this.snapshotId} is not writable`); } this.storeChunk(buffer); }
markComplete(): void { if (this.completed) { return; } if (this.failure || this.disposed) { throw new Error(`Snapshot transcript ${this.snapshotId} cannot be completed`); } this.completed = true; for (const [index, waiters] of this.chunkWaiters) { if (index < this.chunks.length) { continue; } for (const waiter of waiters) { waiter.resolve(undefined); } this.chunkWaiters.delete(index); } }
markFailed(error: Error): void { if (this.failure) { return; } this.failure = error; for (const waiters of this.chunkWaiters.values()) { for (const waiter of waiters) { waiter.reject(error); } } this.chunkWaiters.clear(); }
waitForChunk(index: number, signal?: AbortSignal): Promise<Buffer | undefined> { if (signal?.aborted) return Promise.reject(signal.reason); if (this.failure) { return Promise.reject(this.failure); } if (index < this.chunks.length) { return Promise.resolve(this.readChunk(index)); } if (this.completed) { return Promise.resolve(undefined); } return new Promise((resolve, reject) => { const waiters = this.chunkWaiters.get(index) ?? []; const onAbort = () => { waiters.splice(waiters.indexOf(waiter), 1); if (waiters.length === 0) this.chunkWaiters.delete(index); reject(signal?.reason); }; const waiter = { resolve: (buffer: Buffer | undefined) => { signal?.removeEventListener("abort", onAbort); resolve(buffer); }, reject: (error: Error) => { signal?.removeEventListener("abort", onAbort); reject(error); }, }; waiters.push(waiter); this.chunkWaiters.set(index, waiters); signal?.addEventListener("abort", onAbort, { once: true }); }); }
retain(): () => void { if (this.disposed) { throw new Error(`Snapshot transcript ${this.snapshotId} was disposed`); } this.readers++; let released = false; return () => { if (released) return; released = true; this.readers--; if (this.readers === 0 && this.disposeRequested) { this.disposeNow(); } }; }
dispose(): void { if (this.disposed || this.disposeRequested) return; this.disposeRequested = true; if (this.readers > 0) return; this.disposeNow(); }
private disposeNow(): void { if (this.disposed) return; this.disposed = true; this.markFailed(new Error(`Snapshot transcript ${this.snapshotId} was disposed`)); if (this.cacheDirectory) { rmSync(this.cacheDirectory, { recursive: true, force: true }); this.cacheDirectory = undefined; } this.chunks.length = 0; }
private storeChunk(buffer: Buffer): void { this.totalBytes += buffer.length; const memoryLimit = this.options.memoryCacheBytes ?? SNAPSHOT_MEMORY_CACHE_BYTES; if (!this.cacheDirectory && this.totalBytes > memoryLimit) { this.cacheDirectory = join(this.options.cacheRoot, this.options.snapshotId.replaceAll(/[^a-zA-Z0-9_-]/g, "_")); mkdirSync(this.cacheDirectory, { recursive: true, mode: 0o700 }); for (let index = 0; index < this.chunks.length; index++) { const existing = this.chunks[index]!; if (!existing.buffer) { continue; } const path = join(this.cacheDirectory, `${index}.jsonl`); writeFileSync(path, existing.buffer, { mode: 0o600 }); this.chunks[index] = { path }; } }
if (this.cacheDirectory) { const path = join(this.cacheDirectory, `${this.chunks.length}.jsonl`); writeFileSync(path, buffer, { mode: 0o600 }); this.chunks.push({ path }); } else { this.chunks.push({ buffer }); } const index = this.chunks.length - 1; const waiters = this.chunkWaiters.get(index); if (waiters) { const stored = this.readChunk(index); for (const waiter of waiters) { waiter.resolve(stored); } this.chunkWaiters.delete(index); } }}