Something went wrong. Try again.
snapcompact for prime agent
Something went wrong. Try again.
prime-snapcompact index.ts
44 kB · 1181 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182import { compact as nativeCompact, convertToLlm, getLatestCompactionEntry, type CompactionEntry, type CompactionResult, type ExtensionAPI, type ExtensionContext, type SessionBeforeCompactEvent, type SessionEntry,} from "@earendil-works/pi-coding-agent";import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core";import type { ImageContent, TextContent } from "@earendil-works/pi-ai";import { createHash, randomUUID } from "node:crypto";import { type Dirent, existsSync, mkdirSync, readFileSync, readdirSync, renameSync, rmSync, statSync, writeFileSync,} from "node:fs";import { homedir } from "node:os";import { basename, dirname, join } from "node:path";import { spawn } from "node:child_process";import { fileURLToPath } from "node:url";
const SCHEMA = 1;const PACKAGE_VERSION = "18.0.8";const MAX_FRAME_DATA_BYTES = 3_000_000;const MAX_ARCHIVE_FRAMES = 80;const MAX_PNG_DIMENSION = 4096;const MAX_ARTIFACT_BYTES = 8 * 1024 * 1024;const MAX_WORKER_INPUT_BYTES = 32 * 1024 * 1024;const MAX_WORKER_OUTPUT_BYTES = 8 * 1024 * 1024;const MAX_WORKER_ERROR_BYTES = 64 * 1024;const WORKER_TIMEOUT_MS = 120_000;const CONTEXT_SAFETY_TOKENS = 4096;const UNBOUNDED_TEXT_TOKENS = Number.MAX_SAFE_INTEGER;const MIGRATION_SAFETY_TOKENS = 8192;const NEWLINE_GLYPH = "█";const DIM_MARKERS = /[\u000e\u000f]/g;const WORKER_PATH = fileURLToPath(new URL("./worker.ts", import.meta.url));
interface SnapcompactArchive { frames: Array<{ data: string; mimeType: string; chars?: number; detail?: "low" | "high" | "auto" | "original"; }>; totalChars: number; truncatedChars: number; dataBudgetOmittedFrames?: number; text?: string; textHead?: string; textTail?: string;}
interface SnapcompactArtifact { schema: number; packageVersion: string; preserveData: { snapcompact: SnapcompactArchive; [key: string]: unknown; };}
interface ArtifactRef { file: string; bytes: number; sha256: string;}
interface SnapcompactState { schema: number; packageVersion: string; producer: { provider: string; api: string; id: string; }; frameBudget: number; frameTokenEstimate: number; archiveTokenEstimate: number; shortSummary?: string; artifact: ArtifactRef;}
interface TrackedFileDetails { readFiles: string[]; modifiedFiles: string[];}
interface SnapcompactDetails extends TrackedFileDetails { snapcompact: SnapcompactState;}
interface SnapcompactTextDetails extends TrackedFileDetails { snapcompactText: { schema: number; };}
interface WorkerResult { packageVersion: string; frameBudget: number; frameTokenEstimate: number; archiveTokenEstimate: number; summary: string; shortSummary?: string; firstKeptEntryId: string; tokensBefore: number; readFiles: string[]; modifiedFiles: string[]; preserveData: SnapcompactArtifact["preserveData"];}
export interface SnapcompactDependencies { runWorker?: typeof runWorker; nativeCompact?: typeof nativeCompact;}
const artifactCache = new Map<string, SnapcompactArtifact>();
function bunExecutable(): string { const override = process.env.PRIME_SNAPCOMPACT_BUN; if (override) return override; const userBun = join(homedir(), ".bun", "bin", "bun"); return existsSync(userBun) ? userBun : "bun";}
function abortError(): Error { const error = new Error("SnapCompact cancelled"); error.name = "AbortError"; return error;}
export function runWorker(payload: unknown, signal: AbortSignal): Promise<WorkerResult> { return new Promise((resolve, reject) => { if (signal.aborted) { reject(abortError()); return; } if (!existsSync(WORKER_PATH)) { reject(new Error( `SnapCompact worker source is missing at ${WORKER_PATH}; reinstall or reload the extension from a checkout that still exists`, )); return; } let input: string; try { input = JSON.stringify(payload); } catch (error) { reject(error instanceof Error ? error : new Error(String(error))); return; } if (Buffer.byteLength(input) > MAX_WORKER_INPUT_BYTES) { reject(new Error(`SnapCompact worker input exceeded ${MAX_WORKER_INPUT_BYTES} bytes`)); return; } const env: NodeJS.ProcessEnv = { ...process.env, FORCE_COLOR: "0" }; delete env.NO_COLOR; const child = spawn(bunExecutable(), ["run", WORKER_PATH], { cwd: fileURLToPath(new URL(".", import.meta.url)), env, stdio: ["pipe", "pipe", "pipe"], }); const stdout: Buffer[] = []; const stderr: Buffer[] = []; let stdoutBytes = 0; let stderrBytes = 0; let settled = false; let timeout: NodeJS.Timeout | undefined;
const finish = (error?: Error, result?: WorkerResult) => { if (settled) return; settled = true; if (timeout) clearTimeout(timeout); signal.removeEventListener("abort", onAbort); if (error) reject(error); else resolve(result as WorkerResult); }; const onAbort = () => { child.kill("SIGTERM"); finish(abortError()); }; signal.addEventListener("abort", onAbort, { once: true }); timeout = setTimeout(() => { child.kill("SIGTERM"); finish(new Error(`SnapCompact worker timed out after ${WORKER_TIMEOUT_MS} ms`)); }, WORKER_TIMEOUT_MS); timeout.unref();
child.on("error", (error) => finish(error)); child.stdout.on("data", (chunk: Buffer) => { stdoutBytes += chunk.length; if (stdoutBytes > MAX_WORKER_OUTPUT_BYTES) { child.kill("SIGTERM"); finish(new Error(`SnapCompact worker output exceeded ${MAX_WORKER_OUTPUT_BYTES} bytes`)); return; } stdout.push(chunk); }); child.stderr.on("data", (chunk: Buffer) => { if (stderrBytes >= MAX_WORKER_ERROR_BYTES) return; const remaining = MAX_WORKER_ERROR_BYTES - stderrBytes; stderr.push(chunk.subarray(0, remaining)); stderrBytes += Math.min(chunk.length, remaining); }); child.on("close", (code, closeSignal) => { if (settled) return; const errorText = Buffer.concat(stderr).toString("utf8").trim(); if (code !== 0) { finish(new Error( `SnapCompact worker failed (${closeSignal ? `signal ${closeSignal}` : `exit ${code}`}): ${errorText || "no error output"}`, )); return; } try { const result = JSON.parse(Buffer.concat(stdout).toString("utf8")); validateWorkerResult(result); finish(undefined, result); } catch (error) { finish(error instanceof Error ? error : new Error(String(error))); } });
child.stdin.on("error", (error) => finish(error)); child.stdin.end(input); });}
function isStringArray(value: unknown): value is string[] { return Array.isArray(value) && value.every((item) => typeof item === "string");}
function safePngData(data: string): boolean { if (data.length > MAX_FRAME_DATA_BYTES || !data.startsWith("iVBORw0KGgo")) return false; try { const header = Buffer.from(data.slice(0, 64), "base64"); if (header.length < 24 || header.subarray(0, 8).toString("hex") !== "89504e470d0a1a0a") { return false; } const width = header.readUInt32BE(16); const height = header.readUInt32BE(20); return width > 0 && height > 0 && width <= MAX_PNG_DIMENSION && height <= MAX_PNG_DIMENSION; } catch { return false; }}
function snapcompactArchive(value: unknown): value is SnapcompactArchive { if (!value || typeof value !== "object") return false; const archive = value as Partial<SnapcompactArchive>; return ( Array.isArray(archive.frames) && archive.frames.length <= MAX_ARCHIVE_FRAMES && archive.frames.every((frame) => !!frame && typeof frame.data === "string" && safePngData(frame.data) && frame.mimeType === "image/png" && (frame.detail === undefined || ["low", "high", "auto", "original"].includes(frame.detail)) ) && typeof archive.totalChars === "number" && typeof archive.truncatedChars === "number" && (archive.dataBudgetOmittedFrames === undefined || (Number.isInteger(archive.dataBudgetOmittedFrames) && archive.dataBudgetOmittedFrames >= 0)) && (archive.text === undefined || typeof archive.text === "string") && (archive.textHead === undefined || typeof archive.textHead === "string") && (archive.textTail === undefined || typeof archive.textTail === "string") );}
function validatePreserveData(value: unknown): asserts value is SnapcompactArtifact["preserveData"] { if (!value || typeof value !== "object" || !snapcompactArchive((value as Record<string, unknown>).snapcompact)) { throw new Error("SnapCompact preserve data is missing a valid archive"); } const archive = (value as SnapcompactArtifact["preserveData"]).snapcompact; const frameBytes = archive.frames.reduce((sum, frame) => sum + frame.data.length, 0); if (frameBytes > MAX_FRAME_DATA_BYTES) { throw new Error("SnapCompact archive exceeded its frame-data budget"); }}
function validateWorkerResult(value: unknown): asserts value is WorkerResult { if (!value || typeof value !== "object") throw new Error("SnapCompact worker returned a non-object"); const result = value as Partial<WorkerResult>; if (result.packageVersion !== PACKAGE_VERSION) throw new Error("SnapCompact worker version mismatch"); if (typeof result.summary !== "string" || typeof result.firstKeptEntryId !== "string") { throw new Error("SnapCompact worker result is missing summary metadata"); } if ( typeof result.tokensBefore !== "number" || typeof result.frameBudget !== "number" || typeof result.frameTokenEstimate !== "number" || typeof result.archiveTokenEstimate !== "number" ) { throw new Error("SnapCompact worker result is missing numeric metadata"); } if (!isStringArray(result.readFiles) || !isStringArray(result.modifiedFiles)) { throw new Error("SnapCompact worker result has invalid file metadata"); } validatePreserveData(result.preserveData); if (result.preserveData.snapcompact.frames.length > result.frameBudget) { throw new Error("SnapCompact worker exceeded its frame budget"); }}
function validArtifactRef(value: unknown): value is ArtifactRef { if (!value || typeof value !== "object") return false; const ref = value as Partial<ArtifactRef>; return ( typeof ref.file === "string" && basename(ref.file) === ref.file && /^[0-9a-f-]+\.json$/i.test(ref.file) && typeof ref.bytes === "number" && ref.bytes > 0 && ref.bytes <= MAX_ARTIFACT_BYTES && typeof ref.sha256 === "string" && /^[0-9a-f]{64}$/i.test(ref.sha256) );}
function detailsFromEntry(entry: CompactionEntry | null): SnapcompactDetails | undefined { if (!entry?.details || typeof entry.details !== "object") return undefined; const details = entry.details as Partial<SnapcompactDetails>; const state = details.snapcompact; if ( !state || state.schema !== SCHEMA || state.packageVersion !== PACKAGE_VERSION || !state.producer || typeof state.producer.provider !== "string" || typeof state.producer.api !== "string" || typeof state.producer.id !== "string" || typeof state.frameBudget !== "number" || typeof state.frameTokenEstimate !== "number" || typeof state.archiveTokenEstimate !== "number" || !validArtifactRef(state.artifact) || !isStringArray(details.readFiles) || !isStringArray(details.modifiedFiles) ) { return undefined; } return details as SnapcompactDetails;}
function trackedDetailsFromEntry(entry: CompactionEntry | null): TrackedFileDetails | undefined { const snapcompact = detailsFromEntry(entry); if (snapcompact) return snapcompact; if (!entry?.details || typeof entry.details !== "object") return undefined; const details = entry.details as Partial<SnapcompactTextDetails>; if ( details.snapcompactText?.schema !== SCHEMA || !isStringArray(details.readFiles) || !isStringArray(details.modifiedFiles) ) { return undefined; } return details as SnapcompactTextDetails;}
function previousDetails(entries: SessionEntry[]): SnapcompactDetails | undefined { return detailsFromEntry(getLatestCompactionEntry(entries));}
function artifactDirectoryForFile(sessionFile: string, sessionId?: string): string { const id = sessionId ?? basename(sessionFile).replace(/\.jsonl$/, ""); if (!id || basename(id) !== id) throw new Error("invalid Prime session id for SnapCompact artifact"); return join(dirname(dirname(sessionFile)), "session-artifacts", id, "snapcompact");}
function artifactDirectory(ctx: ExtensionContext): string | undefined { const sessionFile = ctx.sessionManager.getSessionFile(); if (!sessionFile) return undefined; return artifactDirectoryForFile(sessionFile, ctx.sessionManager.getSessionId());}
function sha256(data: string | Buffer): string { return createHash("sha256").update(data).digest("hex");}
function cacheArtifact(key: string, artifact: SnapcompactArtifact) { artifactCache.delete(key); artifactCache.set(key, artifact); while (artifactCache.size > 4) { const oldest = artifactCache.keys().next().value; if (oldest === undefined) break; artifactCache.delete(oldest); }}
function atomicWriteArtifact(directory: string, file: string, data: string | Buffer) { mkdirSync(directory, { recursive: true, mode: 0o700 }); const target = join(directory, file); const temporary = `${target}.tmp-${process.pid}-${randomUUID()}`; try { writeFileSync(temporary, data, { mode: 0o600, flag: "wx" }); renameSync(temporary, target); } catch (error) { rmSync(temporary, { force: true }); throw error; }}
function isMissingPathError(error: unknown): boolean { return typeof error === "object" && error !== null && "code" in error && error.code === "ENOENT";}
function writeArtifact( ctx: ExtensionContext, preserveData: SnapcompactArtifact["preserveData"],): ArtifactRef { validatePreserveData(preserveData); const directory = artifactDirectory(ctx); if (!directory) throw new Error("SnapCompact requires a persisted Prime session"); const artifact: SnapcompactArtifact = { schema: SCHEMA, packageVersion: PACKAGE_VERSION, preserveData, }; const data = JSON.stringify(artifact); const bytes = Buffer.byteLength(data); if (bytes > MAX_ARTIFACT_BYTES) { throw new Error(`SnapCompact artifact exceeded ${MAX_ARTIFACT_BYTES} bytes`); } const file = `${randomUUID()}.json`; atomicWriteArtifact(directory, file, data); const ref = { file, bytes, sha256: sha256(data) }; cacheArtifact(ref.sha256, artifact); return ref;}
function readArtifactAt( directory: string, state: SnapcompactState, useCache: boolean,): { artifact: SnapcompactArtifact; data: Buffer } { const cached = useCache ? artifactCache.get(state.artifact.sha256) : undefined; if (cached) return { artifact: cached, data: Buffer.alloc(0) }; const path = join(directory, state.artifact.file); const stat = statSync(path); if (!stat.isFile() || stat.size !== state.artifact.bytes || stat.size > MAX_ARTIFACT_BYTES) { throw new Error("SnapCompact artifact size or type mismatch"); } const data = readFileSync(path); if (sha256(data) !== state.artifact.sha256) { throw new Error("SnapCompact artifact integrity check failed"); } const artifact = JSON.parse(data.toString("utf8")) as SnapcompactArtifact; if (artifact.schema !== SCHEMA || artifact.packageVersion !== PACKAGE_VERSION) { throw new Error("SnapCompact artifact version mismatch"); } validatePreserveData(artifact.preserveData); cacheArtifact(state.artifact.sha256, artifact); return { artifact, data };}
function loadArtifact(ctx: ExtensionContext, state: SnapcompactState): SnapcompactArtifact { const cached = artifactCache.get(state.artifact.sha256); if (cached) return cached; const directory = artifactDirectory(ctx); if (!directory) throw new Error("SnapCompact artifact has no persisted session directory"); return readArtifactAt(directory, state, false).artifact;}
function copyForkArtifacts(previousSessionFile: string, ctx: ExtensionContext) { const targetDirectory = artifactDirectory(ctx); if (!targetDirectory) throw new Error("forked SnapCompact archive has no target session directory"); const sourceDirectory = artifactDirectoryForFile(previousSessionFile); for (const entry of ctx.sessionManager.getEntries()) { if (entry.type !== "compaction") continue; const details = detailsFromEntry(entry); if (!details) continue; const state = details.snapcompact; let source: ReturnType<typeof readArtifactAt>; try { source = readArtifactAt(sourceDirectory, state, false); } catch (error) { if (isMissingPathError(error)) continue; throw error; } const target = join(targetDirectory, state.artifact.file); if (existsSync(target)) { try { readArtifactAt(targetDirectory, state, false); continue; } catch (error) { if (!isMissingPathError(error)) throw error; } } atomicWriteArtifact(targetDirectory, state.artifact.file, source.data); }}
function referencedArtifactFiles(ctx: ExtensionContext): Set<string> { const referenced = new Set<string>(); for (const entry of ctx.sessionManager.getEntries()) { if (entry.type !== "compaction" || !entry.details || typeof entry.details !== "object") continue; const state = (entry.details as { snapcompact?: { artifact?: unknown } }).snapcompact; if (state && validArtifactRef(state.artifact)) referenced.add(state.artifact.file); } return referenced;}
function gcArtifacts(ctx: ExtensionContext): { files: number; bytes: number } { const directory = artifactDirectory(ctx); if (!directory || !existsSync(directory)) return { files: 0, bytes: 0 }; const referenced = referencedArtifactFiles(ctx); let entries: Dirent[]; try { entries = readdirSync(directory, { withFileTypes: true }); } catch (error) { if (isMissingPathError(error)) return { files: 0, bytes: 0 }; throw error; } let files = 0; let bytes = 0; for (const entry of entries) { if (!entry.isFile()) continue; const isArtifact = /^[0-9a-f-]+\.json$/i.test(entry.name); const isTemporary = entry.name.includes(".tmp-"); if ((!isArtifact && !isTemporary) || referenced.has(entry.name)) continue; const path = join(directory, entry.name); let size: number; try { size = statSync(path).size; } catch (error) { if (isMissingPathError(error)) continue; throw error; } rmSync(path, { force: true }); bytes += size; files++; } return { files, bytes };}
function plainText(text: string): string { return text.replace(DIM_MARKERS, "").replaceAll(NEWLINE_GLYPH, "\n");}
function archivePlainText(archive: SnapcompactArchive): string { if (archive.text) return plainText(archive.text); return [archive.textHead, archive.textTail] .filter((part): part is string => typeof part === "string" && part.length > 0) .map(plainText) .join("\n\n-------------- middle history unavailable\n\n");}
function migrationSourceText(archive: SnapcompactArchive): string { if (typeof archive.text !== "string" || archive.text.length === 0) { throw new Error("SnapCompact archive has no complete recoverable source text"); } return plainText(archive.text);}
function conservativeTextTokens(text: string): number { let nonAscii = 0; for (let index = 0; index < text.length; index++) { if (text.charCodeAt(index) > 0x7f) nonAscii++; } return Math.ceil((text.length - nonAscii) / 3 + nonAscii * 2);}
function boundedText(text: string, maxChars: number, label: string): string { if (text.length <= maxChars) return text; const headChars = Math.floor(maxChars * 0.4); const tailChars = maxChars - headChars; const omitted = text.length - maxChars; return ( text.slice(0, headChars) + `
-------------- ${omitted.toLocaleString()} middle ${label} characters omitted for this request; the persisted archive remains intact
` + text.slice(text.length - tailChars) );}
function boundedTextForTokens(text: string, maxTokens: number, label: string): string { if (maxTokens <= 0) return ""; if (conservativeTextTokens(text) <= maxTokens) return text; let low = 0; let high = text.length; let best = ""; while (low <= high) { const chars = Math.floor((low + high) / 2); const candidate = boundedText(text, chars, label); if (conservativeTextTokens(candidate) <= maxTokens) { best = candidate; low = chars + 1; } else { high = chars - 1; } } return best;}
function omissionNotice(frames: number): string { return `-------------- ${frames.toLocaleString()} older SnapCompact image frame${frames === 1 ? "" : "s"} omitted by the current request budget`;}
function archiveBlocks( archive: SnapcompactArchive, maxImages: number, maxTextTokens: number,): Array<TextContent | ImageContent> { const selected: SnapcompactArchive["frames"] = []; let bytes = 0; for (let index = archive.frames.length - 1; index >= 0 && selected.length < maxImages; index--) { const frame = archive.frames[index]; if (bytes + frame.data.length > MAX_FRAME_DATA_BYTES) break; selected.push(frame); bytes += frame.data.length; } selected.reverse(); const omitted = archive.frames.length - selected.length + (archive.dataBudgetOmittedFrames ?? 0); const rawHead = plainText(archive.textHead ?? ""); const rawTail = plainText(archive.textTail ?? ""); const totalEdgeTokens = conservativeTextTokens(rawHead) + conservativeTextTokens(rawTail); const headBudget = totalEdgeTokens <= maxTextTokens ? conservativeTextTokens(rawHead) : Math.min(conservativeTextTokens(rawHead), Math.floor(maxTextTokens * 0.4)); const tailBudget = totalEdgeTokens <= maxTextTokens ? conservativeTextTokens(rawTail) : Math.min(conservativeTextTokens(rawTail), Math.max(0, maxTextTokens - headBudget)); const head = boundedTextForTokens(rawHead, headBudget, "oldest edge"); const tail = boundedTextForTokens(rawTail, tailBudget, "newest edge"); const blocks: Array<TextContent | ImageContent> = []; const hasImages = selected.length > 0; if (head) { const suffix = hasImages ? "\n-------------- imaged middle below\n" : omitted > 0 ? `${omissionNotice(omitted)}` : ""; blocks.push({ type: "text", text: head + suffix }); } if (hasImages && omitted > 0) blocks.push({ type: "text", text: omissionNotice(omitted) }); blocks.push(...selected.map((frame) => ({ type: "image" as const, data: frame.data, mimeType: frame.mimeType, ...(frame.detail ? { detail: frame.detail } : {}), }))); if (tail) { const prefix = hasImages ? "-------------- imaged middle above\n" : omitted > 0 || archive.truncatedChars > 0 ? "\n-------------- middle history omitted above\n" : ""; const text = prefix + tail; const last = blocks.at(-1); if (last?.type === "text") last.text += text; else blocks.push({ type: "text", text }); } if (blocks.length === 0 && archive.text) { const text = boundedTextForTokens(archivePlainText(archive), maxTextTokens, "archive"); if (text) blocks.push({ type: "text", text }); } return blocks;}
function mergeFileOps( event: SessionBeforeCompactEvent, previous: TrackedFileDetails | undefined,) { const read = new Set(event.preparation.fileOps.read); const written = new Set(event.preparation.fileOps.written); const edited = new Set(event.preparation.fileOps.edited); for (const file of previous?.readFiles ?? []) read.add(file); for (const file of previous?.modifiedFiles ?? []) edited.add(file); return { read, written, edited };}
function maxArchiveTokens(event: SessionBeforeCompactEvent, ctx: ExtensionContext): number { const model = ctx.model; if (!model) return 0; const systemTokens = conservativeTextTokens(ctx.getSystemPrompt()); return Math.max( 0, model.contextWindow - systemTokens - event.preparation.settings.keepRecentTokens - event.preparation.settings.reserveTokens - CONTEXT_SAFETY_TOKENS, );}
function workerPayload( event: SessionBeforeCompactEvent, ctx: ExtensionContext, previous: TrackedFileDetails | undefined, previousArtifact: SnapcompactArtifact | undefined,) { const model = ctx.model; if (!model) throw new Error("SnapCompact has no active model"); const fileOps = mergeFileOps(event, previous); return { preparation: { firstKeptEntryId: event.preparation.firstKeptEntryId, messagesToSummarize: convertToLlm(event.preparation.messagesToSummarize), turnPrefixMessages: convertToLlm(event.preparation.turnPrefixMessages), tokensBefore: event.preparation.tokensBefore, previousSummary: event.preparation.previousSummary, previousPreserveData: previousArtifact?.preserveData, fileOps: { read: [...fileOps.read], written: [...fileOps.written], edited: [...fileOps.edited], }, }, model: { api: model.api, id: model.id, provider: model.provider, }, includeThinking: !(model.api === "anthropic-messages" || /claude/i.test(model.id)), maxFrameDataBytes: MAX_FRAME_DATA_BYTES, maxArchiveTokens: maxArchiveTokens(event, ctx), };}
function markTrackedTextResult(result: CompactionResult): CompactionResult<SnapcompactTextDetails> { const details = result.details as Partial<TrackedFileDetails> | undefined; return { ...result, details: { readFiles: isStringArray(details?.readFiles) ? details.readFiles : [], modifiedFiles: isStringArray(details?.modifiedFiles) ? details.modifiedFiles : [], snapcompactText: { schema: SCHEMA }, }, };}
async function nativeCompactTracked( event: SessionBeforeCompactEvent, ctx: ExtensionContext, previous: TrackedFileDetails, compactImpl: typeof nativeCompact, thinkingLevel: ThinkingLevel, messagesToSummarize = event.preparation.messagesToSummarize,) { const model = ctx.model; if (!model) throw new Error("text compaction has no active model"); const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model); if (!auth.ok) throw new Error(auth.error); const result = await compactImpl( { ...event.preparation, messagesToSummarize, fileOps: mergeFileOps(event, previous), }, model, auth.apiKey ?? "", auth.headers, event.customInstructions, event.signal, thinkingLevel, ); return markTrackedTextResult(result);}
async function textCompactWithArchive( event: SessionBeforeCompactEvent, ctx: ExtensionContext, previous: SnapcompactDetails, artifact: SnapcompactArtifact, compactImpl: typeof nativeCompact, thinkingLevel: ThinkingLevel,) { const model = ctx.model; if (!model) throw new Error("text compaction has no active model"); const source = migrationSourceText(artifact.preserveData.snapcompact); const prefix = "The following is the source transcript preserved by an earlier SnapCompact archive. " + "Treat it as prior conversation, not as a new user request.\n\n"; const currentPayload = JSON.stringify(convertToLlm([ ...event.preparation.messagesToSummarize, ...event.preparation.turnPrefixMessages, ])); const currentTokens = conservativeTextTokens(currentPayload); const metadataTokens = conservativeTextTokens( (event.preparation.previousSummary ?? "") + (event.customInstructions ?? ""), ); const inputBudget = Math.max( 0, model.contextWindow - event.preparation.settings.reserveTokens - MIGRATION_SAFETY_TOKENS, ); const prefixTokens = conservativeTextTokens(prefix); const archiveTokens = conservativeTextTokens(source); // One request carries as much archive text as this request's own budget leaves // room for. The bound is that budget, never a fixed ceiling: a constant cap is // what dropped whole archives on models that could have carried them, and the // previous fix removed it. Nothing is staged either, because splitting one // summary across calls changes what the summarizer sees. const availableArchiveTokens = Math.max(0, inputBudget - currentTokens - metadataTokens - prefixTokens); const carriedTokens = Math.min(archiveTokens, availableArchiveTokens); const droppedTokens = archiveTokens - carriedTokens; notify( ctx, `SnapCompact is migrating the ${formatModel(previous.snapcompact.producer)} image archive to text ` + `compaction: about ${(prefixTokens + carriedTokens + currentTokens + metadataTokens).toLocaleString()} ` + `input tokens (archive ${archiveTokens.toLocaleString()}, current messages ${currentTokens.toLocaleString()}) ` + `against a ${inputBudget.toLocaleString()}-token budget; carrying ${carriedTokens.toLocaleString()} of the ` + `archive's ${archiveTokens.toLocaleString()} text tokens` + (droppedTokens > 0 ? ` and dropping ${droppedTokens.toLocaleString()} of them from this request` : ""), );
const messagesToSummarize: AgentMessage[] = []; if (carriedTokens > 0) { const archiveSource = archiveTokens <= availableArchiveTokens ? source : boundedTextForTokens(source, availableArchiveTokens, "archive"); if (archiveSource.length > 0) { messagesToSummarize.push({ role: "user", content: [{ type: "text", text: prefix + archiveSource }], timestamp: Date.now(), }); } } const carriedArchive = messagesToSummarize.length > 0; if (!carriedArchive) { notify( ctx, `SnapCompact cannot fit the archive's ${archiveTokens.toLocaleString()} text tokens next to ` + `${(currentTokens + metadataTokens).toLocaleString()} tokens of current messages in a ` + `${inputBudget.toLocaleString()}-token budget, so this compaction runs without the archive text; ` + `the archive file is left untouched`, "warning", ); } messagesToSummarize.push(...event.preparation.messagesToSummarize);
try { return await nativeCompactTracked( event, ctx, previous, compactImpl, thinkingLevel, messagesToSummarize, ); } catch (error) { if (event.signal.aborted || !carriedArchive) throw error; // The migration call is the only one that can fail on its own request shape. // Compaction still has to happen, so retry the same summarization without the // archive text instead of cancelling and leaving the session uncompactable. notify( ctx, `SnapCompact text migration failed: ${error instanceof Error ? error.message : String(error)}; ` + `retrying this compaction without the archive text`, "warning", ); console.warn("[snapcompact] text migration retry without archive", error); return nativeCompactTracked(event, ctx, previous, compactImpl, thinkingLevel); }}
function formatModel(model: { provider: string; id: string }): string { return `${model.provider}/${model.id}`;}
function notify(ctx: ExtensionContext, message: string, level: "info" | "warning" | "error" = "info") { if (!ctx.hasUI) return; try { ctx.ui.notify(message, level); } catch (error) { console.warn("[snapcompact] notification failed", error); }}
export function createSnapcompactExtension(dependencies: SnapcompactDependencies = {}) { const worker = dependencies.runWorker ?? runWorker; const compactImpl = dependencies.nativeCompact ?? nativeCompact;
return function snapcompactExtension(pi: ExtensionAPI) { let enabled = process.env.PRIME_SNAPCOMPACT !== "0"; const notifiedDegradations = new Set<string>(); pi.registerCommand("snapcompact", { description: "Show, change, or clean automatic SnapCompact compaction", handler: async (args, ctx) => { const action = args.trim().toLowerCase(); if (action === "gc") { const removed = gcArtifacts(ctx); notify(ctx, `SnapCompact removed ${removed.files} unreferenced artifact file${removed.files === 1 ? "" : "s"} (${removed.bytes.toLocaleString()} bytes)`); return; } if (action === "on") enabled = true; else if (action === "off") enabled = false; else if (action && action !== "status") { notify(ctx, "usage: /snapcompact [status|on|off|gc]", "warning"); return; } const vision = ctx.model?.input.includes("image") === true; notify( ctx, `SnapCompact ${enabled ? "enabled" : "disabled"}; active model ${vision ? "accepts" : "does not accept"} images`, ); }, });
pi.on("session_start", (event, ctx) => { if (event.reason === "fork" && event.previousSessionFile) { try { copyForkArtifacts(event.previousSessionFile, ctx); } catch (error) { const message = error instanceof Error ? error.message : String(error); notify(ctx, `SnapCompact could not copy inherited fork artifacts: ${message}`, "error"); console.warn("[snapcompact] fork artifact copy failed", error); } } try { gcArtifacts(ctx); } catch (error) { console.warn("[snapcompact] artifact cleanup failed", error); } });
pi.on("session_compact", (_event, ctx) => { try { gcArtifacts(ctx); } catch (error) { console.warn("[snapcompact] artifact cleanup failed", error); } });
pi.on("session_before_compact", async (event, ctx) => { const latestCompaction = getLatestCompactionEntry(event.branchEntries); const previous = detailsFromEntry(latestCompaction); const tracked = trackedDetailsFromEntry(latestCompaction); let previousArtifact: SnapcompactArtifact | undefined; if (previous) { try { previousArtifact = loadArtifact(ctx, previous.snapcompact); } catch (error) { notify(ctx, `SnapCompact archive cannot be read: ${error instanceof Error ? error.message : String(error)}`, "error"); return { cancel: true }; } } const vision = ctx.model?.input.includes("image") === true; // Model identity never gates the compaction path (the oh-my-pi snapcompact // contract): the archive's frames are images of text that any vision model // reads, so a model switch only changes which producer is recorded next. const needsTextPath = !enabled || !vision || event.customInstructions !== undefined; if (needsTextPath) { if (previous && previousArtifact) { try { return { compaction: await textCompactWithArchive( event, ctx, previous, previousArtifact, compactImpl, pi.getThinkingLevel(), ), }; } catch (error) { if (event.signal.aborted) return { cancel: true }; const message = error instanceof Error ? error.message : String(error); notify(ctx, `SnapCompact archive could not transition to text compaction: ${message}`, "error"); console.warn("[snapcompact] text transition failed", error); return { cancel: true }; } } if (!tracked) return; try { return { compaction: await nativeCompactTracked(event, ctx, tracked, compactImpl, pi.getThinkingLevel()) }; } catch (error) { if (!event.signal.aborted) { notify(ctx, `tracked native compaction failed: ${error instanceof Error ? error.message : String(error)}`, "error"); console.warn("[snapcompact] tracked native compaction failed", error); } return { cancel: true }; } }
try { const result = await worker( workerPayload(event, ctx, tracked, previousArtifact), event.signal, ); const artifact = writeArtifact(ctx, result.preserveData); const details: SnapcompactDetails = { readFiles: result.readFiles, modifiedFiles: result.modifiedFiles, snapcompact: { schema: SCHEMA, packageVersion: result.packageVersion, producer: { provider: ctx.model!.provider, api: ctx.model!.api, id: ctx.model!.id, }, frameBudget: result.frameBudget, frameTokenEstimate: result.frameTokenEstimate, archiveTokenEstimate: result.archiveTokenEstimate, shortSummary: result.shortSummary, artifact, }, }; notify(ctx, result.shortSummary ?? "SnapCompact archived prior history"); return { compaction: { summary: result.summary, firstKeptEntryId: result.firstKeptEntryId, tokensBefore: result.tokensBefore, details, }, }; } catch (error) { if (event.signal.aborted) return { cancel: true }; const message = error instanceof Error ? error.message : String(error); console.warn("[snapcompact] worker failed", error); if (previous && previousArtifact) { try { const compaction = await textCompactWithArchive( event, ctx, previous, previousArtifact, compactImpl, pi.getThinkingLevel(), ); notify(ctx, `SnapCompact unavailable; preserved history with native text compaction: ${message}`, "warning"); return { compaction }; } catch (migrationError) { const fallbackMessage = migrationError instanceof Error ? migrationError.message : String(migrationError); notify( ctx, `SnapCompact visual compaction failed: ${message}. Its text fallback also failed, so the existing archive was retained: ${fallbackMessage}`, "error", ); return { cancel: true }; } } if (tracked) { try { const compaction = await nativeCompactTracked(event, ctx, tracked, compactImpl, pi.getThinkingLevel()); notify(ctx, `SnapCompact unavailable; preserved tracked history with native text compaction: ${message}`, "warning"); return { compaction }; } catch (trackingError) { notify( ctx, `SnapCompact failed and retained tracked history: ${trackingError instanceof Error ? trackingError.message : String(trackingError)}`, "error", ); return { cancel: true }; } } notify(ctx, `SnapCompact unavailable; using native text compaction: ${message}`, "warning"); return; } });
pi.on("context", (event, ctx) => { const branch = ctx.sessionManager.getBranch(); const compaction = getLatestCompactionEntry(branch); const details = detailsFromEntry(compaction); if (!details) return; const summaryIndex = event.messages.findIndex((message) => message.role === "compactionSummary"); if (summaryIndex < 0) return; const summary = event.messages[summaryIndex]; if (summary.role !== "compactionSummary") return;
let artifact: SnapcompactArtifact; try { artifact = loadArtifact(ctx, details.snapcompact); } catch (error) { const message = error instanceof Error ? error.message : String(error); notify(ctx, `SnapCompact archive unavailable: ${message}`, "error"); const messages = [...event.messages]; messages[summaryIndex] = { role: "custom", customType: "snapcompact-context-error", content: [ { type: "text", text: summary.summary }, { type: "text", text: `A prior SnapCompact archive is unavailable (${message}). Do not infer its contents; re-read workspace state before continuing.`, }, ], display: false, details: { artifact: details.snapcompact.artifact }, timestamp: summary.timestamp, }; return { messages }; }
const archive = artifact.preserveData.snapcompact; const vision = ctx.model?.input.includes("image") === true; const isImagedArchive = archive.frames.length > 0; if (isImagedArchive && !vision) { const compactionKey = `${ctx.sessionManager.getSessionId()}:${details.snapcompact.artifact.sha256}`; if (!notifiedDegradations.has(compactionKey)) { notifiedDegradations.add(compactionKey); const frameCount = archive.frames.length; const frameLabel = `${frameCount} image frame${frameCount === 1 ? "" : "s"}`; const producer = formatModel(details.snapcompact.producer); notify( ctx, `SnapCompact archive rendered as ${frameLabel} by ${producer}; this model cannot read images, so the middle of the history is a text placeholder until the next compaction migrates it`, "warning", ); } }
const content: Array<TextContent | ImageContent> = [{ type: "text", text: summary.summary }]; if (vision && ctx.model) { // The archive was already fitted to the context window at compaction // time (maxArchiveTokens plus the worker's frameBudget). Re-fitting // against the live message list would rewrite this message as the tail // grows and bust the prompt cache back to the system prefix on every // budget boundary crossing, so the block layout stays frozen for the // whole compaction epoch: frame count from the stored budget, text // edges verbatim (UNBOUNDED_TEXT_TOKENS makes boundedTextForTokens a // verified passthrough), MAX_FRAME_DATA_BYTES as the only clamp. The // reader model's identity does not enter the layout. content.push(...archiveBlocks(archive, details.snapcompact.frameBudget, UNBOUNDED_TEXT_TOKENS)); } else { // Non-vision model reading an imaged archive: a transitional state. The // next compaction takes the needsTextPath branch and migrates this // archive into Prime's own text compaction, so this block only has to // stay stable until then. Both the budget and the fits/does-not-fit // decision are functions of model identity alone - never of the growing // tail - so the block stays byte-identical for the whole epoch. A tail // that outgrows the window triggers that migrating compaction instead // of this path cannibalizing the archive turn by turn to make room. const capacity = Math.max( 0, (ctx.model?.contextWindow ?? 128_000) - Math.min(ctx.model?.maxTokens ?? 0, 32_768) - CONTEXT_SAFETY_TOKENS, ); const summaryTokens = conservativeTextTokens(summary.summary); const textBudget = Math.max(0, capacity - summaryTokens); if (textBudget > 0) { const text = boundedTextForTokens( archivePlainText(archive), textBudget, "archive", ); if (text) content.push({ type: "text", text }); } }
const messages = [...event.messages]; messages[summaryIndex] = { role: "custom", customType: "snapcompact-context", content, display: false, details: { packageVersion: details.snapcompact.packageVersion, frames: archive.frames.length, totalChars: archive.totalChars, artifact: details.snapcompact.artifact, }, timestamp: summary.timestamp, }; return { messages }; }); };}
export default createSnapcompactExtension();