import type { ExtensionContext, SessionInfo, } from "@earendil-works/pi-coding-agent"; import { SessionManager } from "@earendil-works/pi-coding-agent"; import type { BakArchive } from "./archive.js"; import { type S3ConfigResolution, s3Config } from "./config.js"; import { closeDebug, dbg } from "./src/debug.ts"; import { type CatalogCache, clearPending, type HostIdentity, markPending, type PendingEntry, readPending, type SessionRecord, statePaths, } from "./state.js"; const ICON = "󰁯"; const STATUS_KEY = "bak"; const SHUTDOWN_DRAIN_MS = 2_000; const SUBCOMMANDS = [ "init", "backup", "restore", "search", "refresh", "status", ] as const; const SEARCH_LIMIT = 10; type ArchiveModule = typeof import("./archive.js"); interface Backend { archive: BakArchive; archiveModule: ArchiveModule; } function currentSession(ctx: ExtensionContext): SessionInfo | undefined { const path = ctx.sessionManager.getSessionFile(); const header = ctx.sessionManager.getHeader(); if (!path || !header) return undefined; return { path, id: header.id, cwd: header.cwd, name: ctx.sessionManager.getSessionName(), parentSessionPath: header.parentSession, created: new Date(header.timestamp), modified: new Date(), messageCount: ctx.sessionManager .getEntries() .filter((entry) => entry.type === "message").length, firstMessage: "", allMessagesText: "", }; } function errorText(error: unknown): string { const text = error instanceof Error ? error.message : String(error); return text.replace(/\s+/g, " ").trim().slice(0, 240); } function formatAge(timestamp: string | undefined): string { if (!timestamp) return "never"; const age = Date.now() - new Date(timestamp).getTime(); if (!Number.isFinite(age) || age < 0) return "unknown"; if (age < 60_000) return "now"; if (age < 3_600_000) return `${Math.floor(age / 60_000)}m ago`; if (age < 86_400_000) return `${Math.floor(age / 3_600_000)}h ago`; return `${Math.floor(age / 86_400_000)}d ago`; } function destinationArgs(args: string): { target: string; path?: string } { const match = /^(\S+)(?:\s+(.+))?$/.exec(args.trim()); return { target: match?.[1] ?? "", path: match?.[2]?.trim() || undefined }; } function selectRecords(cache: CatalogCache, target: string): SessionRecord[] { if (target === "all") return cache.indexes.flatMap((index) => index.sessions); const slash = target.indexOf("/"); const alias = slash >= 0 ? target.slice(0, slash) : undefined; const prefix = slash >= 0 ? target.slice(slash + 1) : target; if (!prefix) return []; return cache.indexes .filter((index) => !alias || index.alias === alias) .flatMap((index) => index.sessions) .filter((record) => record.id.startsWith(prefix)); } export function searchRecords( cache: CatalogCache, query: string, ): { alias: string; record: SessionRecord; prefix: string }[] { const needle = query.trim().toLowerCase(); if (!needle) return []; const matches: { alias: string; record: SessionRecord; ids: string[] }[] = []; for (const index of cache.indexes) { const ids = index.sessions.map((record) => record.id); for (const record of index.sessions) { if ( !record.name?.toLowerCase().includes(needle) && !record.cwd.toLowerCase().includes(needle) ) continue; matches.push({ alias: index.alias, record, ids }); } } matches.sort((a, b) => b.record.created.localeCompare(a.record.created)); return matches.slice(0, SEARCH_LIMIT).map(({ alias, record, ids }) => ({ alias, record, prefix: uniquePrefix(record.id, ids), })); } function uniquePrefix(id: string, ids: string[]): string { for (let length = Math.min(8, id.length); length < id.length; length++) { const prefix = id.slice(0, length); if ( !ids.some((candidate) => candidate !== id && candidate.startsWith(prefix)) ) return prefix; } return id; } async function localSessions(ctx: ExtensionContext): Promise { const manager = ctx.sessionManager as typeof ctx.sessionManager & { usesDefaultSessionDir?: () => boolean; }; // Non-persistent sessions report an empty directory; scanning it would read // the working directory instead of the session store. const directory = ctx.sessionManager.getSessionDir(); return !directory || manager.usesDefaultSessionDir?.() ? SessionManager.listAll() : SessionManager.listAll(directory); } export function createBakRuntime( resolved: S3ConfigResolution = s3Config(), missingWarningShown = false, ) { const paths = statePaths(); let backendPromise: Promise | undefined; const loadBackend = (): Promise => { const config = resolved.config; if (!config) return Promise.resolve(undefined); if (backendPromise) return backendPromise; backendPromise = Promise.all([ import("./archive.js"), import("./s3.js"), ]).then(([archiveModule, { S3Store }]) => ({ archive: new archiveModule.BakArchive(new S3Store(config), paths), archiveModule, })); return backendPromise; }; let identity: HostIdentity | undefined; let cache: CatalogCache | undefined; let lastBackupAt: string | undefined; let lastError: string | undefined; let warnedFailure = false; let warnedSetup = missingWarningShown; let lifecycle = 0; let loadedLifecycle = 0; let backgroundLifecycle = 0; let serial: Promise = Promise.resolve(); let readiness: Promise = Promise.resolve(); let lifecycleController: AbortController | undefined; let pending = new Map(); const notify = ( ctx: ExtensionContext, message: string, level: "info" | "warning" | "error" = "info", ) => { if (ctx.hasUI) ctx.ui.notify(`bak: ${message}`, level); }; const runSerial = (operation: () => Promise): Promise => { const result = serial.then(operation, operation); serial = result.then( () => undefined, () => undefined, ); return result; }; const requireArchive = async ( ctx: ExtensionContext, ): Promise => { const backend = await loadBackend(); if (backend) return backend.archive; notify(ctx, `disabled; missing ${resolved.missing.join(", ")}`, "error"); return undefined; }; const setProgress = (ctx: ExtensionContext, text?: string) => { if (ctx.hasUI) ctx.ui.setStatus( STATUS_KEY, text ? ctx.ui.theme.fg("dim", `${ICON} ${text}`) : undefined, ); }; const addPending = async (path: string): Promise => { const entry = await markPending(paths.pendingDir, path); pending.set(entry.id, entry); return entry; }; const removePending = async (entries: PendingEntry[]) => { await Promise.all( entries.map(async (entry) => { await clearPending(paths.pendingDir, entry.id); pending.delete(entry.id); }), ); }; const ensureLoaded = ( ctx: ExtensionContext, epoch = lifecycle, ): Promise => { if (!resolved.config || loadedLifecycle === epoch) return readiness; loadedLifecycle = epoch; readiness = (async () => { const backend = await loadBackend(); if (!backend) return; const [savedIdentity, savedPending] = await Promise.all([ backend.archive.identity(), readPending(paths.pendingDir), ]); let savedCache: CatalogCache | undefined; try { savedCache = await backend.archive.cachedCatalog(); } catch { savedCache = undefined; } if (epoch !== lifecycle) return; identity = savedIdentity; cache = savedCache; pending = new Map(savedPending.map((entry) => [entry.id, entry])); if (!identity && !warnedSetup) { notify(ctx, "run /bak init ", "warning"); warnedSetup = true; } })(); return readiness; }; const refresh = async ( ctx: ExtensionContext, signal?: AbortSignal, ): Promise => { const client = await requireArchive(ctx); if (!client) return undefined; const next = await client.refresh(signal); cache = next; return next; }; const backupBatch = async ( infos: SessionInfo[], onProgress?: (done: number, total: number) => void, replayEntries?: PendingEntry[], ) => { if (infos.length === 0) return { uploaded: 0, indexed: 0 }; const backend = await loadBackend(); if (!backend) return { uploaded: 0, indexed: 0 }; const archive = backend.archive; await readiness; if (!replayEntries) await Promise.all(infos.map((info) => addPending(info.path))); const paths = new Set(infos.map((info) => info.path)); const covered = (replayEntries ?? [...pending.values()]).filter((entry) => paths.has(entry.path), ); if (!identity) throw new Error("Run /bak init first"); try { const result = await archive.backup( infos, onProgress, lifecycleController?.signal, ); await removePending(covered); lastBackupAt = new Date().toISOString(); lastError = undefined; return result; } catch (error) { if (error instanceof backend.archiveModule.PartialBackupError) { const succeeded = new Set(error.succeededPaths); await removePending( covered.filter((entry) => succeeded.has(entry.path)), ); } throw error; } }; const startBackground = (ctx: ExtensionContext) => { const epoch = lifecycle; if (!resolved.config || backgroundLifecycle === epoch) return; backgroundLifecycle = epoch; void ensureLoaded(ctx, epoch) .then(async () => { if (!identity || epoch !== lifecycle) return; const backend = await loadBackend(); if (!backend) return; try { await backend.archive.ensureInitialized( identity, lifecycleController?.signal, ); } catch (error) { if (epoch === lifecycle) notify( ctx, `background init failed: ${errorText(error)}`, "warning", ); } void refresh(ctx, lifecycleController?.signal).catch((error) => { if (epoch === lifecycle && !cache) notify( ctx, `catalog refresh failed: ${errorText(error)}`, "warning", ); }); await retryPending(epoch); }) .catch((error) => { if (epoch === lifecycle) notify( ctx, `background state failed: ${errorText(error)}`, "warning", ); }); }; const automaticBackup = (ctx: ExtensionContext) => { startBackground(ctx); const info = currentSession(ctx); if (!resolved.config || !info) { dbg?.("operation.finish", { operation: "automatic-backup", outcome: "skipped", }); return; } const epoch = lifecycle; void runSerial(async () => { const statusTimer = setTimeout(() => { if (epoch === lifecycle) setProgress(ctx, "↑"); }, 300); try { await backupBatch([info]); dbg?.("operation.finish", { operation: "automatic-backup", outcome: "success", }); if (warnedFailure && epoch === lifecycle) { notify(ctx, "backup recovered"); warnedFailure = false; } } catch (error) { dbg?.("operation.finish", { operation: "automatic-backup", outcome: "failed", }); lastError = errorText(error); if (!warnedFailure && epoch === lifecycle) { notify(ctx, `backup failed: ${lastError}; retry pending`, "warning"); warnedFailure = true; } } finally { clearTimeout(statusTimer); if (epoch === lifecycle) setProgress(ctx); } }); }; const retryPending = async (epoch: number) => { if (!resolved.config || pending.size === 0 || epoch !== lifecycle) return; const backend = await loadBackend(); if (!backend) return; const replayEntries = [...pending.values()]; const infos: SessionInfo[] = []; for (const path of new Set(replayEntries.map((entry) => entry.path))) { try { infos.push(await backend.archiveModule.sessionInfoFromPath(path)); } catch (error) { lastError = `pending ${path}: ${errorText(error)}`; } } if (infos.length === 0) return; await runSerial(() => backupBatch(infos, undefined, replayEntries)); }; const sessionStart = (ctx: ExtensionContext) => { lifecycle++; lifecycleController?.abort(); lifecycleController = new AbortController(); loadedLifecycle = 0; backgroundLifecycle = 0; readiness = Promise.resolve(); identity = undefined; cache = undefined; pending = new Map(); if (!resolved.config && !warnedSetup) { notify( ctx, `disabled; missing ${resolved.missing.join(", ")}`, "warning", ); warnedSetup = true; } }; /** With automatic=false no upload starts; an explicit /bak backup still drains. */ const sessionShutdown = async (ctx: ExtensionContext, automatic = true) => { const info = currentSession(ctx); if (automatic && resolved.config && info) { await addPending(info.path).catch((error) => { lastError = errorText(error); }); } if (automatic) automaticBackup(ctx); await Promise.race([ serial, new Promise((resolve) => setTimeout(resolve, SHUTDOWN_DRAIN_MS)), ]); lifecycleController?.abort(); lifecycle++; setProgress(ctx); }; const command = { getArgumentCompletions(prefix: string) { const normalized = prefix.trimStart(); const restore = /^restore\s+(\S*)$/i.exec(normalized); if (restore) { const query = restore[1]?.toLowerCase() ?? ""; const records = cache?.indexes.flatMap((index) => { const ids = index.sessions.map((record) => record.id); return index.sessions.map((record) => { const prefix = uniquePrefix(record.id, ids); return { value: `restore ${index.alias}/${prefix}`, label: `${index.alias}/${prefix}`, description: record.name ?? record.cwd, }; }); }) ?? []; const items = [ { value: "restore all", label: "all", description: "all sessions" }, ...records, ].filter((item) => item.label.toLowerCase().startsWith(query)); return items.length > 0 ? items.slice(0, 100) : null; } if (normalized.includes(" ")) return null; const items = SUBCOMMANDS.filter((command) => command.startsWith(normalized.toLowerCase()), ).map((command) => ({ value: command, label: command })); return items.length > 0 ? items : null; }, async handler(args: string, ctx: ExtensionContext) { const trimmed = args.trim(); const [command = "status", ...rest] = trimmed.split(/\s+/); const subcommand = SUBCOMMANDS.find((candidate) => candidate === command) ?? "invalid"; let outcome: "success" | "failed" | "skipped" = "success"; try { if (command === "status") { if (!resolved.config) outcome = "skipped"; else await ensureLoaded(ctx).catch((error) => { outcome = "failed"; lastError = errorText(error); }); const indexed = cache?.indexes.reduce( (total, index) => total + index.sessions.length, 0, ); const state = resolved.config ? identity ? "ready" : "not initialized" : `disabled; missing ${resolved.missing.join(", ")}`; notify( ctx, [ identity?.alias, state, indexed === undefined ? undefined : `${indexed} indexed`, `${new Set([...pending.values()].map((entry) => entry.path)).size} pending`, `last backup ${formatAge(lastBackupAt)}`, `catalog ${formatAge(cache?.refreshedAt)}`, lastError ? `error ${lastError}` : undefined, ] .filter(Boolean) .join(" · "), lastError ? "warning" : "info", ); return; } const client = await requireArchive(ctx); if (!client) { outcome = "skipped"; return; } await ensureLoaded(ctx); if (command === "init") { if (rest.length !== 1) throw new Error("Usage: /bak init "); setProgress(ctx, "initializing"); identity = await runSerial(() => client.init(rest[0] ?? "", lifecycleController?.signal), ); cache = await refresh(ctx, lifecycleController?.signal); warnedSetup = false; void retryPending(lifecycle).catch((error) => { lastError = errorText(error); notify(ctx, `pending retry failed: ${lastError}`, "warning"); }); notify(ctx, `initialized as ${identity.alias}`); return; } if (command === "refresh") { if (rest.length > 0) throw new Error("Usage: /bak refresh"); setProgress(ctx, "refreshing"); const next = await runSerial(() => client.refresh(lifecycleController?.signal), ); cache = next; notify(ctx, `refreshed ${next.indexes.length} host(s)`); return; } if (command === "backup") { if (rest.length > 0) throw new Error("Usage: /bak backup"); setProgress(ctx, "scanning"); const infos = await localSessions(ctx); const result = await runSerial(() => backupBatch(infos, (done, total) => setProgress(ctx, `↑ ${done}/${total}`), ), ); notify( ctx, `backed up ${result.uploaded}; ${result.indexed} indexed`, ); return; } if (command === "search") { const query = rest.join(" ").trim(); if (!query) throw new Error("Usage: /bak search "); if (!cache) { setProgress(ctx, "refreshing"); cache = await runSerial(() => client.refresh(lifecycleController?.signal), ); } const results = searchRecords(cache, query); notify( ctx, results.length ? results .map( ({ alias, record, prefix }) => `${alias}/${prefix} ${record.name ?? record.cwd} (${record.created.slice(0, 10)})`, ) .join("\n") + (results.length === SEARCH_LIMIT ? "\nShowing first 10; narrow query for more" : "") : `no sessions match ${query}`, ); return; } if (command === "restore") { const parsed = destinationArgs(rest.join(" ")); if (!parsed.target) throw new Error( "Usage: /bak restore /|all [path]", ); setProgress(ctx, "refreshing"); const current = await runSerial(() => client.refresh(lifecycleController?.signal), ); cache = current; const records = selectRecords(current, parsed.target); if (records.length === 0) throw new Error(`No session matches ${parsed.target}`); if (parsed.target !== "all" && records.length !== 1) throw new Error(`Ambiguous session: ${parsed.target}`); const result = await runSerial(() => client.restore( records, parsed.path ?? ctx.sessionManager.getSessionDir(), (done, total) => setProgress(ctx, `↓ ${done}/${total}`), lifecycleController?.signal, ), ); notify( ctx, `restored ${result.restored}, skipped ${result.skipped}, conflicts ${result.conflicts}`, result.conflicts > 0 ? "warning" : "info", ); return; } throw new Error( "Usage: /bak init | backup | restore /|all [path] | search | refresh | status", ); } catch (error) { outcome = "failed"; lastError = errorText(error); notify(ctx, lastError, "error"); } finally { dbg?.("operation.finish", { operation: "command", outcome, subcommand, }); if (command !== "status") setProgress(ctx); } }, }; return { automaticBackup, closeDebugWhenIdle() { void serial.finally(closeDebug); }, command, sessionShutdown, sessionStart, }; }