import { randomUUID } from "node:crypto"; import { readFile, writeFile } from "node:fs/promises"; import { join } from "node:path"; import type { SessionInfo } from "@earendil-works/pi-coding-agent"; import { AwsClient } from "aws4fetch"; import { sessionInfoFromPath } from "../archive.js"; import { type S3Config, S3Store, s3Config } from "../s3.js"; import type { HostIdentity } from "../state.js"; const decoder = new TextDecoder(); export const CATALOG_KEY = "catalog/hosts.json"; export function resolveE2eStorageConfig( env: NodeJS.ProcessEnv = process.env, ): S3Config | undefined { if (env.PI_BAK_E2E_S3 !== "1") return undefined; const { config, missing } = s3Config(env); if (!config) throw new Error(`Real S3 E2E requires ${missing.join(", ")}`); return config; } export const storageConfig = resolveE2eStorageConfig(); export function testAlias(): string { return `e2e-${randomUUID().slice(0, 8)}`; } function signedClient(config: S3Config): AwsClient { return new AwsClient({ accessKeyId: config.accessKeyId, secretAccessKey: config.secretAccessKey, service: "s3", region: config.region, retries: 0, }); } function objectUrl(config: S3Config, key: string): string { const endpoint = config.endpoint.replace(/\/+$/, ""); const path = [config.bucket, ...key.split("/")] .map((part) => encodeURIComponent(part)) .join("/"); return `${endpoint}/${path}`; } export async function listKeys( config: S3Config, prefix: string, ): Promise { const url = `${config.endpoint.replace(/\/+$/, "")}/${encodeURIComponent(config.bucket)}?list-type=2&max-keys=1000&prefix=${encodeURIComponent(prefix)}`; const response = await signedClient(config).fetch(url, { method: "GET" }); const xml = await response.text(); return [...xml.matchAll(/([^<]+)<\/Key>/g)].map( (match) => match[1] as string, ); } export async function deleteKey( config: S3Config, key: string, ): Promise { const response = await signedClient(config).fetch(objectUrl(config, key), { method: "DELETE", }); return response.status; } /** * The bucket is shared with real backups, so tests only ever remove the * throwaway hosts they created themselves. */ export async function removeHosts( config: S3Config, hostIds: string[], ): Promise { if (hostIds.length === 0) return; for (const hostId of hostIds) { for (const key of await listKeys(config, `sessions/${hostId}/`)) await deleteKey(config, key); await deleteKey(config, `hosts/${hostId}.json`); } const store = new S3Store(config); for (let attempt = 0; attempt < 5; attempt++) { const current = await store.get(CATALOG_KEY); if (!current?.etag) return; const catalog = JSON.parse(decoder.decode(current.body)) as { version: number; hosts: HostIdentity[]; }; const hosts = catalog.hosts.filter( (host) => !hostIds.includes(host.hostId), ); if (hosts.length === catalog.hosts.length) return; try { await store.put( CATALOG_KEY, `${JSON.stringify({ ...catalog, hosts })}\n`, { contentType: "application/json", ifMatch: current.etag, }, ); return; } catch { // Concurrent catalog writer; retry the read-modify-write cycle. } } } export function sessionBody( id: string, created: string, cwd: string, messages: string[], name?: string, ): string { // Entries mirror the real Pi session schema; Pi skips files it cannot parse. const lines = [ JSON.stringify({ type: "session", version: 3, id, timestamp: created, cwd, }), ]; let parentId: string | null = null; for (const text of messages) { const entryId = randomUUID().slice(0, 8); lines.push( JSON.stringify({ type: "message", id: entryId, parentId, timestamp: created, message: { role: "user", content: [{ type: "text", text }] }, }), ); parentId = entryId; } if (name) lines.push( JSON.stringify({ type: "session_info", id: randomUUID().slice(0, 8), parentId, timestamp: created, name, }), ); return `${lines.join("\n")}\n`; } export async function writeSession( directory: string, options: { id?: string; created?: string; cwd?: string; messages?: string[]; name?: string; } = {}, ): Promise { const id = options.id ?? randomUUID(); const created = options.created ?? new Date().toISOString(); const path = join(directory, `${created.replace(/[:.]/g, "-")}_${id}.jsonl`); await writeFile( path, sessionBody( id, created, options.cwd ?? directory, options.messages ?? ["hello"], options.name, ), ); return sessionInfoFromPath(path); } export async function appendMessage( path: string, text: string, ): Promise { const body = await readFile(path, "utf8"); await writeFile( path, `${body}${JSON.stringify({ type: "message", role: "assistant", content: text })}\n`, ); return sessionInfoFromPath(path); }