@katyella/legio 0.1.0

package/src/server/websocket.test.ts

@@ -0,0 +1,288 @@
+import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import type { WebSocket } from "ws";
+import { createWebSocketManager, type WebSocketData } from "./websocket.ts";
+
+let tempDir: string;
+let legioDir: string;
+
+beforeEach(async () => {
+	tempDir = await mkdtemp(join(tmpdir(), "ws-test-"));
+	legioDir = join(tempDir, ".legio");
+	await mkdir(legioDir, { recursive: true });
+	await writeFile(
+		join(legioDir, "config.yaml"),
+		"project:\n  name: test\n  canonicalBranch: main\nagents:\n  maxDepth: 2\ncoordinator:\n  model: sonnet\n",
+	);
+});
+
+afterEach(async () => {
+	await rm(tempDir, { recursive: true, force: true });
+});
+
+describe("createWebSocketManager", () => {
+	it("returns a manager with the expected interface", () => {
+		const manager = createWebSocketManager(legioDir);
+		expect(typeof manager.addClient).toBe("function");
+		expect(typeof manager.removeClient).toBe("function");
+		expect(typeof manager.handleMessage).toBe("function");
+		expect(typeof manager.startPolling).toBe("function");
+		expect(typeof manager.stopPolling).toBe("function");
+	});
+
+	it("startPolling and stopPolling do not throw", () => {
+		const manager = createWebSocketManager(legioDir);
+		expect(() => manager.startPolling()).not.toThrow();
+		expect(() => manager.stopPolling()).not.toThrow();
+	});
+
+	it("stopPolling is idempotent", () => {
+		const manager = createWebSocketManager(legioDir);
+		expect(() => {
+			manager.stopPolling();
+			manager.stopPolling();
+		}).not.toThrow();
+	});
+
+	it("gatherSnapshot returns valid shape with missing dbs", () => {
+		// Use a non-existent dir — all stores will fail to open and fall back to empty data
+		const manager = createWebSocketManager(join(tempDir, "no-such-dir"));
+		// We test gatherSnapshot indirectly via addClient's initial snapshot send
+		let received: string | null = null;
+		const fakeWs = {
+			send(msg: string) {
+				received = msg;
+			},
+		} as unknown as WebSocket;
+		// WebSocketData is unused at runtime but keeps test aligned with implementation
+		const _unused: WebSocketData = { connectedAt: "" };
+
+		manager.addClient(fakeWs);
+
+		expect(received).not.toBeNull();
+		if (received === null) throw new Error("expected a message");
+		const parsed = JSON.parse(received);
+		expect(parsed.type).toBe("snapshot");
+		expect(parsed.data).toBeDefined();
+		expect(Array.isArray(parsed.data.agents)).toBe(true);
+		expect(typeof parsed.data.mail.unreadCount).toBe("number");
+		expect(Array.isArray(parsed.data.mergeQueue)).toBe(true);
+		expect(typeof parsed.data.metricsSummary.totalSessions).toBe("number");
+		expect(parsed.timestamp).toBeDefined();
+	});
+
+	it("addClient sends initial snapshot immediately", () => {
+		const manager = createWebSocketManager(legioDir);
+		let received: string | null = null;
+
+		const fakeWs = {
+			send(msg: string) {
+				received = msg;
+			},
+		} as unknown as WebSocket;
+
+		manager.addClient(fakeWs);
+		expect(received).not.toBeNull();
+		if (received === null) throw new Error("expected a message");
+
+		const snapshot = JSON.parse(received);
+		expect(snapshot.type).toBe("snapshot");
+	});
+
+	it("removeClient stops further sends to that client", () => {
+		const manager = createWebSocketManager(legioDir);
+		let count = 0;
+
+		const fakeWs = {
+			send() {
+				count++;
+			},
+		} as unknown as WebSocket;
+
+		manager.addClient(fakeWs); // triggers 1 send (initial snapshot)
+		const countAfterAdd = count;
+		expect(countAfterAdd).toBe(1);
+
+		manager.removeClient(fakeWs);
+
+		// Trigger broadcast — should not reach removed client
+		const fakeWs2 = {
+			send() {
+				count++;
+			},
+		} as unknown as WebSocket;
+		manager.addClient(fakeWs2); // sends to ws2, not to ws1
+
+		// count should only increase by 1 (for fakeWs2 add), not for fakeWs
+		expect(count).toBe(2);
+	});
+
+	it("handleMessage with {type:'refresh'} sends a snapshot to that client", () => {
+		const manager = createWebSocketManager(legioDir);
+		const received: string[] = [];
+
+		const fakeWs = {
+			send(msg: string) {
+				received.push(msg);
+			},
+		} as unknown as WebSocket;
+
+		manager.addClient(fakeWs); // initial snapshot = received[0]
+		manager.handleMessage(fakeWs, Buffer.from(JSON.stringify({ type: "refresh" })));
+
+		expect(received.length).toBe(2);
+		const refreshed = JSON.parse(received[1] as string);
+		expect(refreshed.type).toBe("snapshot");
+	});
+
+	it("handleMessage with unknown type does not throw", () => {
+		const manager = createWebSocketManager(legioDir);
+		const fakeWs = {
+			send() {},
+		} as unknown as WebSocket;
+
+		manager.addClient(fakeWs);
+		expect(() =>
+			manager.handleMessage(fakeWs, Buffer.from(JSON.stringify({ type: "unknown" }))),
+		).not.toThrow();
+	});
+
+	it("handleMessage with invalid JSON does not throw", () => {
+		const manager = createWebSocketManager(legioDir);
+		const fakeWs = {
+			send() {},
+		} as unknown as WebSocket;
+
+		manager.addClient(fakeWs);
+		expect(() => manager.handleMessage(fakeWs, Buffer.from("not-json{{{"))).not.toThrow();
+	});
+
+	it("broadcastEvent sends event to all connected clients", () => {
+		const manager = createWebSocketManager(legioDir);
+		const messagesA: string[] = [];
+		const messagesB: string[] = [];
+
+		const ws1 = {
+			send(msg: string) {
+				messagesA.push(msg);
+			},
+		} as unknown as WebSocket;
+
+		const ws2 = {
+			send(msg: string) {
+				messagesB.push(msg);
+			},
+		} as unknown as WebSocket;
+
+		manager.addClient(ws1); // initial snapshot
+		manager.addClient(ws2); // initial snapshot
+
+		manager.broadcastEvent({ type: "mail_new", data: { id: "msg-test-001" } });
+
+		// Each client received: 1 initial snapshot + 1 broadcastEvent
+		expect(messagesA.length).toBe(2);
+		expect(messagesB.length).toBe(2);
+
+		const evtA = JSON.parse(messagesA[1] as string);
+		expect(evtA.type).toBe("mail_new");
+		expect(evtA.timestamp).toBeDefined();
+		expect((evtA.data as { id: string }).id).toBe("msg-test-001");
+	});
+
+	it("broadcastEvent with no clients does not throw", () => {
+		const manager = createWebSocketManager(legioDir);
+		expect(() => manager.broadcastEvent({ type: "mail_new" })).not.toThrow();
+	});
+
+	it("broadcastEvent event format includes timestamp", () => {
+		const manager = createWebSocketManager(legioDir);
+		const received: string[] = [];
+
+		const fakeWs = {
+			send(msg: string) {
+				received.push(msg);
+			},
+		} as unknown as WebSocket;
+
+		manager.addClient(fakeWs); // initial snapshot (received[0])
+		manager.broadcastEvent({ type: "test_event", data: { foo: "bar" } });
+
+		expect(received.length).toBe(2);
+		const evt = JSON.parse(received[1] as string);
+		expect(evt.type).toBe("test_event");
+		expect(typeof evt.timestamp).toBe("string");
+		expect(() => new Date(evt.timestamp)).not.toThrow();
+		expect((evt.data as { foo: string }).foo).toBe("bar");
+	});
+
+	it("multiple clients receive snapshot when added", () => {
+		const manager = createWebSocketManager(legioDir);
+		const messages: Record<string, string[]> = { a: [], b: [] };
+
+		const ws1 = {
+			send(msg: string) {
+				messages.a?.push(msg);
+			},
+		} as unknown as WebSocket;
+
+		const ws2 = {
+			send(msg: string) {
+				messages.b?.push(msg);
+			},
+		} as unknown as WebSocket;
+
+		manager.addClient(ws1);
+		manager.addClient(ws2);
+
+		expect(messages.a?.length).toBe(1);
+		expect(messages.b?.length).toBe(1);
+
+		const snap1 = JSON.parse(messages.a?.[0] ?? "{}");
+		const snap2 = JSON.parse(messages.b?.[0] ?? "{}");
+		expect(snap1.type).toBe("snapshot");
+		expect(snap2.type).toBe("snapshot");
+	});
+});
+
+describe("WebSocket integration (real server)", () => {
+	it("client receives initial snapshot on connect and can request refresh", async () => {
+		// Dynamic import to avoid top-level routes.ts dependency issues
+		const { createServer } = await import("./index.ts");
+		const { WebSocket: WsClient } = await import("ws");
+
+		const server = await createServer({ port: 0, host: "localhost", root: tempDir });
+		const port = server.port;
+
+		try {
+			const ws = new WsClient(`ws://localhost:${port}/ws`);
+
+			const messages: string[] = [];
+			await new Promise<void>((resolve, reject) => {
+				const timeout = setTimeout(() => reject(new Error("timeout")), 5000);
+				ws.onmessage = (e) => {
+					messages.push(e.data as string);
+					if (messages.length === 1) {
+						// Got initial snapshot, send refresh
+						ws.send(JSON.stringify({ type: "refresh" }));
+					} else {
+						clearTimeout(timeout);
+						resolve();
+					}
+				};
+				ws.onerror = () => reject(new Error("ws error"));
+			});
+
+			expect(messages.length).toBeGreaterThanOrEqual(2);
+			const initial = JSON.parse(messages[0] as string);
+			expect(initial.type).toBe("snapshot");
+			const refreshed = JSON.parse(messages[1] as string);
+			expect(refreshed.type).toBe("snapshot");
+
+			ws.close();
+		} finally {
+			server.stop(true);
+		}
+	});
+});

