lvmh/plugin/lvmh-agent.ts

/**
 * lvmh-agent — pi extension that mirrors session events to the lvmh daemon.
 *
 * Wire contract: PROTOCOL.md (envelope v1) at the repo root.
 *
 * Environment:
 *   LVMH_URL    full agent websocket URL, e.g. ws://alarm:8686/agent/ws
 *   LVMH_TOKEN  shared bearer secret
 *   LVMH_AGENT  set to "1" by the daemon when spawning container sessions
 *               (marks session.agent=true in the handshake snapshot)
 *   LVMH_REPO   "group/project" spawn metadata for container sessions
 *
 * If LVMH_URL or LVMH_TOKEN is unset the extension is fully inert.
 *
 * Bulletproofing rules implemented here (see selftest.md):
 *   - no npm dependencies: Node built-in global WebSocket (undici) + node:fs
 *   - every event handler is wrapped by sub(): sync throws are caught and
 *     returned promise rejections are swallowed; nothing ever reaches pi
 *   - all sends go through a bounded queue (drop-oldest); persisted-kind
 *     events are additionally kept in a bounded replay buffer and resent
 *     after welcome.lastSeq on reconnect
 *   - reconnect with exponential backoff + jitter, capped, forever
 *   - handshake watchdog: if welcome does not arrive within
 *     LVMH_WELCOME_TIMEOUT_MS (default 15s) of connecting, the socket is
 *     abandoned and retried (covers a daemon that accepts TCP then hangs)
 *   - stall watchdog: undici exposes no ping(), so a daemon that dies
 *     without a TCP close is detected via bufferedAmount never draining
 *     (two LVMH_STALL_CHECK_MS ticks over LVMH_STALL_BYTES, default 30s/16MB)
 *   - errors are logged to ~/.pi/lvmh-agent.log only; stdout stays clean
 *   - session_shutdown closes the socket and stops all timers immediately
 */

import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";

const PROTOCOL_VERSION: number = 1;
const ENV_URL: string = "LVMH_URL";
const ENV_TOKEN: string = "LVMH_TOKEN";
const ENV_AGENT: string = "LVMH_AGENT";
const ENV_REPO: string = "LVMH_REPO";
/** Daemon-assigned session id for spawned containers: takes precedence over
 *  pi's own session id so web prompts route to the id /api/spawn returned. */
const ENV_SESSION_ID: string = "LVMH_SESSION_ID";

const SEND_QUEUE_MAX: number = 1000;
const REPLAY_BUFFER_MAX: number = 10000;
const BACKOFF_BASE_MS: number = 1000;
const BACKOFF_MAX_MS: number = 30000;
const BACKOFF_JITTER_MS: number = 1000;
const RESULT_PREVIEW_MAX_CHARS: number = 2000;
const LOG_LINE_MAX_CHARS: number = 500;

function envPositiveInt(name: string, fallback: number): number {
	const raw: string | undefined = process.env[name];
	if (raw === undefined) return fallback;
	const n: number = Number.parseInt(raw, 10);
	return Number.isFinite(n) && n > 0 ? n : fallback;
}

// Test/debug knobs; production never sets them. Read lazily so harnesses
// can reconfigure per scenario after the module is already imported.
const welcomeTimeoutMs = (): number =>
	envPositiveInt("LVMH_WELCOME_TIMEOUT_MS", 15000);
const stallCheckMs = (): number => envPositiveInt("LVMH_STALL_CHECK_MS", 30000);
const stallBytes = (): number => envPositiveInt("LVMH_STALL_BYTES", 16777216);

const TRANSIENT_TYPES: ReadonlySet<string> = new Set(["message_update"]);

interface Frame {
	v: number;
	sessionId: string;
	seq: number;
	ts: number;
	type: string;
	[key: string]: unknown;
}

interface SessionSnapshot {
	id: string;
	name: string | null;
	cwd: string;
	model: string | null;
	provider: string | null;
	agent: boolean;
	repo: string | null;
	startedAt: number;
}

