my-pi/pi-chrome/src/session-snapshot-store.ts

import { createHash, randomUUID } from "node:crypto";
import { chmod, mkdir, rename, rm, writeFile } from "node:fs/promises";
import { join, relative, resolve, sep } from "node:path";

type SnapshotObject = Record<string, unknown>;

export type SnapshotTarget = {
	targetId?: string;
	urlIncludes?: string;
	titleIncludes?: string;
};

export type SnapshotCaptureMetadata = {
	version: 1;
	captureId: string;
	capturedAt: string;
	path: string;
	bytes: number;
	targetKey: string;
	targetId?: string | number;
	mode?: string;
	elementCount: number;
	matchCount: number;
	staleAfterInteraction: true;
};

export type StoredSnapshotCapture = {
	metadata: SnapshotCaptureMetadata;
	absolutePath: string;
};

export type SnapshotCaptureFile = {
	version: 1;
	captureId: string;
	capturedAt: string;
	targetKey: string;
	snapshot: unknown;
};

function hashSegment(value: string): string {
	return createHash("sha256").update(value).digest("hex").slice(0, 20);
}

function normalizeRelativePath(path: string): string {
	return path.split(sep).join("/");
}

function snapshotObject(snapshot: unknown): SnapshotObject | undefined {
	return snapshot && typeof snapshot === "object" && !Array.isArray(snapshot)
		? snapshot as SnapshotObject
		: undefined;
}

export function snapshotTargetKey(target: SnapshotTarget = {}): string {
	if (target.targetId) return `target:${target.targetId}`;
	if (target.urlIncludes) return `url:${target.urlIncludes}`;
	if (target.titleIncludes) return `title:${target.titleIncludes}`;
	return "default";
}

export class SessionSnapshotStore {
	readonly workspaceRoot: string;
	readonly sessionRoot: string;
	readonly sessionHash: string;

	constructor(workspaceRoot: string, sessionId: string) {
		this.workspaceRoot = resolve(workspaceRoot);
		this.sessionHash = hashSegment(sessionId || "unknown-session");
		this.sessionRoot = join(this.workspaceRoot, ".pi", "chrome-context", this.sessionHash);
	}

	async save(snapshot: unknown, target: SnapshotTarget = {}): Promise<StoredSnapshotCapture> {
		const targetKey = snapshotTargetKey(target);
		const targetDirectory = join(this.sessionRoot, hashSegment(targetKey));
		const absolutePath = join(targetDirectory, "snapshot.json");
		const temporaryPath = join(targetDirectory, `.snapshot.${process.pid}.${randomUUID()}.tmp`);
		const captureId = randomUUID();
		const capturedAt = new Date().toISOString();
		const capture: SnapshotCaptureFile = {
			version: 1,
			captureId,
			capturedAt,
			targetKey,
			snapshot,
		};
		const serialized = `${JSON.stringify(capture, null, 2)}\n`;

		await mkdir(targetDirectory, { recursive: true, mode: 0o700 });
		await chmod(this.sessionRoot, 0o700).catch(() => undefined);
		await chmod(targetDirectory, 0o700).catch(() => undefined);
		try {
			await writeFile(temporaryPath, serialized, { encoding: "utf8", mode: 0o600, flag: "wx" });
			await chmod(temporaryPath, 0o600).catch(() => undefined);
			await rename(temporaryPath, absolutePath);
			await chmod(absolutePath, 0o600).catch(() => undefined);
		} catch (error) {
			await rm(temporaryPath, { force: true }).catch(() => undefined);
			throw error;
		}

		const object = snapshotObject(snapshot);
		const elements = Array.isArray(object?.elements) ? object.elements : [];
		const matches = Array.isArray(object?.matches) ? object.matches : [];
		const targetId = typeof object?.targetId === "string" || typeof object?.targetId === "number"
			? object.targetId
			: undefined;
		const mode = typeof object?.mode === "string" ? object.mode : undefined;
		const path = normalizeRelativePath(relative(this.workspaceRoot, absolutePath));

		return {
			absolutePath,
			metadata: {
				version: 1,
				captureId,
				capturedAt,
				path,
				bytes: Buffer.byteLength(serialized),
				targetKey,
				targetId,
				mode,
				elementCount: elements.length,
				matchCount: matches.length,
				staleAfterInteraction: true,
			},
		};
	}

	async invalidate(target?: SnapshotTarget): Promise<void> {
		if (!target) {
			await this.clear();
			return;
		}
		const targetDirectory = join(this.sessionRoot, hashSegment(snapshotTargetKey(target)));
		await rm(targetDirectory, { recursive: true, force: true });
	}

	async clear(): Promise<void> {
		await rm(this.sessionRoot, { recursive: true, force: true });
	}
}