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.
15 kB · 443 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444import { appendFileSync, chmodSync, type chownSync, existsSync, lstatSync, mkdirSync, mkdtempSync, readdirSync, readFileSync, type renameSync, rmSync, statSync, symlinkSync, type writeFileSync, type writeSync as writeSyncFs,} from "node:fs";import { tmpdir } from "node:os";import { basename, dirname, join } from "node:path";import { afterEach, describe, expect, it, vi } from "vitest";
type ChmodSync = typeof chmodSync;type ChownSync = typeof chownSync;type RenameSync = typeof renameSync;type WriteFileSync = typeof writeFileSync;type WriteSync = typeof writeSyncFs;
const fsMocks = vi.hoisted(() => ({ actualWriteFileSync: undefined as WriteFileSync | undefined, actualWriteSync: undefined as WriteSync | undefined, chmodSync: vi.fn<ChmodSync>(), chownSync: vi.fn<ChownSync>(), renameSync: vi.fn<RenameSync>(), writeFileSync: vi.fn<WriteFileSync>(), writeSync: vi.fn<WriteSync>(),}));vi.mock("node:fs", async (importOriginal) => { const actual = await importOriginal<typeof import("node:fs")>(); fsMocks.actualWriteFileSync = actual.writeFileSync; fsMocks.actualWriteSync = actual.writeSync; fsMocks.chmodSync.mockImplementation(actual.chmodSync); fsMocks.chownSync.mockImplementation(actual.chownSync); fsMocks.renameSync.mockImplementation(actual.renameSync); fsMocks.writeFileSync.mockImplementation(actual.writeFileSync); fsMocks.writeSync.mockImplementation(actual.writeSync); return { ...actual, chmodSync: fsMocks.chmodSync, chownSync: fsMocks.chownSync, renameSync: fsMocks.renameSync, writeFileSync: fsMocks.writeFileSync, writeSync: fsMocks.writeSync, };});
import { SessionManager } from "../src/core/session-manager.js";
const tempDirs: string[] = [];
afterEach(() => { while (tempDirs.length > 0) { const dir = tempDirs.pop()!; rmSync(dir, { recursive: true, force: true }); }});
function createTempDir(): string { const dir = mkdtempSync(join(tmpdir(), "pi-flush-test-")); tempDirs.push(dir); return dir;}
describe("SessionManager.flushNow", () => { it("writes the session file with all in-memory entries before any assistant message", () => { const dir = createTempDir(); const sessionDir = join(dir, "sessions"); const mgr = SessionManager.create(dir, sessionDir);
// Append a custom entry (goal state) — normally suppressed by _persist // because no assistant message exists yet. mgr.appendCustomEntry("thread_goal_state", { active: true, status: "active" });
const file = mgr.getSessionFile()!; expect(existsSync(file)).toBe(false);
// flushNow forces a write despite the no-assistant guard mgr.flushNow(); expect(existsSync(file)).toBe(true);
const content = readFileSync(file, "utf8"); const lines = content.trim().split("\n"); expect(lines.length).toBe(2); // header + custom entry
const header = JSON.parse(lines[0]!); expect(header.type).toBe("session");
const custom = JSON.parse(lines[1]!); expect(custom.type).toBe("custom"); expect(custom.customType).toBe("thread_goal_state"); expect(custom.data).toEqual({ active: true, status: "active" }); });
it("preserves live bytes and cleans up the temp file when a rewrite fails", () => { const dir = createTempDir(); const mgr = SessionManager.create(dir, join(dir, "sessions")); mgr.appendCustomEntry("before_failure"); mgr.flushNow(); const file = mgr.getSessionFile()!; const before = readFileSync(file); const tempPrefix = `${basename(file)}.`;
mgr.appendMessage({ role: "user", content: "pending", timestamp: Date.now() }); fsMocks.writeSync.mockImplementationOnce(((fd: number, data: string) => { fsMocks.actualWriteSync!(fd, Buffer.from(String(data)).subarray(0, 12)); throw new Error("disk full"); }) as unknown as WriteSync);
expect(() => mgr.flushNow()).toThrow("disk full"); expect(readFileSync(file)).toEqual(before); expect(readdirSync(dirname(file)).filter((name) => name.startsWith(tempPrefix) && name.endsWith(".tmp"))).toEqual( [], ); });
it("preserves live-file metadata before atomically replacing it", () => { const dir = createTempDir(); const mgr = SessionManager.create(dir, join(dir, "sessions")); mgr.appendCustomEntry("before_metadata_check"); mgr.flushNow(); const file = mgr.getSessionFile()!; chmodSync(file, 0o660); const before = statSync(file); fsMocks.chownSync.mockClear(); fsMocks.chownSync.mockImplementationOnce(() => undefined); fsMocks.chmodSync.mockClear(); fsMocks.renameSync.mockClear();
mgr.appendMessage({ role: "user", content: "pending", timestamp: Date.now() }); mgr.flushNow();
const tempPath = fsMocks.chownSync.mock.calls[0]?.[0]; expect(tempPath).toEqual(expect.any(String)); expect(fsMocks.chownSync).toHaveBeenCalledWith(tempPath, before.uid, before.gid); expect(fsMocks.chmodSync).toHaveBeenCalledWith(tempPath, before.mode & 0o777); expect(fsMocks.renameSync).toHaveBeenCalledWith(tempPath, join(dirname(tempPath as string), basename(file))); expect(fsMocks.chownSync.mock.invocationCallOrder[0]!).toBeLessThan( fsMocks.renameSync.mock.invocationCallOrder[0]!, ); expect(fsMocks.chmodSync.mock.invocationCallOrder[0]!).toBeLessThan( fsMocks.renameSync.mock.invocationCallOrder[0]!, ); const after = statSync(file); expect({ mode: after.mode & 0o777, uid: after.uid, gid: after.gid }).toEqual({ mode: before.mode & 0o777, uid: before.uid, gid: before.gid, }); });
it("preserves live bytes and cleans the temp file when restoring ownership fails", () => { const dir = createTempDir(); const mgr = SessionManager.create(dir, join(dir, "sessions")); mgr.appendCustomEntry("before_ownership_failure"); mgr.flushNow(); const file = mgr.getSessionFile()!; const before = readFileSync(file); const tempPrefix = `${basename(file)}.`; const permissionError = Object.assign(new Error("operation not permitted"), { code: "EPERM" }); fsMocks.chownSync.mockImplementationOnce(() => { throw permissionError; }); fsMocks.renameSync.mockClear();
mgr.appendMessage({ role: "user", content: "pending", timestamp: Date.now() });
expect(() => mgr.flushNow()).toThrow(permissionError); expect(fsMocks.renameSync).not.toHaveBeenCalled(); expect(readFileSync(file)).toEqual(before); expect(readdirSync(dirname(file)).filter((name) => name.startsWith(tempPrefix) && name.endsWith(".tmp"))).toEqual( [], ); });
it("rewrites a cross-directory symlink target without replacing the alias", () => { const dir = createTempDir(); const targetDir = join(dir, "targets"); const aliasDir = join(dir, "aliases"); mkdirSync(targetDir); mkdirSync(aliasDir); const target = join(targetDir, "session.jsonl"); const alias = join(aliasDir, "session.jsonl"); fsMocks.actualWriteFileSync!( target, `${JSON.stringify({ type: "session", version: 2, id: "symlink-session", timestamp: "2026-01-01T00:00:00.000Z", cwd: dir, })}\n`, ); chmodSync(target, 0o640); symlinkSync(target, alias);
const mgr = SessionManager.open(alias);
expect(mgr.getSessionFile()).toBe(alias); expect(lstatSync(alias).isSymbolicLink()).toBe(true); expect(JSON.parse(readFileSync(target, "utf8")).version).toBe(3); expect(statSync(target).mode & 0o777).toBe(0o640); });
it("is a no-op for in-memory (non-persisted) sessions", () => { const mgr = SessionManager.inMemory("/tmp"); mgr.appendCustomEntry("thread_goal_state", { active: true }); mgr.flushNow(); expect(mgr.getSessionFile()).toBeUndefined(); });
it("preserves subsequent rewrite behavior after flush with user then assistant", () => { const dir = createTempDir(); const sessionDir = join(dir, "sessions"); const mgr = SessionManager.create(dir, sessionDir);
mgr.appendCustomEntry("thread_goal_state", { active: true, status: "active" }); mgr.flushNow();
const file = mgr.getSessionFile()!; expect(existsSync(file)).toBe(true);
// Append a USER message first — this sets flushed=false because // no assistant exists yet, exercising the pending full-rewrite path. mgr.appendMessage({ role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now(), });
// Then append an ASSISTANT message — this triggers the full rewrite // (not appendFileSync) because flushed is false. mgr.appendMessage({ role: "assistant", content: [{ type: "text", text: "hi" }], api: "openai-completions", provider: "openai", model: "test", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), });
const content = readFileSync(file, "utf8"); const lines = content.trim().split("\n"); expect(lines.length).toBe(4);
const header = JSON.parse(lines[0]!); expect(header.type).toBe("session");
const custom = JSON.parse(lines[1]!); expect(custom.type).toBe("custom"); expect(custom.customType).toBe("thread_goal_state");
const userMsg = JSON.parse(lines[2]!); expect(userMsg.type).toBe("message"); expect(userMsg.message.role).toBe("user");
const assistantMsg = JSON.parse(lines[3]!); expect(assistantMsg.type).toBe("message"); expect(assistantMsg.message.role).toBe("assistant");
expect(custom.parentId).toBeNull(); expect(userMsg.parentId).toBe(custom.id); expect(assistantMsg.parentId).toBe(userMsg.id); });
it("does not rewrite an already-flushed session after an appended entry", () => { const dir = createTempDir(); const sessionDir = join(dir, "sessions"); const mgr = SessionManager.create(dir, sessionDir);
mgr.appendMessage({ role: "assistant", content: [{ type: "text", text: "ready" }], api: "openai-completions", provider: "openai", model: "test", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), }); const rewriteFile = vi.spyOn(mgr as unknown as { _rewriteFile(): void }, "_rewriteFile");
mgr.appendCustomEntry("thread_goal_state", { active: true, status: "active" }); mgr.flushNow();
expect(rewriteFile).not.toHaveBeenCalled(); const entries = readFileSync(mgr.getSessionFile()!, "utf8") .trim() .split("\n") .map((line) => JSON.parse(line)); expect(entries.at(-1)).toMatchObject({ type: "custom", customType: "thread_goal_state", data: { active: true, status: "active" }, }); });});
function createPersistedSessionForRollbackTest(): { mgr: SessionManager; file: string; before: Buffer } { const dir = createTempDir(); const mgr = SessionManager.create(dir, join(dir, "sessions")); mgr.appendMessage({ role: "user", content: "hi", timestamp: Date.now() }); mgr.appendMessage({ role: "assistant", content: [{ type: "text", text: "hello" }], api: "openai-completions", provider: "openai", model: "test", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), }); const file = mgr.getSessionFile()!; return { mgr, file, before: readFileSync(file) };}
function failNextOutcomeAppend(mgr: SessionManager, file: string): void { const internals = mgr as unknown as { _persist(entry: unknown): void }; internals._persist = () => { appendFileSync(file, '{"type":"custom_message","truncat'); throw new Error("append failed"); };}
const failAfterPartialTempWrite: WriteSync = (fd: number, data: NodeJS.ArrayBufferView | string) => { const bytes = typeof data === "string" ? Buffer.from(data) : Buffer.from(data.buffer, data.byteOffset, data.byteLength); fsMocks.actualWriteSync!(fd, bytes.subarray(0, 12)); throw new Error("repair failed");};
describe("SessionManager.appendCustomMessageEntryWithRollback", () => { it("flushes rollback-aware value entries immediately", () => { const dir = createTempDir(); const mgr = SessionManager.create(dir, join(dir, "sessions"));
mgr.appendCustomEntryWithRollback("rlm_max_depth_state", { maxDepth: 2 });
expect(readFileSync(mgr.getSessionFile()!, "utf8")).toContain( '"customType":"rlm_max_depth_state","data":{"maxDepth":2}', ); });
it("flushes rollback-aware custom messages before the first assistant response", () => { const dir = createTempDir(); const sessionDir = join(dir, "sessions"); const mgr = SessionManager.create(dir, sessionDir);
mgr.appendCustomMessageEntryWithRollback("compaction_outcome", "failed early", true);
const persisted = readFileSync(mgr.getSessionFile()!, "utf8"); expect(persisted).toContain('"customType":"compaction_outcome"'); });
it("repairs the session file immediately after a torn append", () => { const dir = createTempDir(); const sessionDir = join(dir, "sessions"); const mgr = SessionManager.create(dir, sessionDir); mgr.appendMessage({ role: "user", content: "hi", timestamp: Date.now() }); mgr.appendMessage({ role: "assistant", content: [{ type: "text", text: "hello" }], api: "openai-completions", provider: "openai", model: "test", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), }); const file = mgr.getSessionFile()!; const before = readFileSync(file, "utf8");
const internals = mgr as unknown as { _persist(entry: unknown): void }; const originalPersist = internals._persist.bind(mgr); internals._persist = () => { appendFileSync(file, '{"type":"custom_message","truncat'); throw new Error("disk full"); }; expect(() => mgr.appendCustomMessageEntryWithRollback("test.outcome", "details", false)).toThrow("disk full"); internals._persist = originalPersist;
// The rollback rewrites the file from the restored in-memory entries, so the // durable session has no torn tail even if the process exits right after. expect(readFileSync(file, "utf8")).toBe(before); });
it("preserves old history when rollback repair fails", () => { const { mgr, file, before } = createPersistedSessionForRollbackTest(); failNextOutcomeAppend(mgr, file); fsMocks.writeSync.mockImplementationOnce(failAfterPartialTempWrite);
expect(() => mgr.appendCustomMessageEntryWithRollback("test.outcome", "details", false)).toThrow("append failed"); expect(readFileSync(file).subarray(0, before.length)).toEqual(before); });
it("repairs a torn tail on the next successful retry", () => { const { mgr, file, before } = createPersistedSessionForRollbackTest(); failNextOutcomeAppend(mgr, file); fsMocks.writeSync.mockImplementationOnce(failAfterPartialTempWrite); expect(() => mgr.appendCustomMessageEntryWithRollback("test.outcome", "details", false)).toThrow();
mgr.flushNow(); expect(readFileSync(file)).toEqual(before); });});