import { mkdir, mkdtemp, readdir, readFile, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { SessionInfo } from "@earendil-works/pi-coding-agent"; import { afterAll, beforeAll, describe, expect, it } from "vitest"; import { BakArchive } from "../archive.js"; import bakExtension from "../index.js"; import { type S3Config, S3Store } from "../s3.js"; import { type HostIdentity, statePaths } from "../state.js"; import { removeHosts, storageConfig, testAlias, writeSession, } from "./support.js"; const config = storageConfig; interface Note { message: string; level: string; } type CommandHandler = (args: string, ctx: unknown) => Promise | void; type EventHandler = (event: unknown, ctx: unknown) => Promise | void; const aliases: string[] = []; let workspace: string; let sessionDir: string; let notes: Note[] = []; let run: (args: string) => Promise; let emit: (event: string) => Promise; let context: Record; let currentSession: SessionInfo | undefined; /** * Registers the extension exactly like Pi does, but against a throwaway state * directory and session store so real backups stay untouched. */ function loadExtension(): void { const handlers = new Map(); let command: { handler: CommandHandler } | undefined; const pi = { on(event: string, handler: EventHandler) { handlers.set(event, [...(handlers.get(event) ?? []), handler]); }, registerCommand(_name: string, definition: { handler: CommandHandler }) { command = definition; }, }; bakExtension(pi as never); if (!command) throw new Error("bak command was not registered"); const registered = command; run = async (args: string) => { await registered.handler(args, context); }; emit = async (event: string) => { for (const handler of handlers.get(event) ?? []) await handler({ type: event }, context); }; } function sessionContext(): Record { return { hasUI: true, ui: { notify: (message: string, level = "info") => { notes.push({ message, level }); }, setStatus: () => undefined, theme: { fg: (_style: string, text: string) => text }, }, sessionManager: { getSessionDir: () => sessionDir, usesDefaultSessionDir: () => false, getSessionFile: () => currentSession?.path, getHeader: () => currentSession ? { id: currentSession.id, cwd: currentSession.cwd, timestamp: currentSession.created.toISOString(), } : undefined, getSessionName: () => currentSession?.name, getEntries: () => [], }, }; } function lastNote(): Note { const note = notes.at(-1); if (!note) throw new Error("no bak notification recorded"); return note; } async function archivedIds(alias: string): Promise { const cache = await sandboxArchive().refresh(); return ( cache.indexes .find((index) => index.alias === alias) ?.sessions.map((record) => record.id) ?? [] ); } async function waitForArchived(alias: string, id: string): Promise { for (let attempt = 0; attempt < 30; attempt++) { if ((await archivedIds(alias)).includes(id)) return; await new Promise((resolve) => setTimeout(resolve, 500)); } throw new Error(`session ${id} never reached the archive`); } function sandboxArchive(): BakArchive { return new BakArchive( new S3Store(config as S3Config), statePaths(process.env), ); } describe.skipIf(!config)("bak commands against real S3 storage", () => { const original = { state: process.env.XDG_STATE_HOME, cache: process.env.XDG_CACHE_HOME, }; beforeAll(async () => { workspace = await mkdtemp(join(tmpdir(), "bak-cmd-e2e-")); process.env.XDG_STATE_HOME = join(workspace, "state"); process.env.XDG_CACHE_HOME = join(workspace, "cache"); sessionDir = join(workspace, "sessions"); await mkdir(sessionDir, { recursive: true }); context = sessionContext(); loadExtension(); }); afterAll(async () => { if (config && aliases.length > 0) { const catalog = await new S3Store(config).get("catalog/hosts.json"); const hosts = catalog ? ( JSON.parse(new TextDecoder().decode(catalog.body)) as { hosts: HostIdentity[]; } ).hosts : []; await removeHosts( config, hosts .filter((host) => aliases.includes(host.alias)) .map((host) => host.hostId), ); } process.env.XDG_STATE_HOME = original.state; process.env.XDG_CACHE_HOME = original.cache; if (original.state === undefined) delete process.env.XDG_STATE_HOME; if (original.cache === undefined) delete process.env.XDG_CACHE_HOME; await rm(workspace, { recursive: true, force: true }); }); it("reports an uninitialized host", async () => { notes = []; await run("status"); expect(lastNote().message).toContain("not initialized"); expect(lastNote().message).toContain("0 pending"); }); it("initializes, backs up local sessions, and reports status", async () => { const alias = testAlias(); aliases.push(alias); notes = []; await run(`init ${alias}`); expect(lastNote().message).toBe(`bak: initialized as ${alias}`); const identity = JSON.parse( await readFile(statePaths(process.env).identity, "utf8"), ) as HostIdentity; expect(identity.alias).toBe(alias); const first = await writeSession(sessionDir, { name: "command session" }); const second = await writeSession(sessionDir, { messages: ["x", "y"] }); await run("backup"); expect(lastNote().message).toBe("bak: backed up 2; 2 indexed"); await run("status"); expect(lastNote().message).toContain(alias); expect(lastNote().message).toContain("ready"); // The counter spans every catalog host, including unrelated real hosts. expect(lastNote().message).toMatch(/[1-9]\d* indexed/); expect(lastNote().message).toContain("0 pending"); expect(lastNote().message).not.toContain("error"); const cache = await sandboxArchive().refresh(); const ids = cache.indexes .find((index) => index.alias === alias) ?.sessions.map((record) => record.id) ?? []; expect(ids.sort()).toEqual([first.id, second.id].sort()); }); it("searches remote names from the cached catalog", async () => { const alias = aliases.at(-1) as string; notes = []; await run("refresh"); await run("search COMMAND SESSION"); expect(lastNote().message).toContain(`${alias}/`); expect(lastNote().message).toContain("command session"); }); it("refreshes the catalog and restores by alias prefix", async () => { const alias = aliases.at(-1) as string; notes = []; await run("refresh"); expect(lastNote().message).toMatch(/refreshed \d+ host\(s\)/); const cache = await sandboxArchive().refresh(); const record = cache.indexes.find((index) => index.alias === alias) ?.sessions[0]; expect(record).toBeDefined(); const target = join(workspace, "restore-target"); await run(`restore ${alias}/${record?.id.slice(0, 8)} ${target}`); expect(lastNote().message).toBe("bak: restored 1, skipped 0, conflicts 0"); expect(await readdir(target)).toHaveLength(1); await run(`restore ${alias}/${record?.id.slice(0, 8)} ${target}`); expect(lastNote().message).toBe("bak: restored 0, skipped 1, conflicts 0"); }); it("restores every archived session of the host", async () => { const alias = aliases.at(-1) as string; const target = join(workspace, "restore-all"); notes = []; await run(`restore all ${target}`); expect(lastNote().message).toMatch(/restored \d+, skipped \d+/); const cache = await sandboxArchive().refresh(); const own = cache.indexes.find((index) => index.alias === alias)?.sessions.length ?? 0; expect((await readdir(target)).length).toBeGreaterThanOrEqual(own); }); it("keeps working after session lifecycle events", async () => { notes = []; await emit("session_start"); await writeSession(sessionDir, { messages: ["after restart"] }); await run("backup"); expect(lastNote().message).toBe("bak: backed up 3; 3 indexed"); await run("status"); expect(lastNote().message).toMatch(/[1-9]\d* indexed/); expect(lastNote().message).toContain("0 pending"); expect(lastNote().message).not.toContain("error"); }); it("backs up the active session when the agent settles", async () => { const alias = aliases.at(-1) as string; notes = []; currentSession = await writeSession(sessionDir, { messages: ["settled"], name: "settled session", }); await emit("agent_settled"); await waitForArchived(alias, currentSession.id); expect(notes.filter((note) => note.level !== "info")).toEqual([]); const markers = await readdir(statePaths(process.env).pendingDir); expect(markers).toEqual([]); }); it("retains a pending retry after a failed backup and recovers", async () => { const alias = aliases.at(-1) as string; const bucket = process.env.PI_BAK_S3_BUCKET; notes = []; currentSession = await writeSession(sessionDir, { messages: ["pending"] }); process.env.PI_BAK_S3_BUCKET = `${(config as S3Config).bucket}-missing`; loadExtension(); await run("backup"); expect(lastNote().level).toBe("error"); expect(await readdir(statePaths(process.env).pendingDir)).not.toEqual([]); process.env.PI_BAK_S3_BUCKET = bucket; if (bucket === undefined) delete process.env.PI_BAK_S3_BUCKET; loadExtension(); await emit("session_start"); await emit("agent_settled"); await waitForArchived(alias, currentSession.id); for (let attempt = 0; attempt < 30; attempt++) { if ((await readdir(statePaths(process.env).pendingDir)).length === 0) break; await new Promise((resolve) => setTimeout(resolve, 500)); } expect(await readdir(statePaths(process.env).pendingDir)).toEqual([]); }); it("reports usage and initialization errors", async () => { notes = []; await run("bogus"); expect(lastNote().level).toBe("error"); expect(lastNote().message).toContain("Usage:"); await run(`init ${testAlias()}`); expect(lastNote().level).toBe("error"); expect(lastNote().message).toMatch(/Already initialized as/); }); });