package/src/server/websocket.ts

@@ -0,0 +1,196 @@
+import { join } from "node:path";
+import type { RawData, WebSocket } from "ws";
+import { createMailStore } from "../mail/store.ts";
+import { createMergeQueue } from "../merge/queue.ts";
+import { createMetricsStore } from "../metrics/store.ts";
+import { openSessionStore } from "../sessions/compat.ts";
+import { createRunStore } from "../sessions/store.ts";
+export interface WebSocketData {
+	connectedAt: string;
+}
+
+interface Snapshot {
+	type: "snapshot";
+	data: {
+		agents: unknown[];
+		mail: { unreadCount: number; recent: unknown[] };
+		mergeQueue: unknown[];
+		metricsSummary: { totalSessions: number; avgDuration: number };
+		runs: { active: unknown | null };
+	};
+	timestamp: string;
+}
+
+export interface WebSocketManager {
+	addClient(ws: WebSocket): void;
+	removeClient(ws: WebSocket): void;
+	handleMessage(ws: WebSocket, message: RawData): void;
+	startPolling(): void;
+	stopPolling(): void;
+	broadcastEvent(event: { type: string; data?: unknown }): void;
+}
+
+export function createWebSocketManager(legioDir: string): WebSocketManager {
+	const clients = new Set<WebSocket>();
+	let pollInterval: ReturnType<typeof setInterval> | null = null;
+	let lastSnapshot = "";
+
+	function gatherSnapshot(): Snapshot {
+		const data: Snapshot["data"] = {
+			agents: [],
+			mail: { unreadCount: 0, recent: [] },
+			mergeQueue: [],
+			metricsSummary: { totalSessions: 0, avgDuration: 0 },
+			runs: { active: null },
+		};
+
+		// Session store
+		try {
+			const { store } = openSessionStore(legioDir);
+			try {
+				data.agents = store.getActive();
+			} finally {
+				store.close();
+			}
+		} catch {
+			/* db may not exist */
+		}
+
+		// Mail store
+		try {
+			const mailPath = join(legioDir, "mail.db");
+			const store = createMailStore(mailPath);
+			try {
+				const all = store.getAll();
+				data.mail.unreadCount = all.filter((m: { read: boolean }) => !m.read).length;
+				data.mail.recent = all.slice(0, 20);
+			} finally {
+				store.close();
+			}
+		} catch {
+			/* db may not exist */
+		}
+
+		// Merge queue
+		try {
+			const queue = createMergeQueue(join(legioDir, "merge-queue.db"));
+			try {
+				data.mergeQueue = queue.list();
+			} finally {
+				queue.close();
+			}
+		} catch {
+			/* db may not exist */
+		}
+
+		// Metrics
+		try {
+			const store = createMetricsStore(join(legioDir, "metrics.db"));
+			try {
+				const sessions = store.getRecentSessions(100);
+				data.metricsSummary.totalSessions = sessions.length;
+				data.metricsSummary.avgDuration = store.getAverageDuration();
+			} finally {
+				store.close();
+			}
+		} catch {
+			/* db may not exist */
+		}
+
+		// Runs
+		try {
+			const store = createRunStore(join(legioDir, "sessions.db"));
+			try {
+				data.runs.active = store.getActiveRun();
+			} finally {
+				store.close();
+			}
+		} catch {
+			/* db may not exist */
+		}
+
+		return {
+			type: "snapshot",
+			data,
+			timestamp: new Date().toISOString(),
+		};
+	}
+
+	function broadcast(snapshot: Snapshot): void {
+		const msg = JSON.stringify(snapshot);
+		for (const client of clients) {
+			try {
+				client.send(msg);
+			} catch {
+				clients.delete(client);
+			}
+		}
+	}
+
+	return {
+		addClient(ws) {
+			clients.add(ws);
+			// Send initial snapshot immediately
+			const snapshot = gatherSnapshot();
+			try {
+				ws.send(JSON.stringify(snapshot));
+			} catch {
+				clients.delete(ws);
+			}
+		},
+
+		removeClient(ws) {
+			clients.delete(ws);
+		},
+
+		handleMessage(ws, message) {
+			try {
+				const msgStr =
+					typeof message === "string"
+						? message
+						: Buffer.isBuffer(message)
+							? message.toString("utf8")
+							: Array.isArray(message)
+								? Buffer.concat(message as Buffer[]).toString("utf8")
+								: Buffer.from(message as ArrayBuffer).toString("utf8");
+				const parsed = JSON.parse(msgStr);
+				if (parsed && typeof parsed === "object" && "type" in parsed && parsed.type === "refresh") {
+					const snapshot = gatherSnapshot();
+					ws.send(JSON.stringify(snapshot));
+				}
+			} catch {
+				// Ignore invalid messages
+			}
+		},
+
+		startPolling() {
+			pollInterval = setInterval(() => {
+				if (clients.size === 0) return; // Skip if no clients
+				const snapshot = gatherSnapshot();
+				const snapshotStr = JSON.stringify(snapshot.data);
+				if (snapshotStr !== lastSnapshot) {
+					lastSnapshot = snapshotStr;
+					broadcast(snapshot);
+				}
+			}, 2000);
+		},
+
+		stopPolling() {
+			if (pollInterval) {
+				clearInterval(pollInterval);
+				pollInterval = null;
+			}
+		},
+
+		broadcastEvent(event) {
+			const msg = JSON.stringify({ ...event, timestamp: new Date().toISOString() });
+			for (const client of clients) {
+				try {
+					client.send(msg);
+				} catch {
+					clients.delete(client);
+				}
+			}
+		},
+	};
+}

