Something went wrong. Try again.
snapcompact for prime agent
Something went wrong. Try again.
7.0 kB · 202 lines
JavaScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203import { existsSync, mkdirSync, mkdtempSync, rmSync, symlinkSync } from "node:fs";import { spawnSync } from "node:child_process";import { join } from "node:path";import { tmpdir } from "node:os";import { fileURLToPath, pathToFileURL } from "node:url";import { resolvePrimeRoot } from "./prime-root.mjs";
const primeRoot = resolvePrimeRoot();const prime = await import(pathToFileURL(join(primeRoot, "dist/index.js")));const ai = await import(pathToFileURL(join( primeRoot, "node_modules/@earendil-works/pi-ai/dist/index.js",)));const root = mkdtempSync(join(tmpdir(), "prime-snapcompact-auto-"));const sessionDir = join(root, "sessions");const agentDir = join(root, "agent");mkdirSync(join(agentDir, "extensions"), { recursive: true });symlinkSync( fileURLToPath(new URL("..", import.meta.url)), join(agentDir, "extensions", "snapcompact"), "dir",);const registration = ai.registerFauxProvider({ models: [{ id: "snapcompact-vision", name: "SnapCompact vision test", input: ["text", "image"], reasoning: false, contextWindow: 100_000, maxTokens: 8192, }], tokensPerSecond: 100_000,});const model = registration.getModel();const seenContexts = [];const response = (text) => (context) => { seenContexts.push(context); return ai.fauxAssistantMessage(text);};registration.setResponses([ response("first response"), response("second response"),]);const auth = prime.AuthStorage.inMemory();auth.setRuntimeApiKey(model.provider, "snapcompact-test");const modelRegistry = prime.ModelRegistry.inMemory(auth);const cwd = process.cwd();const sessionManager = prime.SessionManager.create(cwd, sessionDir);const usage = { input: 94_000, output: 1000, cacheRead: 0, cacheWrite: 0, totalTokens: 95_000, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },};for (let turn = 0; turn < 8; turn++) { sessionManager.appendMessage({ role: "user", content: `analyze repository turn ${turn}`, timestamp: Date.now(), }); sessionManager.appendMessage({ role: "assistant", content: [{ type: "text", text: (`turn ${turn} repository analysis src/file.ts [] {}\n`).repeat(1800), }], api: model.api, provider: model.provider, model: model.id, usage, stopReason: "stop", timestamp: Date.now(), });}sessionManager.flushNow();
const sessionOptions = (manager) => ({ cwd, agentDir, sessionManager: manager, model, modelRegistry, noTools: "all", includeGoals: false, includeCompactSkill: false, prewarmIpythonKernel: false, telemetryDisabled: true,});const imageCount = (context) => context.messages.reduce( (count, message) => count + ( Array.isArray(message.content) ? message.content.filter((block) => block.type === "image").length : 0 ), 0,);
let { session, extensionsResult } = await prime.createAgentSession(sessionOptions(sessionManager));await session.bindExtensions({});if (extensionsResult.errors.length) throw new Error(JSON.stringify(extensionsResult.errors));try { await session.prompt("continue after automatic compaction"); const entry = prime.getLatestCompactionEntry(sessionManager.getBranch()); if (!entry?.fromHook || !entry.details?.snapcompact?.artifact) { throw new Error("automatic compaction did not persist SnapCompact details"); } const firstImages = imageCount(seenContexts[0]); if (firstImages < 1 || firstImages > entry.details.snapcompact.frameBudget) { throw new Error("automatic post-compaction request did not receive bounded archive images"); } const detailsBytes = Buffer.byteLength(JSON.stringify(entry.details)); if (detailsBytes >= 1024 * 1024 || JSON.stringify(entry.details).includes("iVBORw0KGgo")) { throw new Error("automatic compaction inlined its visual archive"); } const artifactPath = join( root, "session-artifacts", session.sessionId, "snapcompact", entry.details.snapcompact.artifact.file, ); if (!existsSync(artifactPath)) throw new Error("automatic compaction artifact is missing");
await session.prompt("one more request"); const secondImages = imageCount(seenContexts[1]); if (secondImages < 1) throw new Error("repeated request lost the active visual archive");
const sessionFile = session.sessionFile; await session.disposeAsync(); registration.appendResponses([response("resumed response")]); const resumedManager = prime.SessionManager.open(sessionFile, sessionDir, cwd); ({ session, extensionsResult } = await prime.createAgentSession(sessionOptions(resumedManager))); await session.bindExtensions({}); if (extensionsResult.errors.length) throw new Error(JSON.stringify(extensionsResult.errors)); await session.prompt("continue the resumed archive"); const resumedImages = imageCount(seenContexts[2]); if (resumedImages < 1) throw new Error("resumed session did not rehydrate archive images");
const forkSourceFile = session.sessionFile; const forkFile = resumedManager.createBranchedSession(resumedManager.getLeafId()); if (!forkFile) throw new Error("Prime did not create the fork session file"); await session.disposeAsync(); const forkManager = prime.SessionManager.open(forkFile, sessionDir, cwd); ({ session, extensionsResult } = await prime.createAgentSession({ ...sessionOptions(forkManager), sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile: forkSourceFile, }, })); await session.bindExtensions({}); if (extensionsResult.errors.length) throw new Error(JSON.stringify(extensionsResult.errors)); const forkEntry = prime.getLatestCompactionEntry(forkManager.getBranch()); const forkArtifactPath = join( root, "session-artifacts", session.sessionId, "snapcompact", forkEntry.details.snapcompact.artifact.file, ); if (!existsSync(forkArtifactPath)) { throw new Error(`fork lifecycle did not copy the inherited artifact: ${forkArtifactPath}; source=${forkSourceFile}; fork=${forkFile}; id=${session.sessionId}`); } await session.disposeAsync();
const child = spawnSync( process.execPath, [ fileURLToPath(new URL("./fork-resume-child.mjs", import.meta.url)), forkFile, sessionDir, cwd, agentDir, ], { cwd, encoding: "utf8", env: { ...process.env, FORCE_COLOR: "0", NO_COLOR: undefined }, }, ); if (child.status !== 0) { throw new Error(`fork restart child failed: ${child.stderr || child.stdout}`); } const forkRestart = JSON.parse(child.stdout.trim());
console.log(JSON.stringify({ automaticCompaction: true, detailsBytes, artifactBytes: entry.details.snapcompact.artifact.bytes, firstImages, secondImages, resumedImages, forkRestartImages: forkRestart.forkRestartImages, }, null, 2));} finally { await session.disposeAsync(); registration.unregister(); rmSync(root, { recursive: true, force: true });}