pi-trace-viewer 0.1.0

package/src/index.ts

@@ -0,0 +1,145 @@
+import type { ExtensionAPI, ExtensionContext, ToolInfo } from "@earendil-works/pi-coding-agent";
+import { TraceCollector } from "./collector.ts";
+import { closeViewerController, getViewerController, type ViewerController } from "./server.ts";
+import type { SessionSnapshot } from "./types.ts";
+
+interface BoundSession {
+	id: string;
+	collector: TraceCollector;
+	getSnapshot: () => SessionSnapshot;
+}
+
+export default function piTraceViewer(pi: ExtensionAPI): void {
+	pi.registerFlag("pi-trace-port", {
+		description: "Starting local port for the pi trace viewer (increments automatically if occupied)",
+		type: "string",
+		default: "7890",
+	});
+
+	let controller: ViewerController | undefined;
+	let bound: BoundSession | undefined;
+	let lastSystemPrompt = "";
+	let latestTools: ToolInfo[] = [];
+
+	pi.registerCommand("trace-view", {
+		description: "Show the local session and LLM trace viewer URL",
+		handler: async (_args, ctx) => {
+			if (!controller) {
+				ctx.ui.notify("Trace viewer is not running. Check the startup error above.", "error");
+				return;
+			}
+			const suffix = bound ? `?session=${encodeURIComponent(bound.id)}` : "";
+			ctx.ui.notify(`${controller.url}/${suffix}`, "info");
+		},
+	});
+
+	pi.on("session_start", async (_event, ctx) => {
+		const port = parsePort(pi.getFlag("pi-trace-port"));
+		try {
+			controller = await getViewerController(port);
+		} catch (error) {
+			ctx.ui.notify(
+				`Trace viewer could not bind 127.0.0.1 (starting at port ${port}): ${error instanceof Error ? error.message : String(error)}. Trace capture is disabled for this session.`,
+				"error",
+			);
+			return;
+		}
+
+		latestTools = activeTools(pi);
+		const getSnapshot = (): SessionSnapshot => snapshotFromContext(ctx, latestTools, lastSystemPrompt, true);
+		const initial = getSnapshot();
+		const store = controller.register({
+			id: initial.id,
+			cwd: initial.cwd,
+			file: initial.file,
+			getSnapshot,
+			snapshot: initial,
+		});
+		const collector = new TraceCollector(store, getSnapshot);
+		collector.setTools(latestTools);
+		collector.setSystemPrompt(lastSystemPrompt || ctx.getSystemPrompt());
+		bound = { id: initial.id, collector, getSnapshot };
+		const persistence = store.getPersistence();
+		if (persistence.status === "memory_only") {
+			ctx.ui.notify(`Trace viewer is running in memory only: ${persistence.error ?? "trace directory is unavailable"}`, "warning");
+		}
+		ctx.ui.notify(`Trace viewer: ${controller.url}/?session=${encodeURIComponent(initial.id)}`, "info");
+	});
+
+	pi.on("before_agent_start", (event) => {
+		lastSystemPrompt = event.systemPrompt;
+		latestTools = activeTools(pi);
+		bound?.collector.setSystemPrompt(event.systemPrompt);
+		bound?.collector.setTools(latestTools);
+	});
+
+	pi.on("turn_start", (event) => bound?.collector.startTurn(event.turnIndex));
+	pi.on("context", (event, ctx) => bound?.collector.onContext(event, ctx));
+	pi.on("before_provider_request", (event, ctx) => bound?.collector.onProviderRequest(event, ctx));
+	pi.on("after_provider_response", (event, ctx) => bound?.collector.onProviderResponse(event, ctx));
+	pi.on("message_update", (event) => bound?.collector.onMessageUpdate(event));
+	pi.on("message_end", (event) => bound?.collector.onMessageEnd(event));
+	pi.on("session_before_compact", (event, ctx) => bound?.collector.beginCompaction(event, ctx));
+	pi.on("session_compact", (event) => bound?.collector.onCompaction(event));
+	pi.on("session_compact_failed", (event) =>
+		bound?.collector.onCompactionFailed(event.errorMessage ?? (event.aborted ? "Compaction aborted" : "Compaction failed")),
+	);
+	pi.on("session_before_tree", () => bound?.collector.prepare("branch_summary"));
+	pi.on("session_tree", (event) => bound?.collector.onTree(event));
+
+	const notifySessionUpdated = () => {
+		if (bound) controller?.notify(bound.id, "session-updated");
+	};
+	pi.on("session_info_changed", notifySessionUpdated);
+	pi.on("model_select", notifySessionUpdated);
+	pi.on("thinking_level_select", notifySessionUpdated);
+	pi.on("message_end", notifySessionUpdated);
+	pi.on("session_compact", notifySessionUpdated);
+	pi.on("session_tree", notifySessionUpdated);
+
+	pi.on("session_shutdown", async (event) => {
+		if (bound && controller) {
+			const finalSnapshot = { ...bound.getSnapshot(), active: false, updatedAt: new Date().toISOString() };
+			controller.detach(bound.id, finalSnapshot);
+		}
+		bound = undefined;
+		lastSystemPrompt = "";
+		latestTools = [];
+		if (event.reason === "quit") {
+			await closeViewerController();
+			controller = undefined;
+		}
+	});
+}
+
+function parsePort(value: boolean | string | undefined): number {
+	const port = typeof value === "string" ? Number.parseInt(value, 10) : 7890;
+	return Number.isInteger(port) && port >= 1 && port <= 65535 ? port : 7890;
+}
+
+function activeTools(pi: ExtensionAPI): ToolInfo[] {
+	const active = new Set(pi.getActiveTools());
+	return pi.getAllTools().filter((tool) => active.has(tool.name));
+}
+
+function snapshotFromContext(
+	ctx: ExtensionContext,
+	tools: ToolInfo[],
+	systemPrompt: string,
+	active: boolean,
+): SessionSnapshot {
+	const manager = ctx.sessionManager;
+	return {
+		id: manager.getSessionId(),
+		name: manager.getSessionName(),
+		cwd: manager.getCwd(),
+		file: manager.getSessionFile(),
+		header: structuredClone(manager.getHeader()),
+		entries: structuredClone(manager.getEntries()),
+		leafId: manager.getLeafId(),
+		systemPrompt: systemPrompt || ctx.getSystemPrompt(),
+		tools: structuredClone(tools),
+		active,
+		updatedAt: new Date().toISOString(),
+	};
+}