package/src/sessions/compat.test.ts

@@ -0,0 +1,109 @@
+/**
+ * Tests for the session compat shim.
+ *
+ * Uses real filesystem and better-sqlite3. No mocks.
+ */
+
+import { mkdtemp, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, test } from "vitest";
+import { openSessionStore } from "./compat.ts";
+
+let tempDir: string;
+let legioDir: string;
+
+beforeEach(async () => {
+	tempDir = await mkdtemp(join(tmpdir(), "legio-compat-test-"));
+	legioDir = join(tempDir, ".legio");
+	const { mkdir } = await import("node:fs/promises");
+	await mkdir(legioDir, { recursive: true });
+});
+
+afterEach(async () => {
+	await rm(tempDir, { recursive: true, force: true });
+});
+
+describe("openSessionStore", () => {
+	test("creates empty DB when sessions.db does not exist", () => {
+		const { store, migrated } = openSessionStore(legioDir);
+
+		expect(migrated).toBe(false);
+		expect(store.getAll()).toEqual([]);
+		store.close();
+	});
+
+	test("returns migrated:false always", () => {
+		const { migrated } = openSessionStore(legioDir);
+		expect(migrated).toBe(false);
+	});
+
+	test("returned store supports all SessionStore operations", () => {
+		const { store } = openSessionStore(legioDir);
+
+		const now = new Date().toISOString();
+		store.upsert({
+			id: "s-001",
+			agentName: "test-agent",
+			capability: "builder",
+			worktreePath: "/tmp/worktrees/test-agent",
+			branchName: "legio/test-agent/task-1",
+			beadId: "task-1",
+			tmuxSession: "legio-test-agent",
+			state: "working",
+			pid: 12345,
+			parentAgent: null,
+			depth: 0,
+			runId: null,
+			startedAt: now,
+			lastActivity: now,
+			escalationLevel: 0,
+			stalledSince: null,
+		});
+
+		const all = store.getAll();
+		expect(all).toHaveLength(1);
+		expect(all[0]?.agentName).toBe("test-agent");
+
+		const active = store.getActive();
+		expect(active).toHaveLength(1);
+
+		store.updateState("test-agent", "completed");
+		expect(store.getByName("test-agent")?.state).toBe("completed");
+
+		store.remove("test-agent");
+		expect(store.getByName("test-agent")).toBeNull();
+
+		store.close();
+	});
+
+	test("second call opens existing DB", () => {
+		const { store: store1 } = openSessionStore(legioDir);
+		const now = new Date().toISOString();
+		store1.upsert({
+			id: "s-persist",
+			agentName: "persistent-agent",
+			capability: "scout",
+			worktreePath: "/tmp/worktrees/persistent-agent",
+			branchName: "legio/persistent-agent/task-2",
+			beadId: "task-2",
+			tmuxSession: "legio-persistent-agent",
+			state: "working",
+			pid: null,
+			parentAgent: null,
+			depth: 0,
+			runId: null,
+			startedAt: now,
+			lastActivity: now,
+			escalationLevel: 0,
+			stalledSince: null,
+		});
+		store1.close();
+
+		const { store: store2, migrated } = openSessionStore(legioDir);
+		expect(migrated).toBe(false);
+		expect(store2.getAll()).toHaveLength(1);
+		expect(store2.getByName("persistent-agent")?.id).toBe("s-persist");
+		store2.close();
+	});
+});

package/src/sessions/compat.ts

@@ -0,0 +1,17 @@
+import { join } from "node:path";
+import { createSessionStore, type SessionStore } from "./store.ts";
+
+/**
+ * Open or create a SessionStore at the given .legio directory root.
+ *
+ * @param legioDir - Path to the .legio directory (e.g., /project/.legio)
+ * @returns An object with the SessionStore and whether a migration occurred.
+ */
+export function openSessionStore(legioDir: string): {
+	store: SessionStore;
+	migrated: boolean;
+} {
+	const dbPath = join(legioDir, "sessions.db");
+	const store = createSessionStore(dbPath);
+	return { store, migrated: false };
+}