Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
14 kB · 335 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336import { mkdtempSync, rmSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { afterEach, describe, expect, it, vi } from "vitest";import type { DaemonCommand, DaemonResponse } from "../src/modes/daemon/daemon-protocol.js";import { DaemonSupervisor } from "../src/modes/daemon/daemon-supervisor.js";
type ClientFixture = { id: string; socket: { destroy: ReturnType<typeof vi.fn> }; attachedActiveSessionIds: Set<string>; capabilities: Set<string>;};
type WorkerFixture = { descriptor: { workerId: string }; client?: { close: ReturnType<typeof vi.fn> };};
type PauseEntry = { owner: ClientFixture; worker: WorkerFixture; pauseId: string };
type SupervisorInternals = { clients: Set<ClientFixture>; workers: Map<string, WorkerFixture>; connectionIds: WeakMap<ClientFixture, string>; protocolClientIds: WeakMap<ClientFixture, string>; sessionInputPauseEpochs: WeakMap<ClientFixture, number>; detachingInputPauseSessions: WeakMap<ClientFixture, Set<string>>; sessionInputPauses: Map<string, PauseEntry>; findWorkerForClient: ReturnType<typeof vi.fn>; forwardToWorker: ReturnType<typeof vi.fn>; attachClient: ReturnType<typeof vi.fn>; reserveSnapshotStream: ReturnType<typeof vi.fn>; write: ReturnType<typeof vi.fn>; detachClient: ReturnType<typeof vi.fn>; handleCommand(client: ClientFixture, command: DaemonCommand): Promise<DaemonResponse | undefined>; invalidateWorkerSessionInputPauses(worker: WorkerFixture, reason: string): void; handleWorkerClose(worker: WorkerFixture, client: { close: ReturnType<typeof vi.fn> }, error: Error): Promise<void>;};
const tempDirs: string[] = [];
afterEach(() => { for (const directory of tempDirs.splice(0)) rmSync(directory, { recursive: true, force: true });});
function createHarness() { const directory = mkdtempSync(join(tmpdir(), "prime-supervisor-input-pause-")); tempDirs.push(directory); const supervisor = new DaemonSupervisor(join(directory, "daemon.sock"), { defaultSessionConfig: { agentDir: directory, cwd: directory }, descriptorDir: join(directory, "workers"), }) as unknown as SupervisorInternals; const worker: WorkerFixture = { descriptor: { workerId: "worker-1" }, client: { close: vi.fn() } }; supervisor.workers.set(worker.descriptor.workerId, worker); supervisor.findWorkerForClient = vi.fn(async () => ({ worker, summary: { id: "active-1", activeSessionId: "active-1" }, })); let pauseSequence = 0; supervisor.forwardToWorker = vi.fn(async (_worker: WorkerFixture, command: DaemonCommand) => { if (command.type === "acquire_session_input_pause") { pauseSequence++; return { type: "response", command: command.type, success: true, data: { pauseId: `pause-${pauseSequence}` }, } satisfies DaemonResponse; } return { type: "response", command: command.type, success: true } satisfies DaemonResponse; }); supervisor.attachClient = vi.fn(async () => ({ worker, result: { activeSessionId: "active-1" } })); supervisor.reserveSnapshotStream = vi.fn(() => vi.fn()); supervisor.write = vi.fn(); supervisor.detachClient = vi.fn((client: ClientFixture, activeSessionId?: string) => { if (activeSessionId) client.attachedActiveSessionIds.delete(activeSessionId); else client.attachedActiveSessionIds.clear(); }); supervisor.handleWorkerClose = vi.fn(async (closedWorker: WorkerFixture, workerClient, error: Error) => { if (closedWorker.client !== workerClient) return; closedWorker.client = undefined; supervisor.invalidateWorkerSessionInputPauses(closedWorker, error.message); }); return { supervisor, worker };}
function addClient(supervisor: SupervisorInternals, socketId: string, protocolId: string): ClientFixture { const client = { id: protocolId, socket: { destroy: vi.fn() }, attachedActiveSessionIds: new Set(["active-1"]), capabilities: new Set<string>(), }; supervisor.clients.add(client); supervisor.connectionIds.set(client, socketId); supervisor.protocolClientIds.set(client, protocolId); supervisor.sessionInputPauseEpochs.set(client, 0); supervisor.detachingInputPauseSessions.set(client, new Set()); return client;}
function acquireCommand(id: string): DaemonCommand { return { id, type: "acquire_session_input_pause", activeSessionId: "active-1", leaseKey: "shared-key" };}
describe("daemon supervisor session input pause ownership", () => { it("keeps identical public lease keys independent across client sockets", async () => { const { supervisor } = createHarness(); const first = addClient(supervisor, "socket-a", "protocol-a"); const second = addClient(supervisor, "socket-b", "protocol-b");
const firstResponse = await supervisor.handleCommand(first, acquireCommand("acquire-a")); const secondResponse = await supervisor.handleCommand(second, acquireCommand("acquire-b")); expect(firstResponse).toMatchObject({ success: true, data: { pauseId: "pause-1" } }); expect(secondResponse).toMatchObject({ success: true, data: { pauseId: "pause-2" } }); expect( supervisor.forwardToWorker.mock.calls .map((call) => call[1] as DaemonCommand) .filter((command) => command.type === "acquire_session_input_pause") .map((command) => command.leaseKey), ).toEqual([ JSON.stringify(["socket-a", "protocol-a", "shared-key"]), JSON.stringify(["socket-b", "protocol-b", "shared-key"]), ]);
await supervisor.handleCommand(first, { id: "release-a", type: "release_session_input_pause", activeSessionId: "active-1", pauseId: "pause-1", }); expect(supervisor.sessionInputPauses.has("pause-1")).toBe(false); expect(supervisor.sessionInputPauses.has("pause-2")).toBe(true); });
it("clears stable and live detach selectors after reattach", async () => { const { supervisor } = createHarness(); const client = addClient(supervisor, "socket-a", "protocol-a"); const detachingSessions = supervisor.detachingInputPauseSessions.get(client)!; detachingSessions.add("stable-id");
await supervisor.handleCommand(client, { id: "reattach", type: "reattach", activeSessionId: "old-active", targetActiveSessionId: "stable-id", });
expect(detachingSessions.has("stable-id")).toBe(false); await expect( supervisor.handleCommand(client, { id: "acquire", type: "acquire_session_input_pause", activeSessionId: "stable-id", leaseKey: "shared-key", }), ).resolves.toMatchObject({ success: true, data: { pauseId: "pause-1" } }); });
it("reacquires a fresh worker pause while the prior pause is releasing", async () => { const { supervisor } = createHarness(); const client = addClient(supervisor, "socket-a", "protocol-a"); await supervisor.handleCommand(client, acquireCommand("acquire-1")); let finishRelease!: () => void; const releaseGate = new Promise<void>((resolve) => { finishRelease = resolve; }); supervisor.forwardToWorker.mockImplementationOnce(async (_worker: WorkerFixture, command: DaemonCommand) => { await releaseGate; return { type: "response", command: command.type, success: true } satisfies DaemonResponse; });
const release = supervisor.handleCommand(client, { id: "release-1", type: "release_session_input_pause", activeSessionId: "active-1", pauseId: "pause-1", }); const reacquired = await supervisor.handleCommand(client, acquireCommand("acquire-2")); expect(reacquired).toMatchObject({ success: true, data: { pauseId: "pause-2" } });
finishRelease(); await release; expect(supervisor.sessionInputPauses.has("pause-1")).toBe(false); expect(supervisor.sessionInputPauses.has("pause-2")).toBe(true); });
it("deduplicates retries and releases the client lease on detach", async () => { const { supervisor } = createHarness(); const client = addClient(supervisor, "socket-a", "protocol-a");
const first = await supervisor.handleCommand(client, acquireCommand("acquire-1")); const retry = await supervisor.handleCommand(client, acquireCommand("acquire-2")); expect(first).toMatchObject({ success: true, data: { pauseId: "pause-1" } }); expect(retry).toMatchObject({ success: true, data: { pauseId: "pause-1" } }); expect( supervisor.forwardToWorker.mock.calls.filter( (call) => (call[1] as DaemonCommand).type === "acquire_session_input_pause", ), ).toHaveLength(1);
await supervisor.handleCommand(client, { id: "detach", type: "detach", activeSessionId: "active-1" }); expect(supervisor.sessionInputPauses.size).toBe(0); expect(supervisor.detachClient).toHaveBeenCalledWith(client, "active-1"); const releaseCalls = supervisor.forwardToWorker.mock.calls.filter( (call) => (call[1] as DaemonCommand).type === "release_session_input_pause", ); expect(releaseCalls).toHaveLength(1); expect(releaseCalls[0]?.[2]).toBe(5_000); });
it("blocks new acquisitions before detach waits for cleanup", async () => { const { supervisor } = createHarness(); const client = addClient(supervisor, "socket-a", "protocol-a"); await supervisor.handleCommand(client, acquireCommand("acquire")); let releaseCleanup!: () => void; const cleanupGate = new Promise<void>((resolve) => { releaseCleanup = resolve; }); supervisor.forwardToWorker.mockImplementationOnce(async (_worker: WorkerFixture, command: DaemonCommand) => { await cleanupGate; return { type: "response", command: command.type, success: true } as DaemonResponse; });
const detach = supervisor.handleCommand(client, { id: "detach", type: "detach", activeSessionId: "active-1", }); await vi.waitFor(() => expect(supervisor.detachClient).toHaveBeenCalledWith(client, "active-1")); await expect(supervisor.handleCommand(client, acquireCommand("late-acquire"))).rejects.toThrow( "Session is detaching", ); releaseCleanup(); await expect(detach).resolves.toMatchObject({ success: true }); expect(supervisor.sessionInputPauses.size).toBe(0); });
it("fails closed for every owner when detach cannot release the worker pause", async () => { const { supervisor, worker } = createHarness(); const client = addClient(supervisor, "socket-a", "protocol-a"); const other = addClient(supervisor, "socket-b", "protocol-b"); await supervisor.handleCommand(client, acquireCommand("acquire-a")); await supervisor.handleCommand(other, acquireCommand("acquire-b")); supervisor.forwardToWorker.mockRejectedValueOnce(new Error("worker unavailable"));
await expect( supervisor.handleCommand(client, { id: "detach", type: "detach", activeSessionId: "active-1" }), ).rejects.toThrow("worker unavailable"); expect(supervisor.sessionInputPauses.size).toBe(0); expect(supervisor.handleWorkerClose).toHaveBeenCalledTimes(1); expect(worker.client).toBeUndefined(); expect(client.socket.destroy).toHaveBeenCalled(); expect(other.socket.destroy).toHaveBeenCalled(); });
it("releases all client pauses when detach omits a session id", async () => { const { supervisor } = createHarness(); const client = addClient(supervisor, "socket-a", "protocol-a"); await supervisor.handleCommand(client, acquireCommand("acquire"));
await supervisor.handleCommand(client, { id: "detach-all", type: "detach" });
expect(supervisor.sessionInputPauses.size).toBe(0); expect(client.attachedActiveSessionIds.size).toBe(0); expect(supervisor.detachClient).toHaveBeenCalledWith(client, undefined); });
it("rejects release from a different public client", async () => { const { supervisor } = createHarness(); const owner = addClient(supervisor, "socket-a", "protocol-a"); const other = addClient(supervisor, "socket-b", "protocol-b"); await supervisor.handleCommand(owner, acquireCommand("acquire"));
await expect( supervisor.handleCommand(other, { id: "release", type: "release_session_input_pause", activeSessionId: "active-1", pauseId: "pause-1", }), ).rejects.toThrow("owned by another client"); });
it("invalidates every public owner when the worker connection closes", async () => { const { supervisor, worker } = createHarness(); const first = addClient(supervisor, "socket-a", "protocol-a"); const second = addClient(supervisor, "socket-b", "protocol-b"); await supervisor.handleCommand(first, acquireCommand("acquire-a")); await supervisor.handleCommand(second, acquireCommand("acquire-b"));
supervisor.invalidateWorkerSessionInputPauses(worker, "worker pause invalidated");
expect(supervisor.sessionInputPauses.size).toBe(0); expect(first.socket.destroy).toHaveBeenCalledWith( expect.objectContaining({ message: "worker pause invalidated" }), ); expect(second.socket.destroy).toHaveBeenCalledWith( expect.objectContaining({ message: "worker pause invalidated" }), ); });
it("releases an acquire that completes after its socket is invalidated", async () => { const { supervisor } = createHarness(); const client = addClient(supervisor, "socket-a", "protocol-a"); let releaseAcquire!: () => void; const gate = new Promise<void>((resolve) => { releaseAcquire = resolve; }); const forward = supervisor.forwardToWorker; forward.mockImplementationOnce(async (_worker: WorkerFixture, command: DaemonCommand) => { await gate; return { type: "response", command: command.type, success: true, data: { pauseId: "pause-late" }, } as DaemonResponse; });
const acquisition = supervisor.handleCommand(client, acquireCommand("acquire")); await vi.waitFor(() => expect(forward).toHaveBeenCalledTimes(1)); supervisor.sessionInputPauseEpochs.set(client, 1); supervisor.clients.delete(client); releaseAcquire();
await expect(acquisition).rejects.toThrow("invalidated before completion"); expect(forward.mock.calls.map((call) => (call[1] as DaemonCommand).type)).toEqual([ "acquire_session_input_pause", "release_session_input_pause", ]); expect(supervisor.sessionInputPauses.size).toBe(0); });});