package/src/security.ts

@@ -0,0 +1,20 @@
+const SENSITIVE_KEY = /^(authorization|proxy-authorization|cookie|set-cookie|x-api-key|api[-_]?key|access[-_]?token|refresh[-_]?token|client[-_]?secret|password)$/i;
+
+export function redactSensitive(value: unknown, seen = new WeakSet<object>()): unknown {
+	if (Array.isArray(value)) return value.map((item) => redactSensitive(item, seen));
+	if (!value || typeof value !== "object") return value;
+	if (seen.has(value)) return "[Circular]";
+	seen.add(value);
+
+	const result: Record<string, unknown> = {};
+	for (const [key, item] of Object.entries(value)) {
+		result[key] = SENSITIVE_KEY.test(key) ? "[REDACTED]" : redactSensitive(item, seen);
+	}
+	return result;
+}
+
+export function redactHeaders(headers: Record<string, string>): Record<string, string> {
+	return Object.fromEntries(
+		Object.entries(headers).map(([key, value]) => [key, SENSITIVE_KEY.test(key) ? "[REDACTED]" : value]),
+	);
+}

package/src/server.ts

@@ -0,0 +1,244 @@
+import { createServer, type Server, type ServerResponse } from "node:http";
+import { createRequire } from "node:module";
+import { dirname, join } from "node:path";
+import { fileURLToPath } from "node:url";
+import { readFileSync } from "node:fs";
+import type { SessionRegistration, SessionSnapshot, TraceRecord } from "./types.ts";
+import { TraceStore } from "./store.ts";
+
+const require = createRequire(import.meta.url);
+const webRoot = fileURLToPath(new URL("../web", import.meta.url));
+const markedPath = join(dirname(require.resolve("marked")), "marked.umd.js");
+const highlightPath = require.resolve("@highlightjs/cdn-assets/highlight.min.js");
+
+interface RegisteredSession {
+	registration: SessionRegistration;
+	store: TraceStore;
+	unlisten: () => void;
+}
+
+export interface ViewerController {
+	readonly port: number;
+	readonly url: string;
+	register(registration: SessionRegistration): TraceStore;
+	detach(sessionId: string, snapshot: SessionSnapshot): void;
+	notify(sessionId: string, event: string, payload?: unknown): void;
+	close(): Promise<void>;
+}
+
+class LocalViewerController implements ViewerController {
+	readonly port: number;
+	readonly url: string;
+	private server: Server;
+	private sessions = new Map<string, RegisteredSession>();
+	private clients = new Set<ServerResponse>();
+
+	private constructor(server: Server, port: number) {
+		this.server = server;
+		this.port = port;
+		this.url = `http://127.0.0.1:${port}`;
+	}
+
+	static async start(startPort: number): Promise<LocalViewerController> {
+		for (let port = startPort; port <= 65535; port++) {
+			try {
+				return await new Promise<LocalViewerController>((resolve, reject) => {
+					const server = createServer();
+					const controller = new LocalViewerController(server, port);
+					server.on("request", (request, response) => controller.handle(request.method ?? "GET", request.url ?? "/", response));
+					const onError = (err: unknown) => {
+						server.close();
+						reject(err);
+					};
+					server.once("error", onError);
+					server.listen(port, "127.0.0.1", () => {
+						server.off("error", onError);
+						resolve(controller);
+					});
+				});
+			} catch (error: unknown) {
+				const isAddrInUse =
+					typeof error === "object" && error !== null && "code" in error && (error as { code?: unknown }).code === "EADDRINUSE";
+				if (isAddrInUse) {
+					continue;
+				}
+				throw error;
+			}
+		}
+		throw new Error(`Could not find an available port from ${startPort} to 65535`);
+	}
+
+	register(registration: SessionRegistration): TraceStore {
+		const previous = this.sessions.get(registration.id);
+		if (previous) {
+			previous.registration = registration;
+			this.notify(registration.id, "session-updated");
+			return previous.store;
+		}
+		const store = new TraceStore(registration.snapshot);
+		const unlisten = store.onRecord((record) => this.broadcast("trace-record", record));
+		this.sessions.set(registration.id, { registration, store, unlisten });
+		this.notify(registration.id, "session-added");
+		return store;
+	}
+
+	detach(sessionId: string, snapshot: SessionSnapshot): void {
+		const session = this.sessions.get(sessionId);
+		if (!session) return;
+		session.registration = { ...session.registration, getSnapshot: undefined, snapshot: { ...snapshot, active: false } };
+		this.notify(sessionId, "session-detached");
+	}
+
+	notify(sessionId: string, event: string, payload?: unknown): void {
+		this.broadcast(event, { sessionId, payload });
+	}
+
+	async close(): Promise<void> {
+		for (const client of this.clients) client.end();
+		this.clients.clear();
+		for (const session of this.sessions.values()) session.unlisten();
+		await new Promise<void>((resolve) => this.server.close(() => resolve()));
+	}
+
+	private snapshot(session: RegisteredSession): SessionSnapshot {
+		if (session.registration.getSnapshot) {
+			try {
+				const snapshot = session.registration.getSnapshot();
+				session.registration.snapshot = snapshot;
+				return snapshot;
+			} catch {
+				// Keep the last stable snapshot while a session is switching.
+			}
+		}
+		return session.registration.snapshot;
+	}
+
+	private handle(method: string, rawUrl: string, response: ServerResponse): void {
+		const url = new URL(rawUrl, this.url);
+		if (method === "GET" && url.pathname === "/api/events") {
+			response.writeHead(200, {
+				"Content-Type": "text/event-stream",
+				"Cache-Control": "no-cache",
+				Connection: "keep-alive",
+				"X-Content-Type-Options": "nosniff",
+			});
+			response.write("event: ready\ndata: {}\n\n");
+			this.clients.add(response);
+			response.on("close", () => this.clients.delete(response));
+			return;
+		}
+
+		if (method === "GET" && url.pathname === "/api/sessions") {
+			const sessions = Array.from(this.sessions.values()).map((session) => {
+				const snapshot = this.snapshot(session);
+				return {
+					id: snapshot.id,
+					name: snapshot.name,
+					cwd: snapshot.cwd,
+					active: snapshot.active,
+					updatedAt: snapshot.updatedAt,
+					callCount: session.store.getCalls().length,
+					tracePersistence: session.store.getPersistence(),
+				};
+			});
+			this.json(response, 200, sessions);
+			return;
+		}
+
+		const match = url.pathname.match(/^\/api\/sessions\/([^/]+)(?:\/(calls)(?:\/([^/]+))?|\/(download))?$/);
+		if (method === "GET" && match) {
+			const sessionId = decodeURIComponent(match[1]);
+			const session = this.sessions.get(sessionId);
+			if (!session) return this.json(response, 404, { error: "Session not found" });
+			if (match[4] === "download") {
+				const snapshot = this.snapshot(session);
+				response.writeHead(200, {
+					"Content-Type": "application/x-ndjson; charset=utf-8",
+					"Content-Disposition": `attachment; filename="${sessionId}.jsonl"`,
+				});
+				response.end([JSON.stringify(snapshot.header), ...snapshot.entries.map((entry) => JSON.stringify(entry))].join("\n"));
+				return;
+			}
+			if (match[2] === "calls") {
+				const calls = session.store.getCalls();
+				if (match[3]) {
+					const call = calls.find((candidate) => candidate.callId === decodeURIComponent(match[3]));
+					return this.json(response, call ? 200 : 404, call ?? { error: "Call not found" });
+				}
+				return this.json(response, 200, calls);
+			}
+			return this.json(response, 200, { ...this.snapshot(session), tracePersistence: session.store.getPersistence() });
+		}
+
+		if (method === "POST" && url.pathname === "/api/refresh") {
+			this.broadcast("refresh", {});
+			return this.json(response, 200, { ok: true });
+		}
+
+		if (method === "GET" && url.pathname === "/vendor/marked.js") return this.file(response, markedPath, "text/javascript");
+		if (method === "GET" && url.pathname === "/vendor/highlight.js") return this.file(response, highlightPath, "text/javascript");
+		if (method === "GET" && url.pathname === "/favicon.ico") {
+			response.writeHead(204, { "Cache-Control": "public, max-age=86400" });
+			response.end();
+			return;
+		}
+		if (method === "GET" && (url.pathname === "/" || url.pathname === "/index.html")) return this.file(response, join(webRoot, "index.html"), "text/html");
+		if (method === "GET" && url.pathname === "/app.js") return this.file(response, join(webRoot, "app.js"), "text/javascript");
+		if (method === "GET" && url.pathname === "/render-helpers.js") return this.file(response, join(webRoot, "render-helpers.js"), "text/javascript");
+		if (method === "GET" && url.pathname === "/styles.css") return this.file(response, join(webRoot, "styles.css"), "text/css");
+		this.json(response, 404, { error: "Not found" });
+	}
+
+	private file(response: ServerResponse, path: string, contentType: string): void {
+		try {
+			const content = readFileSync(path);
+			response.writeHead(200, {
+				"Content-Type": `${contentType}; charset=utf-8`,
+				"Cache-Control": "no-cache",
+				"Content-Security-Policy": "default-src 'self'; img-src 'self' data:; style-src 'self' 'unsafe-inline'; script-src 'self'; connect-src 'self'",
+				"X-Content-Type-Options": "nosniff",
+				"Referrer-Policy": "no-referrer",
+			});
+			response.end(content);
+		} catch (error) {
+			this.json(response, 500, { error: error instanceof Error ? error.message : String(error) });
+		}
+	}
+
+	private json(response: ServerResponse, status: number, value: unknown): void {
+		response.writeHead(status, {
+			"Content-Type": "application/json; charset=utf-8",
+			"Cache-Control": "no-store",
+			"X-Content-Type-Options": "nosniff",
+		});
+		response.end(JSON.stringify(value));
+	}
+
+	private broadcast(event: string, data: unknown): void {
+		const chunk = `event: ${event}\ndata: ${JSON.stringify(data)}\n\n`;
+		for (const client of this.clients) client.write(chunk);
+	}
+}
+
+const CONTROLLER_KEY = Symbol.for("pi-trace-viewer.controller");
+interface ViewerGlobal {
+	[CONTROLLER_KEY]?: Promise<ViewerController>;
+}
+
+export function getViewerController(port: number): Promise<ViewerController> {
+	const globalState = globalThis as ViewerGlobal;
+	globalState[CONTROLLER_KEY] ??= LocalViewerController.start(port).catch((error) => {
+		delete globalState[CONTROLLER_KEY];
+		throw error;
+	});
+	return globalState[CONTROLLER_KEY];
+}
+
+export async function closeViewerController(): Promise<void> {
+	const globalState = globalThis as ViewerGlobal;
+	const controller = globalState[CONTROLLER_KEY];
+	if (!controller) return;
+	delete globalState[CONTROLLER_KEY];
+	const resolved = await controller.catch(() => undefined);
+	await resolved?.close();
+}