/** Per-session seq counters survive reconnects and instance rebinds in-process. */
const seqCounters: Map<string, number> = new Map();
const maxAssignedSeq: Map<string, number> = new Map();
const sessionStartTs: Map<string, number> = new Map();

const logPath: string = path.join(os.homedir(), ".pi", "lvmh-agent.log");

/** Best-effort append to the log file. Never throws, never touches stdout. */
function log(message: unknown): void {
	try {
		const line = `${new Date().toISOString()} ${String(message).slice(0, LOG_LINE_MAX_CHARS)}\n`;
		fs.appendFile(logPath, line, () => undefined);
	} catch {
		// even appendFile argument validation failures must stay silent
	}
}

function nextSeq(sessionId: string): number {
	const n: number = (seqCounters.get(sessionId) ?? 0) + 1;
	seqCounters.set(sessionId, n);
	// High-water mark of seqs THIS process produced — distinguishes a benign
	// reconnect (daemon lastSeq <= maxAssigned: frames were delivered by us)
	// from a foreign high-water after a plugin restart (real gaps).
	const seen = maxAssignedSeq.get(sessionId) ?? 0;
	if (n > seen) maxAssignedSeq.set(sessionId, n);
	return n;
}

/** Concatenate text blocks of a pi content array (or pass through a string). */
function textOfContent(content: unknown): string {
	if (typeof content === "string") return content;
	if (!Array.isArray(content)) return "";
	let out = "";
	for (const block of content) {
		if (
			block !== null &&
			typeof block === "object" &&
			(block as { type?: unknown }).type === "text"
		) {
			const text = (block as { text?: unknown }).text;
			if (typeof text === "string") out += text;
		}
	}
	return out;
}

function safeJsonStringify(value: unknown): string {
	try {
		return JSON.stringify(value) ?? "null";
	} catch {
		return "null";
	}
}

/** Map a pi AgentMessage to the protocol Message shape (message_end payload). */
function mapMessage(message: unknown): Record<string, unknown> | null {
	if (message === null || typeof message !== "object") return null;
	const msg = message as {
		role?: unknown;
		content?: unknown;
		timestamp?: unknown;
		toolCallId?: unknown;
	};
	const role = typeof msg.role === "string" ? msg.role : "unknown";
	const content = msg.content;
	const toolCalls: Array<{ id: string; name: string; argsJson: string }> = [];
	let thinking = "";
	if (Array.isArray(content)) {
		for (const block of content) {
			if (block === null || typeof block !== "object") continue;
			const b = block as {
				type?: unknown;
				id?: unknown;
				name?: unknown;
				arguments?: unknown;
			};
			if (b.type === "toolCall") {
				toolCalls.push({
					id: String(b.id ?? ""),
					name: String(b.name ?? ""),
					argsJson: safeJsonStringify(b.arguments),
				});
			} else if (
				b.type === "thinking" &&
				typeof (b as { thinking?: unknown }).thinking === "string"
			) {
				thinking += (b as { thinking: string }).thinking;
			}
		}
	}
	return {
		role,
		id: `${role}-${typeof msg.timestamp === "number" ? msg.timestamp : Date.now()}`,
		text: textOfContent(content),
		thinking: thinking.length > 0 ? thinking : null,
		toolCalls,
		toolCallId:
			role === "toolResult" && typeof msg.toolCallId === "string"
				? msg.toolCallId
				: null,
	};
}

