pi-gui/test/agent-registry.test.js

import assert from "node:assert/strict";
import { mkdir, mkdtemp, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import test from "node:test";
import { createAgentRegistry } from "../src/bridge/agent-registry.js";

async function createWorktrees() {
	const root = await mkdtemp(join(tmpdir(), "pi-status-bridge-worktrees-"));
	const home = join(root, "home");
	const feature = join(root, "feature");
	await Promise.all([mkdir(home), mkdir(feature)]);
	return { root, home, feature, sessionRoot: join(root, "sessions") };
}

function createAdapterFactory() {
	const calls = [];
	return {
		calls,
		startAdapter: (options) => {
			const adapter = {
				sent: [],
				extensionResponses: [],
				stopped: false,
				send(command) {
					this.sent.push(command);
					return Promise.resolve({
						type: "response",
						command: command.type,
						success: true,
					});
				},
				respondToExtension(requestId, response) {
					this.extensionResponses.push({ requestId, response });
				},
				async stop() {
					this.stopped = true;
				},
			};
			calls.push({ options, adapter });
			return adapter;
		},
	};
}

test("starts only the home agent and creates other worktree agents on explicit selection", async () => {
	const worktrees = await createWorktrees();
	const fixture = createAdapterFactory();
	const registry = createAgentRegistry({
		homeWorktree: worktrees.home,
		sessionRoot: worktrees.sessionRoot,
		startAdapter: fixture.startAdapter,
	});

	const home = await registry.start();
	assert.equal(registry.listAgents().length, 1);
	assert.equal(home.worktreePath, worktrees.home);

	const feature = await registry.selectWorktree(worktrees.feature);
	const selectedAgain = await registry.selectWorktree(
		`${worktrees.feature}/../feature`,
	);
	assert.equal(feature.id, selectedAgain.id);
	assert.notEqual(home.id, feature.id);
	assert.deepEqual(
		fixture.calls.map(({ options }) => options.cwd),
		[worktrees.home, worktrees.feature],
	);
	assert.notEqual(home.sessionDir, feature.sessionDir);

	await registry.stop();
	assert.ok(fixture.calls.every(({ adapter }) => adapter.stopped));
});

test("catalogues managed directories and exposes only their sessions", async () => {
	const worktrees = await createWorktrees();
	const fixture = createAdapterFactory();
	const registry = createAgentRegistry({
		homeWorktree: worktrees.home,
		sessionRoot: worktrees.sessionRoot,
		startAdapter: fixture.startAdapter,
	});
	await registry.start();
	const feature = await registry.selectWorktree(worktrees.feature);
	const olderSession = join(feature.sessionDir, "older.jsonl");
	const newestSession = join(feature.sessionDir, "newest.jsonl");
	await writeFile(
		olderSession,
		[
			JSON.stringify({
				type: "session",
				id: "older",
				timestamp: "2025-01-01T00:00:00.000Z",
				cwd: worktrees.feature,
			}),
			JSON.stringify({
				type: "message",
				message: { role: "user", content: "Older task" },
			}),
		].join("\n"),
	);
	await writeFile(
		newestSession,
		[
			JSON.stringify({
				type: "session",
				id: "newest",
				timestamp: "2025-01-02T00:00:00.000Z",
				cwd: worktrees.feature,
			}),
			JSON.stringify({ type: "session_info", name: "Newest work" }),
			JSON.stringify({
				type: "message",
				message: { role: "user", content: "Newest task" },
			}),
		].join("\n"),
	);
	await writeFile(
		join(feature.sessionDir, "bridge-agent.json"),
		`${JSON.stringify({ sessionPath: newestSession, worktreePath: worktrees.feature })}\n`,
	);

	const sessions = await registry.listSessions(feature.id);
	assert.deepEqual(
		sessions.map(({ id, name, isCurrent }) => ({ id, name, isCurrent })),
		[
			{ id: "newest", name: "Newest work", isCurrent: false },
			{ id: "older", name: undefined, isCurrent: false },
		],
	);
	await registry.route(feature.id, "switch_session", {
		sessionPath: newestSession,
	});
	await registry.route(feature.id, "new_session");
	assert.deepEqual(
		fixture.calls[1].adapter.sent.map((command) => command.type),
		["get_state", "switch_session", "get_state", "new_session", "get_state"],
	);
	await assert.rejects(
		registry.route(feature.id, "switch_session", {
			sessionPath: olderSession.replace("older", "../outside"),
		}),
		/ENOENT|belong/,
	);
	await registry.stop();

	const restored = createAgentRegistry({
		homeWorktree: worktrees.home,
		sessionRoot: worktrees.sessionRoot,
		startAdapter: createAdapterFactory().startAdapter,
	});
	await restored.start();
	assert.deepEqual(await restored.listDirectories(), [
		{ worktreePath: worktrees.feature, state: "inactive" },
		{
			worktreePath: worktrees.home,
			state: "idle",
			agentId: restored.listAgents()[0].id,
		},
	]);
	await restored.stop();
});

test("routes commands by explicit agent ID and replays only events after the cursor", async () => {
	const worktrees = await createWorktrees();
	const fixture = createAdapterFactory();
	const registry = createAgentRegistry({
		homeWorktree: worktrees.home,
		sessionRoot: worktrees.sessionRoot,
		startAdapter: fixture.startAdapter,
	});
	const agent = await registry.start();
	const [{ options, adapter }] = fixture.calls;

	const response = await registry.route(agent.id, "prompt", {
		message: "Stay in this worktree",
	});
	assert.deepEqual(response, {
		type: "response",
		command: "prompt",
		success: true,
	});
	const stats = await registry.route(agent.id, "get_session_stats");
	assert.deepEqual(stats, {
		type: "response",
		command: "get_session_stats",
		success: true,
	});
	const commands = await registry.route(agent.id, "get_commands");
	assert.deepEqual(commands, {
		type: "response",
		command: "get_commands",
		success: true,
	});
	assert.deepEqual(adapter.sent, [
		{ type: "get_state" },
		{ type: "prompt", message: "Stay in this worktree" },
		{ type: "get_session_stats" },
		{ type: "get_commands" },
	]);

	options.onEvent({
		type: "agent_state",
		data: { state: "streaming" },
	});
	await registry.route(agent.id, "submit_prompt", {
		message: "After this turn",
	});
	assert.deepEqual(adapter.sent.at(-1), {
		type: "follow_up",
		message: "After this turn",
	});

	options.onEvent({
		type: "stream",
		data: { event: { type: "message_update" } },
	});
	options.onEvent({ type: "queue", data: { event: { type: "queue_update" } } });
	assert.deepEqual(
		registry.eventsAfter(agent.id, 0).map(({ seq, type }) => ({ seq, type })),
		[
			{ seq: 1, type: "agent_state" },
			{ seq: 2, type: "stream" },
			{ seq: 3, type: "queue" },
		],
	);
	assert.deepEqual(
		registry.eventsAfter(agent.id, 1).map(({ seq, type }) => ({ seq, type })),
		[
			{ seq: 2, type: "stream" },
			{ seq: 3, type: "queue" },
		],
	);

	await assert.rejects(
		registry.route("missing-agent", "abort"),
		/unknown agent/,
	);
	await registry.stop();
});