package/src/store.ts

@@ -0,0 +1,200 @@
+import { appendFileSync, chmodSync, existsSync, mkdirSync, readFileSync, statSync } from "node:fs";
+import { dirname, join } from "node:path";
+import type { CallView, SessionSnapshot, TracePersistence, TraceRecord } from "./types.ts";
+import { TRACE_SCHEMA_VERSION } from "./types.ts";
+
+export type TraceListener = (record: TraceRecord) => void;
+type TraceRecordInput = TraceRecord extends infer RecordType
+	? RecordType extends TraceRecord
+		? Omit<RecordType, "schemaVersion" | "sessionId" | "sequence" | "timestamp"> & { timestamp?: string }
+		: never
+	: never;
+
+export class TraceStore {
+	readonly sessionId: string;
+	filePath: string | undefined;
+	private records: TraceRecord[] = [];
+	private sequence = 0;
+	private listeners = new Set<TraceListener>();
+	private persistence: TracePersistence;
+
+	constructor(snapshot: SessionSnapshot) {
+		this.sessionId = snapshot.id;
+		const filePath = join(snapshot.cwd, ".pi-traces", `${snapshot.id}.jsonl`);
+		const error = validateTraceDirectory(snapshot.cwd);
+		this.filePath = error ? undefined : filePath;
+		this.persistence = error ? { status: "memory_only", error } : { status: "persisted", filePath };
+		this.load();
+		if (!this.records.some((record) => record.type === "trace_header")) {
+			this.append({ type: "trace_header", sessionFile: snapshot.file, cwd: snapshot.cwd });
+		}
+		this.reconcileHistoricalCompactions(snapshot);
+	}
+
+	onRecord(listener: TraceListener): () => void {
+		this.listeners.add(listener);
+		return () => this.listeners.delete(listener);
+	}
+
+	append(record: TraceRecordInput): TraceRecord {
+		const complete = {
+			...record,
+			schemaVersion: TRACE_SCHEMA_VERSION,
+			sessionId: this.sessionId,
+			sequence: ++this.sequence,
+			timestamp: record.timestamp ?? new Date().toISOString(),
+		} as TraceRecord;
+		this.records.push(complete);
+		if (this.filePath) {
+			try {
+				mkdirSync(dirname(this.filePath), { recursive: true, mode: 0o700 });
+				appendFileSync(this.filePath, `${JSON.stringify(complete)}\n`, { encoding: "utf8", mode: 0o600 });
+				chmodSync(this.filePath, 0o600);
+			} catch (error) {
+				this.persistence = {
+					status: "memory_only",
+					error: error instanceof Error ? error.message : String(error),
+				};
+				this.filePath = undefined;
+			}
+		}
+		for (const listener of this.listeners) listener(complete);
+		return complete;
+	}
+
+	getPersistence(): TracePersistence {
+		return { ...this.persistence };
+	}
+
+	getRecords(): readonly TraceRecord[] {
+		return this.records;
+	}
+
+	getCalls(): CallView[] {
+		const calls = new Map<string, CallView>();
+		for (const record of this.records) {
+			if (!("callId" in record)) continue;
+			if (record.type === "call_started") {
+				calls.set(record.callId, {
+					callId: record.callId,
+					kind: record.kind,
+					startedAt: record.timestamp,
+					turnIndex: record.turnIndex,
+					leafId: record.leafId,
+					model: record.model,
+					captureSource: record.captureSource ?? "live",
+					sourceEntryId: record.sourceEntryId,
+					status: "running",
+					providerRequests: [],
+					providerResponses: [],
+					outputEvents: [],
+				});
+				continue;
+			}
+			const call = calls.get(record.callId);
+			if (!call) continue;
+			switch (record.type) {
+				case "generic_context":
+					call.context = record.context;
+					break;
+				case "compaction_context":
+					call.compactionContext = record.context;
+					break;
+				case "provider_request":
+					call.providerRequests.push({ attempt: record.attempt, timestamp: record.timestamp, payload: record.payload });
+					break;
+				case "provider_response":
+					call.providerResponses.push({
+						attempt: record.attempt,
+						timestamp: record.timestamp,
+						status: record.status,
+						headers: record.headers,
+					});
+					break;
+				case "output_event":
+					call.outputEvents.push({ timestamp: record.timestamp, event: record.event });
+					break;
+				case "call_completed":
+					call.status = "success";
+					call.completedAt = record.timestamp;
+					call.finalMessage = record.message;
+					call.leafId = record.leafId ?? call.leafId;
+					break;
+				case "call_failed":
+					call.status = "error";
+					call.completedAt = record.timestamp;
+					call.error = record.error;
+					break;
+				case "compaction_completed":
+					call.status = "success";
+					call.completedAt = record.timestamp;
+					call.finalMessage = { role: "assistant", content: [{ type: "text", text: record.summary }] };
+					break;
+				case "branch_summary_completed":
+					call.status = "success";
+					call.completedAt = record.timestamp;
+					call.leafId = record.leafId ?? call.leafId;
+					call.finalMessage = { role: "assistant", content: [{ type: "text", text: record.summary }] };
+					break;
+			}
+		}
+		return Array.from(calls.values()).sort((a, b) => b.startedAt.localeCompare(a.startedAt));
+	}
+
+	private load(): void {
+		if (!this.filePath || !existsSync(this.filePath)) return;
+		const loaded: TraceRecord[] = [];
+		for (const line of readFileSync(this.filePath, "utf8").split("\n")) {
+			if (!line.trim()) continue;
+			try {
+				const record = JSON.parse(line) as TraceRecord;
+				if (record.schemaVersion === TRACE_SCHEMA_VERSION && record.sessionId === this.sessionId) loaded.push(record);
+			} catch {
+				// A truncated final line is expected after an abrupt process exit.
+			}
+		}
+		this.records = loaded;
+		this.sequence = loaded.reduce((max, record) => Math.max(max, record.sequence), 0);
+	}
+
+	private reconcileHistoricalCompactions(snapshot: SessionSnapshot): void {
+		const completed = this.records.filter((record) => record.type === "compaction_completed");
+		for (const entry of snapshot.entries) {
+			if (entry.type !== "compaction") continue;
+			const alreadyCaptured = completed.some(
+				(record) =>
+					record.sourceEntryId === entry.id ||
+					(record.summary === entry.summary && record.tokensBefore === entry.tokensBefore),
+			);
+			if (alreadyCaptured) continue;
+			const callId = `session-compaction-${entry.id}`;
+			this.append({
+				type: "call_started",
+				callId,
+				kind: "compaction",
+				leafId: entry.parentId,
+				captureSource: "session_entry",
+				sourceEntryId: entry.id,
+				timestamp: entry.timestamp,
+			});
+			this.append({
+				type: "compaction_completed",
+				callId,
+				summary: entry.summary,
+				tokensBefore: entry.tokensBefore,
+				sourceEntryId: entry.id,
+				timestamp: entry.timestamp,
+			});
+		}
+	}
+}
+
+function validateTraceDirectory(cwd: string): string | undefined {
+	try {
+		const stat = statSync(cwd);
+		if (!stat.isDirectory()) return `Trace cwd is not a directory: ${cwd}`;
+	} catch (error) {
+		return error instanceof Error ? error.message : String(error);
+	}
+	return undefined;
+}