export default function (pi: ExtensionAPI): void {
	const rawUrl: string | undefined = process.env[ENV_URL];
	const rawToken: string | undefined = process.env[ENV_TOKEN];
	if (!rawUrl || !rawToken) return; // inert: no daemon configured
	const url: string = rawUrl;
	const token: string = rawToken;
	if (typeof WebSocket === "undefined") {
		log("global WebSocket unavailable (need Node >= 22); extension inert");
		return;
	}

	let stopped: boolean = false;
	let ws: WebSocket | null = null;
	let wsOpen: boolean = false;
	let greeted: boolean = false;
	let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
	let backoffAttempt: number = 0;
	let welcomeTimer: ReturnType<typeof setTimeout> | null = null;
	let stallTimer: ReturnType<typeof setInterval> | null = null;
	let lastStallBytes: number | null = null;

	let currentSessionId: string | null = null;
	let snapshot: SessionSnapshot | null = null;
	let lastStreamText: string = "";
	let lastCtx: ExtensionContext | undefined;

	let sendQueue: Frame[] = [];
	let replayBuf: Frame[] = [];
	let droppedEvents: number = 0;

	/**
	 * Subscribe with a bulletproof wrapper: the inner handler may be sync or
	 * async; sync throws are caught here and returned promises that reject are
	 * swallowed. A failing handler logs and yields undefined — it can never
	 * throw into or reject on pi.
	 */
	function sub(
		event: string,
		handler: (event: any, ctx: ExtensionContext) => unknown,
	): void {
		(
			pi as {
				on: (
					e: string,
					h: (event: any, ctx: ExtensionContext) => unknown,
				) => void;
			}
		).on(event, (evt: any, ctx: ExtensionContext): undefined => {
			lastCtx = ctx;
			try {
				const result = handler(evt, ctx);
				if (
					result !== null &&
					typeof result === "object" &&
					typeof (result as Promise<unknown>).then === "function"
				) {
					(result as Promise<unknown>).then(undefined, (err: unknown) => {
						log(`async handler error (${event}): ${err}`);
					});
				}
			} catch (err) {
				log(`handler error (${event}): ${err}`);
			}
			return undefined;
		});
	}

	function enqueue(frame: Frame): void {
		sendQueue.push(frame);
		while (sendQueue.length > SEND_QUEUE_MAX) {
			const dropped: Frame | undefined = sendQueue.shift();
			if (dropped !== undefined) {
				droppedEvents++;
				log(`send queue overflow, dropped ${dropped.type} seq=${dropped.seq}`);
			}
		}
		flush();
	}

	function flush(): void {
		while (
			!stopped &&
			ws !== null &&
			wsOpen &&
			greeted &&
			sendQueue.length > 0
		) {
			const frame: Frame | undefined = sendQueue.shift();
			if (frame === undefined) return;
			let text: string;
			try {
				text = JSON.stringify(frame);
			} catch (err) {
				log(`serialize error (${frame.type}): ${err}`);
				continue;
			}
			try {
				ws.send(text);
			} catch (err) {
				log(`send error (${frame.type}): ${err}`);
				handleDisconnect();
				scheduleReconnect(); // in case no onclose follows the failed send
				return;
			}
		}
	}

	/** Persisted-kind events only (message_update is live-only per protocol). */
	function emit(type: string, payload: Record<string, unknown>): void {
		if (stopped || currentSessionId === null) return;
		const frame: Frame = {
			v: PROTOCOL_VERSION,
			sessionId: currentSessionId,
			seq: nextSeq(currentSessionId),
			ts: Date.now(),
			type,
			...payload,
		};
		if (!TRANSIENT_TYPES.has(type)) {
			replayBuf.push(frame);
			while (replayBuf.length > REPLAY_BUFFER_MAX) {
				replayBuf.shift();
				droppedEvents++;
				log(
					`replay buffer overflow (cap ${REPLAY_BUFFER_MAX}), dropped oldest event`,
				);
			}
		}
		enqueue(frame);
	}

	function handleDisconnect(): void {
		const socket: WebSocket | null = ws;
		ws = null;
		wsOpen = false;
		greeted = false;
		clearWelcomeTimer();
		lastStallBytes = null;
		try {
			socket?.close();
		} catch {
			// close failures are meaningless here
		}
	}

	function clearWelcomeTimer(): void {
		if (welcomeTimer === null) return;
		try {
			clearTimeout(welcomeTimer);
		} catch {
			// clearable timers never throw in practice
		}
		welcomeTimer = null;
	}

	/** A daemon that accepts TCP but never completes the handshake would hang
	 *  the socket (and all mirroring) forever: undici has no handshake timeout
	 *  of its own. Arm at connect time, clear at welcome/disconnect. */
	function armWelcomeTimeout(socket: WebSocket): void {
		clearWelcomeTimer();
		welcomeTimer = setTimeout(() => {
			welcomeTimer = null;
			if (ws !== socket || greeted) return;
			log("welcome timeout; abandoning socket");
			handleDisconnect();
			scheduleReconnect();
		}, welcomeTimeoutMs());
		(welcomeTimer as { unref?: () => void }).unref?.();
	}

	/** undici's WHATWG WebSocket has no ping(). Detect a daemon that died
	 *  without a TCP close by watching bufferedAmount: if a backlog over
	 *  STALL_BYTES makes no progress between two ticks, the peer is not
	 *  reading — kill the socket and reconnect. Draining resets the state. */
	function checkStall(): void {
		const socket: WebSocket | null = ws;
		if (socket === null || !wsOpen) {
			lastStallBytes = null;
			return;
		}
		const buffered: number = socket.bufferedAmount;
		if (!Number.isFinite(buffered) || buffered < stallBytes()) {
			lastStallBytes = null;
			return;
		}
		if (lastStallBytes !== null && buffered >= lastStallBytes) {
			log(`ws stalled: bufferedAmount=${buffered} not draining; reconnecting`);
			lastStallBytes = null;
			handleDisconnect();
			scheduleReconnect();
			return;
		}
		lastStallBytes = buffered;
	}

	function scheduleReconnect(): void {
		if (stopped || reconnectTimer !== null || ws !== null) return;
		const expMs: number = BACKOFF_BASE_MS * 2 ** backoffAttempt;
		const waitMs: number =
			Math.min(
				Number.isFinite(expMs) ? expMs : BACKOFF_MAX_MS,
				BACKOFF_MAX_MS,
			) + Math.floor(Math.random() * BACKOFF_JITTER_MS);
		backoffAttempt++;
		reconnectTimer = setTimeout(() => {
			reconnectTimer = null;
			connect();
		}, waitMs);
		(reconnectTimer as { unref?: () => void }).unref?.();
	}

	function connect(): void {
		if (stopped || ws !== null || reconnectTimer !== null) return;
		let socket: WebSocket;
		try {
			socket = new WebSocket(url, {
				headers: { Authorization: `Bearer ${token}` },
			});
		} catch (err) {
			log(`ws construct error: ${err}`);
			scheduleReconnect();
			return;
		}
		ws = socket;
		armWelcomeTimeout(socket);
		if (stallTimer === null) {
			stallTimer = setInterval(checkStall, stallCheckMs());
			(stallTimer as { unref?: () => void }).unref?.();
		}
		// Stale-socket guards: after a session switch an old socket's callbacks
		// can still fire; they must never touch the state of the new socket.
		socket.onopen = () => {
			if (ws !== socket) return;
			try {
				wsOpen = true;
				sendHello();
				flush();
			} catch (err) {
				log(`onopen error: ${err}`);
				handleDisconnect();
				scheduleReconnect();
			}
		};
		socket.onmessage = (ev: unknown) => {
			if (ws !== socket) return;
			try {
				onMessage(ev);
			} catch (err) {
				log(`onmessage error: ${err}`);
			}
		};
		socket.onerror = () => {
			// undici always follows onerror with onclose; logging happens there
		};
		socket.onclose = () => {
			if (ws !== socket) return;
			try {
				log(`ws closed after backoff attempt #${backoffAttempt}`);
				handleDisconnect();
				scheduleReconnect();
			} catch (err) {
				log(`onclose error: ${err}`);
			}
		};
	}

	function sendHello(): void {
		if (ws === null || currentSessionId === null || snapshot === null) return;
		const hello: Frame = {
			v: PROTOCOL_VERSION,
			sessionId: currentSessionId,
			seq: 0,
			ts: Date.now(),
			type: "hello",
			session: snapshot,
		};
		ws.send(JSON.stringify(hello));
	}

	function onWelcome(frame: Record<string, unknown>): void {
		if (
			typeof frame.sessionId === "string" &&
			frame.sessionId !== currentSessionId
		) {
			log(`welcome for foreign session ${String(frame.sessionId)}, ignoring`);
			return;
		}
		const lastSeq: number =
			typeof frame.lastSeq === "number" && Number.isFinite(frame.lastSeq)
				? frame.lastSeq
				: 0;
		backoffAttempt = 0;
		greeted = true;
		clearWelcomeTimer();
		// Post-restart safety: never reuse seq numbers the daemon already has.
		if (currentSessionId !== null) {
			const counter: number = seqCounters.get(currentSessionId) ?? 0;
			if (counter <= lastSeq) seqCounters.set(currentSessionId, lastSeq);
		}
		// Gap accounting: only a daemon lastSeq ABOVE anything this process ever
		// assigned indicates real gaps (plugin restart against a daemon that has
		// history). lastSeq <= maxAssigned means the daemon acked frames we sent
		// ourselves — benign reconnect, no notice. replayBuf alone holds every
		// persisted-kind frame (sendQueue copies are a subset), so counting it
		// once cannot double-count.
		const maxAssigned: number =
			currentSessionId !== null
				? (maxAssignedSeq.get(currentSessionId) ?? 0)
				: 0;
		if (lastSeq > maxAssigned) {
			const covered: number = replayBuf.filter((f) => f.seq <= lastSeq).length;
			if (covered > 0) {
				droppedEvents += covered;
				log(
					`welcome lastSeq=${lastSeq} > maxAssigned=${maxAssigned}: ${covered} replayed frames never acknowledged; flagging`,
				);
			}
		}
		// Replay persisted events the daemon is missing, before queued live
		// frames. Persisted-kind frames still sitting in the queue are dropped
		// here: they are already in replayBuf, so sending both would duplicate
		// them (only transient message_update frames survive the merge).
		const replay: Frame[] = replayBuf.filter((f) => f.seq > lastSeq);
		sendQueue = replay.concat(
			sendQueue.filter((f) => TRANSIENT_TYPES.has(f.type)),
		);
		while (sendQueue.length > SEND_QUEUE_MAX) {
			const dropped: Frame | undefined = sendQueue.shift();
			if (dropped !== undefined) {
				droppedEvents++;
				log(
					`send queue overflow during replay, dropped ${dropped.type} seq=${dropped.seq}`,
				);
			}
		}
		if (droppedEvents > 0) {
			emit("buffer_overflow", { dropped: droppedEvents });
			droppedEvents = 0;
		}
		flush();
	}

	function onMessage(ev: unknown): void {
		const data = (ev as { data?: unknown }).data;
		let frame: unknown;
		try {
			frame = JSON.parse(typeof data === "string" ? data : "");
		} catch {
			log("malformed json frame from daemon");
			return;
		}
		if (frame === null || typeof frame !== "object") return;
		const f = frame as {
			v?: unknown;
			type?: unknown;
			message?: unknown;
			sessionId?: unknown;
		};
		if (f.v !== PROTOCOL_VERSION) return;
		if (f.type === "welcome") onWelcome(f as Record<string, unknown>);
		else if (f.type === "prompt") deliverPrompt(f);
		else if (f.type === "abort") doAbort();
		else if (f.type === "set_model") applySetModel(f);
		else if (f.type === "rename") applyRename(f);
		// unknown types are ignored (forward compatibility)
	}

	function deliverPrompt(frame: {
		message?: unknown;
		sessionId?: unknown;
	}): void {
		const message = frame.message;
		if (typeof message !== "string" || message.length === 0) return;
		if (
			typeof frame.sessionId === "string" &&
			frame.sessionId !== currentSessionId
		)
			return;
		try {
			pi.sendUserMessage(message, { deliverAs: "steer" });
		} catch (err) {
			log(`sendUserMessage failed: ${err}`);
		}
	}

	function doAbort(): void {
		try {
			const maybeAbort = (pi as unknown as { abort?: () => void }).abort;
			if (typeof maybeAbort === "function") maybeAbort.call(pi);
			else lastCtx?.abort();
		} catch (err) {
			log(`abort failed: ${err}`);
		}
	}

	/** Daemon-requested model switch. Success is mirrored by the resulting
	 *  model_select event (which emits session_info); every failure path
	 *  emits error_notice so the web sees why nothing changed. */
	function applySetModel(frame: {
		modelId?: unknown;
		provider?: unknown;
		sessionId?: unknown;
	}): void {
		const modelId = frame.modelId;
		const provider = frame.provider;
		if (
			typeof frame.sessionId === "string" &&
			frame.sessionId !== currentSessionId
		)
			return;
		if (typeof modelId !== "string" || modelId.length === 0) {
			emit("error_notice", { reason: "set_model without modelId" });
			return;
		}
		if (typeof provider !== "string" || provider.length === 0) {
			emit("error_notice", { reason: "set_model without provider" });
			return;
		}
		const registry = lastCtx?.modelRegistry as
			| { find?: (provider: string, modelId: string) => unknown }
			| undefined;
		if (registry === undefined || typeof registry.find !== "function") {
			emit("error_notice", { reason: "model registry unavailable" });
			return;
		}
		let model: unknown;
		try {
			model = registry.find(provider, modelId);
		} catch (err) {
			log(`model lookup failed: ${err}`);
			emit("error_notice", {
				reason: `model lookup failed: ${String(err)}`,
			});
			return;
		}
		if (model === undefined || model === null) {
			emit("error_notice", {
				reason: `model not found: ${provider}/${modelId}`,
			});
			return;
		}
		const maybeSet = (pi as unknown as { setModel?: (m: unknown) => unknown })
			.setModel;
		if (typeof maybeSet !== "function") {
			emit("error_notice", { reason: "pi.setModel unavailable" });
			return;
		}
		try {
			Promise.resolve(maybeSet.call(pi, model)).then(
				(ok: unknown) => {
					if (ok !== true)
						emit("error_notice", {
							reason: `setModel rejected ${provider}/${modelId} (no API key?)`,
						});
				},
				(err: unknown) => {
					log(`setModel failed: ${err}`);
					emit("error_notice", {
						reason: `setModel failed: ${String(err)}`,
					});
				},
			);
		} catch (err) {
			log(`setModel threw: ${err}`);
			emit("error_notice", { reason: `setModel failed: ${String(err)}` });
		}
	}

	/** Daemon-requested rename. pi has no extension rename API: only the
	 *  mirrored snapshot changes (PROTOCOL.md); the TUI session keeps its
	 *  own name. session_info propagates the change to the daemon/web. */
	function applyRename(frame: { name?: unknown; sessionId?: unknown }): void {
		const name = frame.name;
		if (typeof name !== "string" || name.length === 0) return;
		if (
			typeof frame.sessionId === "string" &&
			frame.sessionId !== currentSessionId
		)
			return;
		if (snapshot === null) return;
		if (name === snapshot.name) return;
		snapshot.name = name;
		emit("session_info", { session: snapshot });
	}

	function buildSnapshot(
		ctx: ExtensionContext,
		sessionId: string,
	): SessionSnapshot {
		const sm = ctx.sessionManager as {
			getSessionName?: () => string | undefined;
			getCwd?: () => string;
			getHeader?: () => { timestamp?: unknown } | null;
		};
		let name: string | null = null;
		let cwd: string = ctx.cwd;
		let startedAt: number | undefined = sessionStartTs.get(sessionId);
		try {
			name = sm.getSessionName?.() ?? null;
			cwd = sm.getCwd?.() ?? ctx.cwd;
			if (startedAt === undefined) {
				const headerTs = sm.getHeader?.()?.timestamp;
				startedAt = typeof headerTs === "string" ? Date.parse(headerTs) : NaN;
				if (!Number.isFinite(startedAt)) startedAt = Date.now();
				sessionStartTs.set(sessionId, startedAt);
			}
		} catch (err) {
			log(`snapshot fallback (${String(err).slice(0, 80)}): using ctx.cwd/now`);
			if (startedAt === undefined) {
				startedAt = Date.now();
				sessionStartTs.set(sessionId, startedAt);
			}
		}
		const model = ctx.model as { id?: unknown; provider?: unknown } | undefined;
		return {
			id: sessionId,
			name,
			cwd,
			model:
				model !== undefined && typeof model.id === "string" ? model.id : null,
			provider:
				model !== undefined && typeof model.provider === "string"
					? model.provider
					: null,
			agent: process.env[ENV_AGENT] === "1",
			repo: process.env[ENV_REPO] ?? null,
			startedAt,
		};
	}

	sub("session_start", (_event: any, ctx: ExtensionContext) => {
		stopped = false;
		let sessionId: string | null = null;
		try {
			const id = (
				ctx.sessionManager as { getSessionId?: () => string }
			).getSessionId?.();
			sessionId = typeof id === "string" && id.length > 0 ? id : null;
		} catch (err) {
			log(`getSessionId failed: ${err}`);
		}
		if (sessionId === null) {
			log("no session id; lvmh mirroring disabled for this session");
			// Enforce the logged contract: stop mirroring rather than
			// attributing this session's events to the previous session id.
			currentSessionId = null;
			handleDisconnect();
			try {
				if (reconnectTimer !== null) clearTimeout(reconnectTimer);
			} catch {
				// clearable timers never throw in practice
			}
			reconnectTimer = null;
			return;
		}
		const assignedId = process.env[ENV_SESSION_ID];
		if (typeof assignedId === "string" && assignedId.length > 0) {
			// Spawned container session: report the daemon-assigned id.
			sessionId = assignedId;
		}
		if (sessionId !== currentSessionId) {
			sendQueue = [];
			replayBuf = [];
			droppedEvents = 0;
			lastStreamText = "";
			handleDisconnect(); // drop any socket bound to the previous session
			try {
				if (reconnectTimer !== null) clearTimeout(reconnectTimer);
			} catch {
				// clearable timers never throw in practice
			}
			reconnectTimer = null;
		}
		currentSessionId = sessionId;
		snapshot = buildSnapshot(ctx, sessionId);
		connect();
	});

	sub("session_shutdown", () => {
		stopped = true;
		try {
			if (reconnectTimer !== null) clearTimeout(reconnectTimer);
		} catch {
			// ignore
		}
		reconnectTimer = null;
		clearWelcomeTimer();
		try {
			if (stallTimer !== null) clearInterval(stallTimer);
		} catch {
			// ignore
		}
		stallTimer = null;
		lastStallBytes = null;
		const socket: WebSocket | null = ws;
		ws = null;
		wsOpen = false;
		greeted = false;
		try {
			if (
				socket !== null &&
				socket.readyState === WebSocket.OPEN &&
				currentSessionId !== null
			) {
				socket.send(
					JSON.stringify({
						v: PROTOCOL_VERSION,
						sessionId: currentSessionId,
						seq: 0,
						ts: Date.now(),
						type: "bye",
						reason: "shutdown",
					}),
				);
			}
			socket?.close();
		} catch {
			// best-effort goodbye; nothing to flush, exit fast
		}
	});

	sub("session_info_changed", (event: { name?: unknown }) => {
		if (snapshot === null) return;
		const name = typeof event.name === "string" ? event.name : null;
		if (name === snapshot.name) return;
		snapshot.name = name;
		emit("session_info", { session: snapshot });
	});

	sub(
		"model_select",
		(event: { model?: { id?: unknown; provider?: unknown } }) => {
			if (snapshot === null) return;
			const model = event.model;
			const id =
				model !== undefined && typeof model.id === "string" ? model.id : null;
			const provider =
				model !== undefined && typeof model.provider === "string"
					? model.provider
					: null;
			if (id === snapshot.model && provider === snapshot.provider) return;
			snapshot.model = id;
			snapshot.provider = provider;
			emit("session_info", { session: snapshot });
		},
	);

	sub("message_start", (event: { message?: unknown }) => {
		const mapped = mapMessage(event.message);
		if (mapped === null) return;
		if (mapped.role === "assistant") lastStreamText = "";
		emit("message_start", { message: { role: mapped.role, id: mapped.id } });
	});

	sub("message_update", (event: { message?: unknown }) => {
		const mapped = mapMessage(event.message);
		if (mapped === null || mapped.role !== "assistant") return;
		const text = typeof mapped.text === "string" ? mapped.text : "";
		// Delta only: pi hands us accumulated text; diff against what we sent.
		const delta: string = text.startsWith(lastStreamText)
			? text.slice(lastStreamText.length)
			: text;
		lastStreamText = text;
		if (delta.length > 0) emit("message_update", { delta });
	});

	sub("message_end", (event: { message?: unknown }) => {
		const mapped = mapMessage(event.message);
		if (mapped === null) return;
		lastStreamText = "";
		emit("message_end", { message: mapped });
	});

	sub(
		"tool_execution_start",
		(event: { toolCallId?: unknown; toolName?: unknown; args?: unknown }) => {
			emit("tool_execution_start", {
				toolCallId: String(event.toolCallId ?? ""),
				toolName: String(event.toolName ?? ""),
				args: event.args ?? null,
			});
		},
	);

	sub(
		"tool_execution_update",
		(event: {
			toolCallId?: unknown;
			toolName?: unknown;
			partialResult?: { content?: unknown };
		}) => {
			const partial: string = textOfContent(event.partialResult?.content).slice(
				0,
				RESULT_PREVIEW_MAX_CHARS,
			);
			emit("tool_execution_update", {
				toolCallId: String(event.toolCallId ?? ""),
				toolName: String(event.toolName ?? ""),
				partial,
			});
		},
	);

	sub(
		"tool_execution_end",
		(event: {
			toolCallId?: unknown;
			toolName?: unknown;
			result?: { content?: unknown };
			isError?: unknown;
		}) => {
			const preview: string = textOfContent(event.result?.content).slice(
				0,
				RESULT_PREVIEW_MAX_CHARS,
			);
			emit("tool_execution_end", {
				toolCallId: String(event.toolCallId ?? ""),
				toolName: String(event.toolName ?? ""),
				isError: event.isError === true,
				resultPreview: preview,
			});
		},
	);

	sub("agent_start", () => {
		emit("agent_start", {});
	});

	sub("agent_end", (event: { messages?: unknown }) => {
		let inputTokens = 0;
		let outputTokens = 0;
		let totalCost = 0;
		let seen = false;
		if (Array.isArray(event.messages)) {
			for (const m of event.messages) {
				if (
					m === null ||
					typeof m !== "object" ||
					(m as { role?: unknown }).role !== "assistant"
				)
					continue;
				const u = (
					m as {
						usage?: {
							input?: unknown;
							output?: unknown;
							cost?: { total?: unknown };
						};
					}
				).usage;
				if (u === undefined) continue;
				seen = true;
				inputTokens += typeof u.input === "number" ? u.input : 0;
				outputTokens += typeof u.output === "number" ? u.output : 0;
				totalCost +=
					u.cost !== undefined && typeof u.cost.total === "number"
						? u.cost.total
						: 0;
			}
		}
		emit("agent_end", {
			usage: seen ? { inputTokens, outputTokens, totalCost } : {},
		});
	});

	sub("agent_settled", () => {
		emit("agent_settled", {});
	});
}