pi-export-my-chat 1.0.0

package/helpers.mjs

@@ -0,0 +1,509 @@
+/**
+ * Pure helpers for export-my-chat.
+ *
+ * Everything here is dependency-free (node builtins only) and side-effect-light
+ * so it can be unit-tested directly with `node --test`. The extension's index.ts
+ * imports these; the tests import them too. No pi imports live in this file —
+ * that keeps the logic portable and testable outside a pi runtime.
+ */
+
+import { createHash, randomBytes } from "node:crypto";
+import { lstat, open, stat, unlink } from "node:fs/promises";
+import { dirname, extname, isAbsolute, join, resolve } from "node:path";
+
+// ─── Constants ─────────────────────────────────────────────────────────────
+
+export const EXPORT_FORMAT = "pi-my-chat-export";
+export const SCHEMA_VERSION = 1;
+const TIMING_ENTRY_TYPE = "operation-timing";
+
+// ─── Small utilities ────────────────────────────────────────────────────────
+
+export function sha256Hex(text) {
+    return createHash("sha256").update(text, "utf8").digest("hex");
+}
+
+export function errorMessage(error) {
+    return error instanceof Error ? error.message : String(error);
+}
+
+/** 8-char hex entry id, matching pi's SessionEntryBase.id convention. */
+export function makeEntryId() {
+    return randomBytes(4).toString("hex");
+}
+
+/** Filesystem-safe UTC ISO timestamp for export names: 20260413T123456-789Z */
+export function fileTimestamp(date = new Date()) {
+    return date.toISOString().replace(/[-:]/g, "").replace(".", "-");
+}
+
+export function timestampedExportName(date = new Date()) {
+    return `pi-my-chat-export-${fileTimestamp(date)}.json`;
+}
+
+/**
+ * The timestamp portion of pi's own session filename convention, byte-for-byte:
+ * `new Date().toISOString().replace(/[:.]/g, "-")` (e.g. 2026-04-13T12-34-56-789Z).
+ * Revived sessions must be named exactly like pi names its own so they are
+ * indistinguishable in the session directory.
+ */
+export function sessionFileTimestamp(date = new Date()) {
+    return date.toISOString().replace(/[:.]/g, "-");
+}
+
+function finiteNonNegative(value) {
+    return typeof value === "number" && Number.isFinite(value) ? Math.max(0, value) : 0;
+}
+
+// ─── Usage accounting ───────────────────────────────────────────────────────
+
+function emptyUsageTotals() {
+    return {
+        usageRecords: 0,
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 0,
+        cost: {
+            input: 0,
+            output: 0,
+            cacheRead: 0,
+            cacheWrite: 0,
+            total: 0,
+        },
+    };
+}
+
+function addUsage(totals, raw) {
+    if (!raw || typeof raw !== "object") return;
+    const input = finiteNonNegative(raw.input);
+    const output = finiteNonNegative(raw.output);
+    const cacheRead = finiteNonNegative(raw.cacheRead);
+    const cacheWrite = finiteNonNegative(raw.cacheWrite);
+    const calculatedTokens = input + output + cacheRead + cacheWrite;
+    const reportedTokens = finiteNonNegative(raw.totalTokens);
+    const rawCost = raw.cost && typeof raw.cost === "object" ? raw.cost : {};
+    const inputCost = finiteNonNegative(rawCost.input);
+    const outputCost = finiteNonNegative(rawCost.output);
+    const cacheReadCost = finiteNonNegative(rawCost.cacheRead);
+    const cacheWriteCost = finiteNonNegative(rawCost.cacheWrite);
+    const calculatedCost = inputCost + outputCost + cacheReadCost + cacheWriteCost;
+    const reportedCost = finiteNonNegative(rawCost.total);
+
+    totals.usageRecords++;
+    totals.input += input;
+    totals.output += output;
+    totals.cacheRead += cacheRead;
+    totals.cacheWrite += cacheWrite;
+    // Prefer the provider-reported total when present; fall back to the sum of
+    // components. Either way the value is clamped to a finite non-negative.
+    totals.totalTokens += reportedTokens || calculatedTokens;
+    totals.cost.input += inputCost;
+    totals.cost.output += outputCost;
+    totals.cost.cacheRead += cacheReadCost;
+    totals.cost.cacheWrite += cacheWriteCost;
+    totals.cost.total += reportedCost || calculatedCost;
+}
+
+/**
+ * Sum authoritative persisted usage without counting timing records, whose
+ * agent rows overlap child operations. Counts assistant messages, nested model
+ * usage reported by tool results, compaction, and branch summaries — each of
+ * which carries a `usage` object in the session tree.
+ */
+export function collectUsage(entries) {
+    const totals = emptyUsageTotals();
+    for (const entry of entries) {
+        if (entry?.type === "message") {
+            const message = entry.message;
+            if (message?.role === "assistant" || message?.role === "toolResult") addUsage(totals, message.usage);
+        } else if (entry?.type === "compaction" || entry?.type === "branch_summary") {
+            addUsage(totals, entry.usage);
+        }
+    }
+    return totals;
+}
+
+// ─── Timing telemetry ───────────────────────────────────────────────────────
+
+export function isTimingRecord(value) {
+    return Boolean(
+        value &&
+        typeof value === "object" &&
+        value.schemaVersion === 1 &&
+        typeof value.operationId === "string" &&
+        typeof value.kind === "string" &&
+        typeof value.startedAt === "number" &&
+        typeof value.endedAt === "number" &&
+        typeof value.durationMs === "number",
+    );
+}
+
+export function collectTimingRecords(entries) {
+    return entries
+        .filter((entry) => entry?.type === "custom" && entry.customType === TIMING_ENTRY_TYPE && isTimingRecord(entry.data))
+        .map((entry) => entry.data);
+}
+
+/**
+ * Per-kind duration summaries. Kinds are kept separate on purpose: an agent
+ * total overlaps its child model/tool records, so summing kinds would
+ * double-count elapsed wall time.
+ */
+export function summarizeTimings(records) {
+    const byKind = {};
+    for (const record of records) {
+        const kind = record.kind;
+        const durationMs = finiteNonNegative(record.durationMs);
+        const current = byKind[kind] ?? {
+            count: 0,
+            totalDurationMs: 0,
+            averageDurationMs: 0,
+            minDurationMs: durationMs,
+            maxDurationMs: durationMs,
+            statuses: {},
+        };
+        current.count++;
+        current.totalDurationMs += durationMs;
+        current.minDurationMs = Math.min(current.minDurationMs, durationMs);
+        current.maxDurationMs = Math.max(current.maxDurationMs, durationMs);
+        const status = typeof record.status === "string" ? record.status : "unknown";
+        current.statuses[status] = (current.statuses[status] ?? 0) + 1;
+        byKind[kind] = current;
+    }
+    for (const summary of Object.values(byKind)) {
+        summary.averageDurationMs = summary.count > 0 ? summary.totalDurationMs / summary.count : 0;
+    }
+    return {
+        recordCount: records.length,
+        byKind,
+        note: "Durations for agent, model, tool, compaction, and user-wait records can overlap; do not sum different kinds as elapsed wall time.",
+    };
+}
+
+// ─── Context window snapshot ─────────────────────────────────────────────────
+
+/**
+ * Live context-window snapshot: what the provider would receive right now,
+ * versus what has been persisted historically. `contextUsage` is pi's
+ * getContextUsage() shape ({ tokens }); `contextWindow` is the model's window.
+ */
+export function contextUsageSnapshot(contextUsage, contextWindow) {
+    // The window may come from the model or from the usage snapshot itself;
+    // prefer the explicit model window, fall back to the snapshot's own.
+    const candidates = [contextWindow, contextUsage?.contextWindow];
+    const window = candidates.find(
+        (value) => typeof value === "number" && Number.isFinite(value) && value > 0,
+    ) ?? 0;
+    // pi reports tokens === null when it cannot estimate (e.g. right after
+    // compaction, before the next response). Preserve that as null — recording
+    // 0 would understate real usage.
+    const rawUsed = contextUsage?.tokens;
+    const used = typeof rawUsed === "number" && Number.isFinite(rawUsed) ? Math.max(0, rawUsed) : null;
+    return {
+        used,
+        window,
+        remaining: used === null ? null : Math.max(0, window - used),
+        percentUsed: used !== null && window > 0 ? Math.round((used / window) * 100) : null,
+    };
+}
+
+// ─── Destination path resolution (export output) ─────────────────────────────
+
+function unwrapPathArgument(rawArgument) {
+    const trimmed = rawArgument.trim();
+    if (!trimmed) return undefined;
+    const first = trimmed[0];
+    const last = trimmed[trimmed.length - 1];
+    if (first === '"' || first === "'") {
+        if (last !== first || trimmed.length < 2) throw new Error("Path argument has an unmatched quote.");
+        const unwrapped = trimmed.slice(1, -1);
+        if (!unwrapped) throw new Error("Path argument cannot be empty.");
+        return unwrapped;
+    }
+    if (last === "'" || last === '"') throw new Error("Path argument has an unmatched quote.");
+    return trimmed;
+}
+
+/**
+ * Resolve the export destination. Rules (strict by design — this path WRITES):
+ *   (none)               -> <cwd>/pi-my-chat-export-<UTC timestamp>.json
+ *   <absolute dir>       -> invalid: must be a .json file path, not a directory
+ *   <absolute .json>      -> exactly that file; must not already exist
+ *   relative path         -> invalid: absolute paths required
+ * Never overwrites; never creates parent directories.
+ */
+export function resolveExportPath(rawArgument, cwd, date = new Date()) {
+    const supplied = unwrapPathArgument(rawArgument ?? "");
+    if (!supplied) return join(resolve(cwd), timestampedExportName(date));
+    if (supplied.includes("\0")) throw new Error("Export destination cannot contain a NUL byte.");
+    if (!isAbsolute(supplied)) {
+        throw new Error(`Export destination must be an absolute path: ${supplied}`);
+    }
+    const resolved = resolve(supplied);
+    if (!resolved.toLowerCase().endsWith(".json")) {
+        throw new Error(`Export destination must have a .json extension: ${supplied}`);
+    }
+    return resolved;
+}
+
+/**
+ * Parse the revive input path (this path is READ, so it is forgiving where the
+ * export destination is strict): relative paths resolve against cwd, quoted
+ * paths unwrap. Returns { path, force, ok, error }.
+ */
+export function parseReviveArgs(rawArgument, cwd) {
+    let force = false;
+    let rest = (rawArgument ?? "").trim();
+
+    // Pull --force out wherever it appears (space-separated).
+    const tokens = [];
+    for (const part of rest.split(/\s+/)) {
+        if (part === "--force") force = true;
+        else if (part) tokens.push(part);
+    }
+    rest = tokens.join(" ");
+
+    const supplied = unwrapPathArgument(rest);
+    if (!supplied) return { ok: false, force, error: "Provide the path to an /export-my-chat JSON file." };
+    if (supplied.includes("\0")) return { ok: false, force, error: "Path cannot contain a NUL byte." };
+    return { ok: true, force, path: isAbsolute(supplied) ? resolve(supplied) : resolve(cwd, supplied) };
+}
+
+// ─── Exclusive private JSON write ────────────────────────────────────────────
+
+async function destinationExists(path) {
+    try {
+        return await lstat(path);
+    } catch (error) {
+        if (error?.code === "ENOENT") return undefined;
+        throw error;
+    }
+}
+
+/** Validate and exclusively create a private JSON export. Partial files are removed after write failures. */
+export async function writeExclusiveJson(destination, value) {
+    const existing = await destinationExists(destination);
+    if (existing) {
+        if (existing.isDirectory()) throw new Error(`Export destination is a directory, not a JSON file: ${destination}`);
+        if (existing.isSymbolicLink()) throw new Error(`Export destination already exists as a symbolic link: ${destination}`);
+        throw new Error(`Export destination already exists: ${destination}`);
+    }
+    if (extname(destination) !== ".json") {
+        throw new Error(`Export destination must have a .json extension: ${destination}`);
+    }
+
+    let parent;
+    try {
+        parent = await stat(dirname(destination));
+    } catch (error) {
+        if (error?.code === "ENOENT") throw new Error(`Export parent directory does not exist: ${dirname(destination)}`);
+        throw error;
+    }
+    if (!parent.isDirectory()) throw new Error(`Export parent path is not a directory: ${dirname(destination)}`);
+
+    let serialized;
+    try {
+        serialized = `${JSON.stringify(value, null, 2)}\n`;
+    } catch (error) {
+        throw new Error(`Export data could not be serialized as JSON: ${errorMessage(error)}`);
+    }
+
+    let handle;
+    let created = false;
+    try {
+        handle = await open(destination, "wx", 0o600);
+        created = true;
+        await handle.writeFile(serialized, "utf8");
+        await handle.sync();
+        await handle.close();
+        handle = undefined;
+        return Buffer.byteLength(serialized);
+    } catch (error) {
+        await handle?.close().catch(() => {});
+        if (created) await unlink(destination).catch(() => {});
+        if (error?.code === "EEXIST") throw new Error(`Export destination already exists: ${destination}`);
+        throw error;
+    }
+}
+
+// ─── Lossless session core: canonical JSONL + checksum ────────────────────────
+
+/**
+ * Rebuild the session JSONL exactly as pi persists it: header line, then each
+ * entry as one compact JSON line, in insertion order. Order matters — the last
+ * line is the active leaf, which is the session's branch position.
+ *
+ * This is the single serialization used by BOTH export (to compute the
+ * checksum) and revive (to rebuild the file), so the checksum is consistent by
+ * construction, not by convention.
+ */
+export function buildSessionJsonl(header, entries) {
+    const lines = [JSON.stringify(header)];
+    for (const entry of entries) {
+        lines.push(JSON.stringify(entry));
+    }
+    return `${lines.join("\n")}\n`;
+}
+
+/**
+ * Structural validation of a session core. Checks every invariant pi's loader
+ * assumes: header shape, unique ids, resolvable parentIds, and (when present)
+ * that the declared leaf is the last entry.
+ */
+export function validateSessionCore(header, entries, { expectedChecksum, expectedLastEntryId } = {}) {
+    const errors = [];
+    if (!header || typeof header !== "object" || header.type !== "session") {
+        errors.push("Session header is missing or malformed.");
+    }
+    if (!Array.isArray(entries)) {
+        errors.push("Session entries are missing or not an array.");
+        return { ok: false, errors };
+    }
+
+    const ids = new Set();
+    for (const [index, entry] of entries.entries()) {
+        if (!entry || typeof entry !== "object" || typeof entry.type !== "string") {
+            errors.push(`Entry at index ${index} is not an object with a type.`);
+            continue;
+        }
+        if (typeof entry.id !== "string" || entry.id.length === 0) {
+            errors.push(`Entry at index ${index} is missing a string id.`);
+            continue;
+        }
+        if (ids.has(entry.id)) errors.push(`Duplicate entry id: ${entry.id}`);
+        ids.add(entry.id);
+    }
+
+    for (const [index, entry] of entries.entries()) {
+        if (!entry || typeof entry !== "object" || typeof entry.id !== "string") continue;
+        const { parentId } = entry;
+        if (parentId !== null && parentId !== undefined) {
+            if (typeof parentId !== "string" || !ids.has(parentId)) {
+                errors.push(`Entry ${entry.id} (index ${index}) references unknown parentId: ${String(parentId)}`);
+            }
+        }
+    }
+
+    const lastEntryId = entries.length > 0 ? entries[entries.length - 1]?.id : null;
+    if (typeof expectedLastEntryId === "string" && expectedLastEntryId !== lastEntryId) {
+        errors.push(
+            `Leaf mismatch: export declares leaf "${expectedLastEntryId}" but the last entry is "${String(lastEntryId)}".`,
+        );
+    }
+
+    if (typeof expectedChecksum === "string") {
+        const actual = sha256Hex(buildSessionJsonl(header, entries));
+        if (actual !== expectedChecksum) {
+            errors.push(
+                `Checksum mismatch: export declares ${expectedChecksum.slice(0, 16)}… but the entries serialize to ${actual.slice(0, 16)}… — the file is corrupt or was edited.`,
+            );
+        }
+    }
+
+    return { ok: errors.length === 0, errors, lastEntryId };
+}
+
+/**
+ * Full validation of an export document against this extension's expectations.
+ * `maxHeaderVersion` is pi's CURRENT_SESSION_VERSION on the reviving machine;
+ * an export made by a NEWER pi (higher session header version) is refused rather
+ * than half-loaded.
+ */
+export function validateExportDocument(doc, { maxHeaderVersion }) {
+    const errors = [];
+    if (!doc || typeof doc !== "object") {
+        return { ok: false, errors: ["File does not contain a JSON object."] };
+    }
+    if (doc.format !== EXPORT_FORMAT) {
+        errors.push(`Not a ${EXPORT_FORMAT} document (found format: ${JSON.stringify(doc.format)}).`);
+        return { ok: false, errors };
+    }
+    if (typeof doc.schemaVersion !== "number" || doc.schemaVersion > SCHEMA_VERSION) {
+        errors.push(
+            `Export schema version ${JSON.stringify(doc.schemaVersion)} is newer than this extension understands (${SCHEMA_VERSION}). Upgrade the extension (and pi), or re-export from the source machine.`,
+        );
+        return { ok: false, errors };
+    }
+
+    const session = doc.session;
+    if (!session || typeof session !== "object" || !Array.isArray(session.entries)) {
+        errors.push("Document has no session.entries array.");
+        return { ok: false, errors };
+    }
+
+    const header = session.header;
+    const headerVersion = typeof header?.version === "number" ? header.version : undefined;
+    if (headerVersion === undefined) {
+        errors.push("Session header has no numeric version.");
+    } else if (typeof maxHeaderVersion === "number" && headerVersion > maxHeaderVersion) {
+        errors.push(
+            `Session header version ${headerVersion} is newer than this pi supports (max ${maxHeaderVersion}). Upgrade pi on this machine.`,
+        );
+    }
+
+    const revive = session.revive;
+    if (typeof revive?.entryCount === "number" && revive.entryCount !== session.entries.length) {
+        errors.push(
+            `Entry count mismatch: export declares ${revive.entryCount} but the entries array has ${session.entries.length}.`,
+        );
+    }
+    const core = validateSessionCore(header, session.entries, {
+        expectedChecksum: typeof revive?.checksum === "string" ? revive.checksum : undefined,
+        expectedLastEntryId: typeof revive?.lastEntryId === "string" ? revive.lastEntryId : undefined,
+    });
+    errors.push(...core.errors);
+
+    return { ok: errors.length === 0, errors };
+}
+
+/**
+ * Build the revival plan: a rebuilt .jsonl (header re-rooted to the target cwd,
+ * original session UUID kept unless it already exists on this machine, optional
+ * session name re-attached as a session_info entry when the exported entries
+ * lost it) plus the destination path in pi's session directory.
+ *
+ * Pure: `uuidExists(uuid) -> boolean`, `newUUID() -> string`, `now() -> Date`,
+ * and `newEntryId() -> string` are injected so tests can pin them.
+ */
+export function buildRevivePlan(doc, { cwd, sessionDir, uuidExists, newUUID, now = new Date(), newEntryId = makeEntryId }) {
+    const header = { ...doc.session.header, cwd };
+    const entries = [...doc.session.entries];
+
+    // Identity continuity across machines; fresh id at home where the original
+    // session file still exists (avoids ambiguous `pi --session <id>` matches).
+    const originalId = typeof doc.session.header?.id === "string" ? doc.session.header.id : undefined;
+    let sessionUUID = originalId && !uuidExists(originalId) ? originalId : newUUID();
+
+    // Name fallback: the name normally round-trips as a session_info entry in
+    // the tree; if it somehow didn't, re-attach it the same way pi itself does
+    // for /name. This runs AFTER checksum verification (see README: it is a
+    // declared revival transformation, not part of the lossless core).
+    let appendedName = null;
+    const hasNameEntry = entries.some((e) => e?.type === "session_info");
+    const name = typeof doc.session.name === "string" && doc.session.name.length > 0 ? doc.session.name : null;
+    if (name && !hasNameEntry) {
+        const lastEntryId = entries.length > 0 ? entries[entries.length - 1].id : null;
+        entries.push({
+            type: "session_info",
+            id: newEntryId(),
+            parentId: lastEntryId,
+            timestamp: now.toISOString(),
+            name,
+        });
+        appendedName = name;
+    }
+
+    const jsonl = buildSessionJsonl(header, entries);
+    const filePath = join(sessionDir, `${sessionFileTimestamp(now)}_${sessionUUID}.jsonl`);
+
+    return {
+        filePath,
+        header,
+        sessionUUID,
+        appendedName,
+        entryCount: entries.length,
+        jsonl,
+    };
+}