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.
41 kB · 1029 lines
TypeScript
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030import type { Socket } from "node:net";import { tmpdir } from "node:os";import { join } from "node:path";import { PassThrough } from "node:stream";import type { AgentMessage } from "@earendil-works/pi-agent-core";import { fauxAssistantMessage } from "@earendil-works/pi-ai";import { describe, expect, it, vi } from "vitest";import type { ActiveSessionState, DaemonSocketClient } from "../src/modes/daemon/active-session-state.js";import { AgentDaemon, markClientSnapshotStreaming } from "../src/modes/daemon/daemon-mode.js";import { DAEMON_PROTOCOL_INFO, type DaemonAttachResult, type DaemonCommand, type DaemonOutbound, type DaemonResponse,} from "../src/modes/daemon/daemon-protocol.js";import type { SessionSummary } from "../src/modes/daemon/daemon-session-list.js";import { DaemonSupervisor } from "../src/modes/daemon/daemon-supervisor.js";import { type DaemonWorkerFrameHeader, isDaemonWorkerFrameHeader } from "../src/modes/daemon/daemon-worker-protocol.js";import { SnapshotTranscriptCache } from "../src/modes/daemon/snapshot-transcript-cache.js";import { type PrivateFrame, PrivateFrameDecoder } from "../src/modes/session-worker/private-framing.js";
const activeSessionId = "active-deferred";const snapshotId = "snapshot-deferred";
function summary(): SessionSummary { return { id: activeSessionId, activeSessionId, lifecycle: "live", activity: "idle", isSessionActive: false, sessionId: "session-deferred", cwd: "/tmp", isStreaming: false, isCompacting: false, attachedClients: 0, messageCount: 1, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, };}
function streamedResult(): DaemonAttachResult { return { protocol: DAEMON_PROTOCOL_INFO, activeSessionId, snapshot: { activeSessionId, summary: summary(), state: { activeSessionId, sessionId: "session-deferred" } as DaemonAttachResult["snapshot"]["state"], messages: [], lastEventSequence: 1, lastEventCursor: { generation: "generation-deferred", sequence: 1 }, }, replay: { status: "complete", toSequence: 1 }, lastEventSequence: 1, lastEventCursor: { generation: "generation-deferred", sequence: 1 }, snapshotStream: { id: snapshotId, messageCount: 0, targetChunkBytes: 512 * 1024 }, client: { id: "client", capabilities: ["chunked_snapshot"] }, };}
function sessionEventMessage(sequence: number): DaemonOutbound { return { type: "session_event", activeSessionId, event: { type: "session_info_changed", name: `change-${sequence}` }, meta: { id: `${activeSessionId}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId, sequence, cursor: { generation: "generation-deferred", sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, };}
function socketClient(socket: PassThrough, extra: Partial<DaemonSocketClient> = {}): DaemonSocketClient { return { id: "client", socket: socket as unknown as Socket, attachedActiveSessionIds: new Set([activeSessionId]), catchupActiveSessionIds: new Set(), detachInput: () => {}, supportsExtensionUi: false, capabilities: new Set(["chunked_snapshot"]), ...extra, } as DaemonSocketClient;}
async function nextMacroTaskTurn(): Promise<void> { await new Promise<void>((resolve) => setImmediate(resolve)); await new Promise<void>((resolve) => setImmediate(resolve));}
interface SupervisorWorkerHarness { descriptor: { workerId: string; lifecycle: "ready"; pid: number }; client?: { close: ReturnType<typeof vi.fn>; request: ReturnType<typeof vi.fn> }; summaries: Map<string, SessionSummary>; snapshotCache: Map<string, DaemonAttachResult>; transcriptCaches: Map<string, SnapshotTranscriptCache>; snapshotGenerations: Map<string, Map<string, unknown>>; snapshotLoads: Map<string, Promise<DaemonAttachResult>>; intentionalStop: boolean; stopRevision: number;}
describe("deferred session frames during snapshot streams", () => { it.each(["snapshot", "dropped", "backpressure"])( "delivers extension prompts through worker and supervisor during %s", (phase) => { const daemon = new AgentDaemon(join(tmpdir(), "extension-frames-worker.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, createRuntime: async () => { throw new Error("unexpected runtime creation"); }, }); const supervisor = new DaemonSupervisor(join(tmpdir(), "extension-frames-supervisor.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, descriptorDir: join(tmpdir(), "extension-frames-supervisor-state"), }); const workerSocket = new PassThrough(); const supervisorSocket = new PassThrough(); const workerClient = socketClient(workerSocket, { supportsExtensionUi: true, transport: "private-framed" }); const supervisorClient = socketClient(supervisorSocket, { supportsExtensionUi: true }); for (const client of [workerClient, supervisorClient]) { if (phase === "snapshot") client.snapshotActiveSessionIds = new Set([activeSessionId]); if (phase === "backpressure") client.backpressured = true; if (phase === "dropped") { client.deferredSessionFramesDropped = new Set([activeSessionId]); client.deferredSessionPayloadsDropped = new Set([activeSessionId]); } } const worker = { snapshotCache: new Map(), snapshotLoads: new Map(), transcriptCaches: new Map(), } as unknown as SupervisorWorkerHarness; const routing = supervisor as unknown as { clients: Set<DaemonSocketClient>; handleWorkerFrame(worker: SupervisorWorkerHarness, frame: PrivateFrame<DaemonWorkerFrameHeader>): void; }; routing.clients.add(supervisorClient); const written: DaemonOutbound[] = []; supervisorSocket.on("data", (payload: Buffer) => written.push(JSON.parse(payload.toString()))); const decoder = new PrivateFrameDecoder(isDaemonWorkerFrameHeader); workerSocket.on("data", (payload: Buffer) => { for (const frame of decoder.push(payload)) routing.handleWorkerFrame(worker, frame); }); try { ( daemon as unknown as { broadcastToSession(state: ActiveSessionState, message: DaemonOutbound): void } ).broadcastToSession( { activeSessionId, clients: new Set([workerClient]), lastEventSequence: 1, eventGeneration: "generation-deferred", } as ActiveSessionState, { type: "extension_ui_request", activeSessionId, id: "editor-request", method: "editor", payload: { title: "Edit" }, }, ); expect(written).toEqual([ expect.objectContaining({ type: "extension_ui_request", id: "editor-request", method: "editor" }), ]); expect(workerClient.deferredSessionOutbounds?.size ?? 0).toBe(0); expect(supervisorClient.deferredSessionPayloads?.size ?? 0).toBe(0); } finally { workerSocket.destroy(); supervisorSocket.destroy(); } }, );
it.each([ [1, false, "stream"], [2, false, "stream"], [2, true, "stream"], [1, false, "replacement-failure"], [1, false, "replacement-replay-failure"], [1, false, "replacement-write-failure"], [1, false, "catchup"], [1, false, "inline-catchup"], ] as const)("worker: replay (%i, fail=%s, %s)", async (streamCount, failFinal, preparation) => { const replayFailed = preparation === "replacement-replay-failure" || preparation === "replacement-write-failure"; const daemon = new AgentDaemon(join(tmpdir(), "deferred-frames-worker.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, createRuntime: async () => { throw new Error("unexpected runtime creation"); }, }); const socket = new PassThrough(); socket.on("error", () => {}); const written: Buffer[] = []; socket.on("data", (chunk: Buffer) => written.push(Buffer.from(chunk))); const client = socketClient(socket, { transport: "private-framed" }); if (preparation === "inline-catchup") client.capabilities.clear(); const state = { activeSessionId, clients: new Set([client]), eventGeneration: "generation-deferred", lastEventSequence: 1, runtime: { metadata: { kind: "top-level" as const, createdAt: 1 } }, } as unknown as ActiveSessionState; const internals = daemon as unknown as { sessions: Map<string, ActiveSessionState>; createAttachResult(): Promise<DaemonAttachResult>; flushDeferredSessionFrames(client: DaemonSocketClient, activeSessionId: string): void; write(client: DaemonSocketClient, message: DaemonOutbound): boolean; drainBackpressuredClientCatchups(client: DaemonSocketClient): Promise<"drained" | "retry-later">; catchUpBackpressuredClient(client: DaemonSocketClient): Promise<void>; broadcastToSession(state: ActiveSessionState, message: DaemonOutbound): void; streamWorkerSnapshot( client: DaemonSocketClient, result: DaemonAttachResult, transcript: AsyncIterable<Buffer> & { dispose?(): void }, purpose: "attach", signal: AbortSignal, snapshotAlreadyMarked: boolean, ): Promise<void>; }; vi.spyOn(internals, "catchUpBackpressuredClient").mockResolvedValue(); internals.sessions.set(activeSessionId, state); if (replayFailed) { vi.spyOn( internals, preparation === "replacement-replay-failure" ? "flushDeferredSessionFrames" : "write", ).mockImplementationOnce(() => { throw new Error("replacement delivery failed"); }); }
let releaseChunk: () => void = () => {}; const chunkGate = new Promise<void>((resolve) => { releaseChunk = resolve; }); async function* gatedTranscript(fail = false): AsyncGenerator<Buffer> { yield Buffer.from(`{"type":"session_snapshot_chunk"}`); await chunkGate; if (fail) throw new Error("last snapshot failed"); yield Buffer.from(`{"type":"session_snapshot_chunk"}`); } const transcript = gatedTranscript();
let stream: Promise<unknown>; if (preparation !== "stream") { vi.spyOn(internals, "createAttachResult").mockImplementation(async () => { const result = streamedResult(); await chunkGate; if (preparation.startsWith("replacement")) throw new Error("snapshot preparation failed"); return result; }); if (preparation.endsWith("catchup")) { client.deferredSessionFramesDropped = new Set([activeSessionId]); client.catchupActiveSessionIds?.add(activeSessionId); if (preparation === "inline-catchup") client.catchupPurposes = new Map([[activeSessionId, "replacement"]]); stream = internals.drainBackpressuredClientCatchups(client); } else { internals.broadcastToSession(state, { type: "session_replaced", activeSessionId, state: streamedResult().snapshot.state, messages: [], }); stream = chunkGate.then(nextMacroTaskTurn); } } else { stream = internals.streamWorkerSnapshot( client, streamedResult(), transcript, "attach", markClientSnapshotStreaming(client, activeSessionId), true, ); } await nextMacroTaskTurn(); const streams = [stream]; for (let index = 1; index < streamCount; index++) { streams.push( internals.streamWorkerSnapshot( client, streamedResult(), gatedTranscript(failFinal), "attach", markClientSnapshotStreaming(client, activeSessionId), true, ), ); } // The stream is parked between chunks; events broadcast now are withheld. internals.broadcastToSession(state, sessionEventMessage(2)); internals.broadcastToSession(state, sessionEventMessage(3)); if (preparation === "inline-catchup") state.lastEventSequence = 3; if (preparation !== "stream") { expect(written).toHaveLength(0); expect(client.deferredSessionOutbounds?.get(activeSessionId)?.frames).toHaveLength(2); } else expect(written.length).toBeGreaterThan(0);
releaseChunk(); const outcomes = await Promise.allSettled(streams); expect(outcomes.filter((outcome) => outcome.status === "rejected")).toHaveLength(failFinal ? 1 : 0);
const decoder = new PrivateFrameDecoder(isDaemonWorkerFrameHeader); const frames = decoder.push(Buffer.concat(written)); if (preparation === "inline-catchup") expect(JSON.parse(frames[0]!.payload.toString()).meta.sequence).toBe(1); const outboundTypes = frames.map((frame) => (frame.header.kind === "outbound" ? frame.header.outboundType : "")); expect(outboundTypes).toEqual([ ...(preparation.startsWith("replacement") ? preparation === "replacement-write-failure" ? [] : ["session_replaced"] : preparation === "catchup" ? ["session_snapshot_begin", "session_snapshot_end"] : preparation === "inline-catchup" ? ["session_replaced"] : Array.from({ length: streamCount }, (_, index) => failFinal && index === streamCount - 1 ? ["session_snapshot_begin", "session_snapshot_chunk", "session_snapshot_failed"] : [ "session_snapshot_begin", "session_snapshot_chunk", "session_snapshot_chunk", "session_snapshot_end", ], ).flat()), ...(failFinal || replayFailed ? [] : ["session_event", "session_event"]), ]); const replayed = frames .filter((frame) => frame.header.kind === "outbound" && frame.header.outboundType === "session_event") .map((frame) => JSON.parse(frame.payload.toString("utf8"))); expect(replayed).toEqual( failFinal || replayFailed ? [] : [ expect.objectContaining({ event: { type: "session_info_changed", name: "change-2" } }), expect.objectContaining({ event: { type: "session_info_changed", name: "change-3" } }), ], ); expect(client.catchupActiveSessionIds?.size ?? 0).toBe(failFinal || replayFailed ? 1 : 0); if (replayFailed) { const before = written.length; internals.broadcastToSession(state, sessionEventMessage(4)); expect(written).toHaveLength(before); expect(client.catchupPurposes?.get(activeSessionId)).toBe("replacement"); expect(internals.catchUpBackpressuredClient).toHaveBeenCalledOnce(); } expect(client.snapshotActiveSessionCounts?.size ?? 0).toBe(0); expect(client.deferredSessionOutbounds?.size ?? 0).toBe(0); socket.destroy(); });
it.each(["jsonl", "private-framed"] as const)( "worker: freezes mutable streaming frames for %s replay", (transport) => { const daemon = new AgentDaemon(join(tmpdir(), "deferred-mutable-frames.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, createRuntime: async () => { throw new Error("unexpected runtime creation"); }, }); const socket = new PassThrough(); const written: Buffer[] = []; socket.on("data", (chunk: Buffer) => written.push(Buffer.from(chunk))); const client = socketClient(socket, { transport, snapshotActiveSessionIds: new Set([activeSessionId]) }); const state = { activeSessionId, clients: new Set([client]), eventGeneration: "generation-deferred", lastEventSequence: 1, } as ActiveSessionState; const internals = daemon as unknown as { broadcastToSession(state: ActiveSessionState, message: DaemonOutbound): void; flushDeferredSessionFrames(client: DaemonSocketClient, activeSessionId: string): void; }; try { const message = fauxAssistantMessage("first"); const text = message.content[0]; if (text?.type !== "text") throw new Error("Expected a text block"); internals.broadcastToSession(state, { type: "session_event", activeSessionId, event: { type: "message_start", message: { ...message } }, }); const toolCall = { type: "toolCall" as const, id: "tool", name: "lookup", arguments: { query: "first" } }; message.content.push(toolCall); for (const query of ["second", "third"]) { text.text = query; toolCall.arguments.query = query; internals.broadcastToSession(state, { type: "session_event", activeSessionId, event: { type: "message_update", message: { ...message }, assistantMessageEvent: { type: "toolcall_delta", contentIndex: 1, delta: query, partial: message }, }, }); } const bufferedBytes = client.deferredSessionOutbounds?.get(activeSessionId)?.bytes; text.text = "future".repeat(1000); toolCall.arguments.query = "future"; expect(written).toHaveLength(0); internals.flushDeferredSessionFrames(client, activeSessionId); const payloads = transport === "private-framed" ? new PrivateFrameDecoder(isDaemonWorkerFrameHeader) .push(Buffer.concat(written)) .map((frame) => frame.payload) : written; const replayed = payloads.map((payload) => JSON.parse(payload.toString())); expect(replayed[0]).toMatchObject({ event: { message: { content: [{ type: "text", text: "first" }] } } }); if (transport === "private-framed") { expect(replayed.slice(1)).toMatchObject([ { type: "assistant_stream_delta", toolCallArguments: { query: "second" } }, { type: "assistant_stream_delta", toolCallArguments: { query: "third" } }, ]); } else { expect(Buffer.concat(written).length).toBe(bufferedBytes); for (const [index, query] of ["second", "third"].entries()) { expect(replayed[index + 1]).toMatchObject({ event: { message: { content: [ { type: "text", text: query }, { type: "toolCall", arguments: { query } }, ], }, }, }); } } } finally { socket.destroy(); } }, );
it.each([ ["count", 257, 0], ["single payload", 1, 2 * 1024 * 1024], ["total bytes", 2, 1024 * 1024], ] as const)("worker: catches up when the deferral %s limit overflows", async (_limit, count, textLength) => { const daemon = new AgentDaemon(join(tmpdir(), "deferred-frames-overflow.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, createRuntime: async () => { throw new Error("unexpected runtime creation"); }, }); const socket = new PassThrough(); socket.on("error", () => {}); const written: Buffer[] = []; socket.on("data", (chunk: Buffer) => written.push(Buffer.from(chunk))); const client = socketClient(socket, { transport: "private-framed" }); const state = { activeSessionId, clients: new Set([client]), eventGeneration: "generation-deferred", lastEventSequence: 1, runtime: { metadata: { kind: "top-level" as const, createdAt: 1 } }, } as unknown as ActiveSessionState; const internals = daemon as unknown as { sessions: Map<string, ActiveSessionState>; broadcastToSession(state: ActiveSessionState, message: DaemonOutbound): void; streamWorkerSnapshot( client: DaemonSocketClient, result: DaemonAttachResult, transcript: AsyncIterable<Buffer> & { dispose?(): void }, purpose: "attach", signal: AbortSignal, snapshotAlreadyMarked: boolean, ): Promise<void>; queueClientCatchup( client: DaemonSocketClient, activeSessionId: string, purpose: "replacement" | "resync", ): void; }; const queueClientCatchup = vi.fn(); internals.queueClientCatchup = queueClientCatchup; internals.sessions.set(activeSessionId, state); const snapshotSignal = markClientSnapshotStreaming(client, activeSessionId);
let releaseChunk: () => void = () => {}; const chunkGate = new Promise<void>((resolve) => { releaseChunk = resolve; }); async function* gatedTranscript(): AsyncGenerator<Buffer> { yield Buffer.from(`{"type":"session_snapshot_chunk"}`); await chunkGate; yield Buffer.from(`{"type":"session_snapshot_chunk"}`); } const transcript = gatedTranscript();
const stream = internals.streamWorkerSnapshot( client, streamedResult(), transcript, "attach", snapshotSignal, true, ); await nextMacroTaskTurn(); for (let index = 0; index < count; index++) { const outbound = sessionEventMessage(index + 2); if (textLength && outbound.type === "session_event") { const message = fauxAssistantMessage("é".repeat(textLength)); outbound.event = { type: "message_update", message, assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: "é", partial: message }, }; } internals.broadcastToSession(state, outbound); if (index < count - 1) expect(client.deferredSessionOutbounds?.get(activeSessionId)?.frames).toHaveLength(index + 1); } expect(client.deferredSessionOutbounds?.size ?? 0).toBe(0); // Once overflow requests a snapshot, later events cannot restart a partial replay. internals.broadcastToSession(state, sessionEventMessage(count + 2)); expect(client.deferredSessionOutbounds?.size ?? 0).toBe(0); releaseChunk(); await stream;
const decoder = new PrivateFrameDecoder(isDaemonWorkerFrameHeader); const frames = decoder.push(Buffer.concat(written)); const outboundTypes = frames.map((frame) => (frame.header.kind === "outbound" ? frame.header.outboundType : "")); expect(outboundTypes).toEqual([ "session_snapshot_begin", "session_snapshot_chunk", "session_snapshot_chunk", "session_snapshot_end", ]); expect(queueClientCatchup).toHaveBeenCalledWith(client, activeSessionId, "resync"); socket.destroy(); });
it.each(["resync", "replacement"] as const)( "supervisor: failed %s does not block healthy sessions", async (purpose) => { const supervisor = new DaemonSupervisor(join(tmpdir(), "catchup-fairness.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, descriptorDir: join(tmpdir(), "catchup-fairness-state"), }); const socket = new PassThrough(); const written: DaemonOutbound[] = []; socket.on("data", (chunk: Buffer) => written.push(JSON.parse(chunk.toString()))); const sessions = [activeSessionId, "also-failing", "healthy"]; const client = socketClient(socket, { capabilities: new Set(), attachedActiveSessionIds: new Set(sessions), catchupActiveSessionIds: new Set(sessions), catchupPurposes: new Map(sessions.map((id) => [id, purpose])), }); const internals = supervisor as unknown as { attachClient( client: DaemonSocketClient, command: { activeSessionId: string }, ): Promise<{ result: DaemonAttachResult }>; catchUpClient(client: DaemonSocketClient): Promise<void>; }; const attach = vi.spyOn(internals, "attachClient").mockImplementation(async (_client, command) => { if (command.activeSessionId !== "healthy") throw new Error("worker remains unavailable"); const result = streamedResult(); result.activeSessionId = result.snapshot.activeSessionId = result.snapshot.state.activeSessionId = "healthy"; return { result }; }); try { await internals.catchUpClient(client); expect(attach.mock.calls.map(([, command]) => command.activeSessionId)).toEqual(sessions); expect(written).toEqual([ expect.objectContaining({ type: purpose === "replacement" ? "session_replaced" : "session_resynced", activeSessionId: "healthy", }), ]); expect(client.catchupActiveSessionIds).toEqual(new Set(sessions.slice(0, 2))); expect(client.catchupRetryTimer).toBeDefined(); client.catchupActiveSessionIds?.add("healthy"); client.catchupPurposes?.set("healthy", purpose); await vi.waitFor(() => expect(written).toHaveLength(2)); expect(attach.mock.calls.map(([, command]) => command.activeSessionId)).toEqual([...sessions, ...sessions]); expect(written[1]).toMatchObject({ type: written[0]?.type, activeSessionId: "healthy" }); expect(client.catchupPurposes).toEqual(new Map(sessions.slice(0, 2).map((id) => [id, purpose]))); expect(client.snapshotActiveSessionCounts?.size ?? 0).toBe(0); } finally { clearTimeout(client.catchupRetryTimer); socket.destroy(); } }, );
const supervisorScenarios = [ "attached", "catchup", "inline-catchup", "failed-catchup", "failed-replacement-catchup", "failed-inline-catchup", "inline-reattach", "failed-reattach", "failed-new-reattach", "detached", "replaced", "failed", "failed-detached", "oversized", "byte-limit", "count-limit", "closed", "oversized-closed", "replaced-closed", "pending-closed", "backpressured-closed", ]; it.each(supervisorScenarios)("supervisor: replay (%s)", async (scenario) => { const overflow = scenario.startsWith("oversized") || scenario.endsWith("-limit"); const closed = scenario.endsWith("closed"); const reattach = scenario.endsWith("reattach"); const catchup = scenario.endsWith("catchup"); const failedCatchup = catchup && scenario.startsWith("failed"); const purpose = scenario.includes("replacement") ? "replacement" : "resync"; const supervisor = new DaemonSupervisor(join(tmpdir(), "deferred-frames-supervisor.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, descriptorDir: join(tmpdir(), "deferred-frames-supervisor-state"), }); const worker: SupervisorWorkerHarness = { descriptor: { workerId: "worker-deferred", lifecycle: "ready", pid: 987_654 }, client: { close: vi.fn(), request: vi.fn(async () => ({ success: true })) }, summaries: new Map([[activeSessionId, summary()]]), snapshotCache: new Map<string, DaemonAttachResult>(), transcriptCaches: new Map<string, SnapshotTranscriptCache>(), snapshotGenerations: new Map<string, Map<string, unknown>>(), snapshotLoads: new Map<string, Promise<DaemonAttachResult>>(), intentionalStop: false, stopRevision: 0, }; const socket = new PassThrough(); socket.on("error", () => {}); const written: string[] = []; socket.on("data", (chunk: Buffer) => { written.push(chunk.toString("utf8")); if (scenario === "backpressured-closed" && written.length === 1) socket.pause(); }); const client = socketClient(socket); const internals = supervisor as unknown as { clients: Set<DaemonSocketClient>; workers: Map<string, SupervisorWorkerHarness>; handleCommand(client: DaemonSocketClient, command: DaemonCommand): Promise<DaemonResponse | undefined>; findWorkerForClient(): Promise<{ worker: SupervisorWorkerHarness; summary: SessionSummary }>; attachClient(): Promise<{ worker: SupervisorWorkerHarness; result: DaemonAttachResult; transcript?: SnapshotTranscriptCache; }>; drainClientCatchups(client: DaemonSocketClient): Promise<void>; catchUpClient(client: DaemonSocketClient): Promise<void>; streamSnapshot( client: DaemonSocketClient, worker: SupervisorWorkerHarness, result: DaemonAttachResult, transcript: SnapshotTranscriptCache, purpose: "attach", ): Promise<void>; handleWorkerFrame(worker: SupervisorWorkerHarness, frame: PrivateFrame<DaemonWorkerFrameHeader>): void; }; if (!failedCatchup) vi.spyOn(internals, "catchUpClient").mockResolvedValue(); internals.clients.add(client); internals.workers.set(worker.descriptor.workerId, worker);
const messages: AgentMessage[] = [ { role: "user", content: scenario === "backpressured-closed" ? "x".repeat(128 * 1024) : "stable", timestamp: 1, }, ]; const transcript = new SnapshotTranscriptCache({ activeSessionId, snapshotId, messages: scenario === "pending-closed" ? undefined : messages, cacheRoot: tmpdir(), }); let releaseChunk: () => void = () => {}; const chunkGate = new Promise<void>((resolve) => { releaseChunk = resolve; }); const waitForChunk = transcript.waitForChunk.bind(transcript); transcript.waitForChunk = async (index: number, signal?: AbortSignal) => { if (index === 1) { await chunkGate; if (scenario.startsWith("failed")) throw new Error("snapshot failed"); } return waitForChunk(index, signal); };
let stream: Promise<unknown>; if (reattach || catchup) { if (reattach || scenario.includes("inline")) client.capabilities.clear(); if (scenario === "failed-new-reattach") client.attachedActiveSessionIds.delete(activeSessionId); vi.spyOn(internals, "findWorkerForClient").mockResolvedValue({ worker, summary: summary() }); vi.spyOn(internals, "attachClient").mockImplementation(async () => { const result = { ...streamedResult(), snapshotStream: undefined }; await chunkGate; if (scenario.startsWith("failed")) throw new Error("reattach failed"); return { worker, result, transcript }; }); if (catchup) { client.deferredSessionPayloadsDropped = new Set([activeSessionId]); client.catchupActiveSessionIds?.add(activeSessionId); client.catchupPurposes = new Map([[activeSessionId, purpose]]); stream = failedCatchup ? internals.catchUpClient(client) : internals.drainClientCatchups(client); } else { client.attachedActiveSessionIds.add("old-session"); stream = internals.handleCommand(client, { type: "reattach", activeSessionId: "old-session", targetActiveSessionId: activeSessionId, }); } } else { stream = internals.streamSnapshot(client, worker, streamedResult(), transcript, "attach"); } const completion = stream.catch((error: unknown) => error); await nextMacroTaskTurn(); // The stream is parked before the end record; relayed frames are withheld. if (scenario.startsWith("replaced")) { internals.handleWorkerFrame(worker, { header: { kind: "outbound", outboundType: "session_replaced", activeSessionId, payloadEncoding: "jsonl", }, payload: Buffer.from( `${JSON.stringify({ type: "session_replaced", activeSessionId, state: streamedResult().snapshot.state, messages: [] })}\n`, ), }); } const eventFrame: PrivateFrame<DaemonWorkerFrameHeader> = { header: { kind: "outbound", outboundType: "session_event", activeSessionId, sessionEventType: "session_info_changed", payloadEncoding: "jsonl", }, payload: Buffer.from( `${JSON.stringify({ ...sessionEventMessage(2), padding: "x".repeat( scenario.startsWith("oversized") ? 8 * 1024 * 1024 : scenario === "byte-limit" ? 4 * 1024 * 1024 : 0, ), })}\n`, ), }; internals.handleWorkerFrame(worker, eventFrame); if (scenario === "byte-limit") internals.handleWorkerFrame(worker, eventFrame); if (scenario === "count-limit") { for (let index = 1; index < 257; index++) internals.handleWorkerFrame(worker, eventFrame); } if (overflow) { expect(client.deferredSessionPayloads?.size ?? 0).toBe(0); expect(client.catchupActiveSessionIds?.has(activeSessionId)).toBe(true); } internals.handleWorkerFrame(worker, { header: { kind: "outbound", outboundType: "session_event", activeSessionId, sessionEventType: "session_info_changed", payloadEncoding: "jsonl", }, payload: Buffer.from(`${JSON.stringify(sessionEventMessage(3))}\n`), });
const closeFrame: PrivateFrame<DaemonWorkerFrameHeader> = { header: { kind: "outbound", outboundType: "session_closed", activeSessionId, payloadEncoding: "jsonl" }, payload: Buffer.from(`${JSON.stringify({ type: "session_closed", activeSessionId, reason: "killed" })}\n`), }; if (catchup) expect(client.deferredSessionPayloads?.get(activeSessionId)?.payloads).toHaveLength(2); if (closed && !overflow) internals.handleWorkerFrame(worker, closeFrame); if (scenario.endsWith("detached")) client.attachedActiveSessionIds.delete(activeSessionId); releaseChunk(); const error = await completion; if (scenario === "backpressured-closed") { socket.resume(); await nextMacroTaskTurn(); } if (closed && overflow) internals.handleWorkerFrame(worker, closeFrame); if (scenario.startsWith("failed") && !catchup) expect(error).toBeInstanceOf(Error); else expect(error).toBeUndefined();
const lines = written.join("").split("\n").filter(Boolean); const parsed = lines.map((line) => JSON.parse(line) as { type: string }); expect(parsed.map((entry) => entry.type)).toEqual( failedCatchup ? [] : scenario === "inline-catchup" ? ["session_resynced", "session_event", "session_event"] : reattach ? [ ...(scenario === "inline-reattach" ? ["response", "session_detached"] : []), ...(scenario === "failed-new-reattach" ? ["session_detached"] : ["session_event", "session_event"]), ] : [ "session_snapshot_begin", ...(scenario === "pending-closed" ? [] : ["session_snapshot_chunk"]), ...(closed && !overflow ? ["session_closed"] : []), ...(closed && !overflow ? [] : [scenario.startsWith("failed") ? "session_snapshot_failed" : "session_snapshot_end"]), ...(closed && overflow ? ["session_closed"] : []), ...(scenario === "attached" || scenario === "catchup" ? ["session_event", "session_event"] : []), ], ); expect(client.deferredSessionPayloads?.size ?? 0).toBe(0); expect(client.snapshotActiveSessionCounts?.size ?? 0).toBe(0); expect(client.snapshotTransferAbortControllers?.size ?? 0).toBe(0); expect(client.catchupActiveSessionIds?.size ?? 0).toBe( !closed && (scenario === "replaced" || scenario === "failed" || failedCatchup || overflow) ? 1 : 0, ); if (failedCatchup) { expect(client.catchupPurposes?.get(activeSessionId)).toBe(purpose); expect(client.catchupRetryTimer).toBeDefined(); internals.handleWorkerFrame(worker, eventFrame); await internals.catchUpClient(client); expect(written).toHaveLength(0); expect(internals.attachClient).toHaveBeenCalledOnce(); transcript.waitForChunk = waitForChunk; const result = streamedResult(); result.lastEventSequence = result.snapshot.lastEventSequence = 3; result.lastEventCursor = result.snapshot.lastEventCursor = { generation: "generation-deferred", sequence: 3 }; vi.mocked(internals.attachClient).mockResolvedValue({ worker, result, transcript }); const expected = client.capabilities.has("chunked_snapshot") ? [ ...(purpose === "replacement" ? ["session_replaced"] : []), "session_snapshot_begin", "session_snapshot_chunk", "session_snapshot_end", ] : ["session_resynced"]; await vi.waitFor(() => expect(written.map((line) => JSON.parse(line).type)).toEqual(expected)); expect(client.catchupActiveSessionIds?.size).toBe(0); expect(client.catchupRetryTimer).toBeUndefined(); expect(client.snapshotActiveSessionCounts?.size ?? 0).toBe(0); internals.handleWorkerFrame(worker, { ...eventFrame, payload: Buffer.from(`${JSON.stringify(sessionEventMessage(4))}\n`), }); expect(JSON.parse(written.at(-1)!).type).toBe("session_event"); } if (!closed && (scenario === "replaced" || overflow)) { const before = written.length; const frame: PrivateFrame<DaemonWorkerFrameHeader> = { header: { kind: "outbound", outboundType: "session_event", activeSessionId, payloadEncoding: "jsonl" }, payload: Buffer.from(`${JSON.stringify(sessionEventMessage(4))}\n`), }; internals.handleWorkerFrame(worker, frame); expect(written).toHaveLength(before); const replacement = internals.streamSnapshot(client, worker, streamedResult(), transcript, "attach"); internals.handleWorkerFrame(worker, frame); await replacement; expect(written.map((line) => JSON.parse(line).type).slice(before)).toEqual([ "session_snapshot_begin", "session_snapshot_chunk", "session_snapshot_end", "session_event", ]); } transcript.dispose(); socket.destroy(); });
it.each(["backpressure", "closed", "detached", "replaced"] as const)("worker: stop after %s", async (scenario) => { const daemon = new AgentDaemon(join(tmpdir(), "deferred-frames-stop.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, createRuntime: async () => { throw new Error("unexpected runtime creation"); }, }); const socket = new PassThrough(); const client = socketClient(socket); const written: DaemonOutbound[] = []; socket.on("data", (chunk: Buffer) => { const message = JSON.parse(chunk.toString("utf8")) as DaemonOutbound; written.push(message); if (scenario === "backpressure" && message.type === "session_snapshot_end") socket.pause(); }); const write = vi.spyOn(socket, "write"); const state = { activeSessionId, clients: new Set([client]), eventGeneration: "generation-deferred", lastEventSequence: 1, runtime: { metadata: { kind: "top-level", createdAt: 1 } }, } as unknown as ActiveSessionState; const internals = daemon as unknown as { sessions: Map<string, ActiveSessionState>; catchUpBackpressuredClient(client: DaemonSocketClient): Promise<void>; broadcastToSession(state: ActiveSessionState, message: DaemonOutbound): void; streamWorkerSnapshot( client: DaemonSocketClient, result: DaemonAttachResult, transcript: AsyncIterable<Buffer>, purpose: "attach", signal: AbortSignal, marked: boolean, ): Promise<void>; }; vi.spyOn(internals, "catchUpBackpressuredClient").mockResolvedValue(); internals.sessions.set(activeSessionId, state); let release!: () => void; const gate = new Promise<void>((resolve) => { release = resolve; }); async function* transcript(): AsyncGenerator<Buffer> { await gate; yield Buffer.from('{"type":"session_snapshot_chunk"}\n'); } const stream = internals.streamWorkerSnapshot( client, streamedResult(), transcript(), "attach", markClientSnapshotStreaming(client, activeSessionId), true, ); try { await nextMacroTaskTurn(); internals.broadcastToSession(state, { ...sessionEventMessage(2), type: "session_event", activeSessionId, event: { type: "session_info_changed", name: "x".repeat(128 * 1024) }, }); internals.broadcastToSession(state, sessionEventMessage(3)); if (scenario === "closed") internals.broadcastToSession(state, { type: "session_closed", activeSessionId, reason: "killed" }); if (scenario === "detached") client.attachedActiveSessionIds.delete(activeSessionId); if (scenario === "replaced") { internals.broadcastToSession(state, { type: "session_replaced", activeSessionId, state: streamedResult().snapshot.state, messages: [], }); internals.broadcastToSession(state, sessionEventMessage(4)); } release(); await stream; const eventWrites = write.mock.calls.filter(([chunk]) => JSON.parse(String(chunk)).type === "session_event"); expect(eventWrites).toHaveLength(scenario === "backpressure" ? 1 : 0); expect(client.deferredSessionOutbounds?.size ?? 0).toBe(0); if (scenario === "backpressure") { expect(socket.writableNeedDrain).toBe(true); expect(client.backpressured).toBe(true); expect(client.catchupActiveSessionIds?.has(activeSessionId)).toBe(true); internals.broadcastToSession(state, sessionEventMessage(4)); expect( write.mock.calls.filter(([chunk]) => JSON.parse(String(chunk)).type === "session_event"), ).toHaveLength(1); } if (scenario === "closed") expect(written.some((message) => message.type === "session_closed")).toBe(true); if (scenario === "replaced") { internals.broadcastToSession(state, sessionEventMessage(5)); expect(written.some((message) => message.type === "session_event")).toBe(false); expect(client.catchupPurposes?.get(activeSessionId)).toBe("replacement"); const replacement = internals.streamWorkerSnapshot( client, streamedResult(), transcript(), "attach", markClientSnapshotStreaming(client, activeSessionId), true, ); internals.broadcastToSession(state, sessionEventMessage(6)); await replacement; expect(written.filter((message) => message.type === "session_event")).toEqual([sessionEventMessage(6)]); } } finally { release(); socket.destroy(); await stream; } });
it("supervisor: preserves replay backpressure when another session's snapshot finishes", () => { const supervisor = new DaemonSupervisor(join(tmpdir(), "deferred-backpressure-supervisor.sock"), { defaultSessionConfig: { agentDir: "/tmp", cwd: "/tmp" }, descriptorDir: join(tmpdir(), "deferred-backpressure-state"), }); const socket = new PassThrough(); const client = socketClient(socket, { deferredSessionPayloads: new Map([ [activeSessionId, { payloads: [Buffer.alloc(128 * 1024), Buffer.from("later")], bytes: 128 * 1024 + 5 }], ]), }); const internals = supervisor as unknown as { reserveSnapshotStream(client: DaemonSocketClient, activeSessionId: string): () => void; flushDeferredSessionPayloads(client: DaemonSocketClient, activeSessionId: string): void; }; try { const finishOtherSnapshot = internals.reserveSnapshotStream(client, "other-session"); const write = vi.spyOn(socket, "write"); internals.flushDeferredSessionPayloads(client, activeSessionId); finishOtherSnapshot(); expect(write).toHaveBeenCalledOnce(); expect(client.backpressured).toBe(true); expect(client.catchupActiveSessionIds?.has(activeSessionId)).toBe(true); expect(client.deferredSessionPayloads?.size ?? 0).toBe(0); } finally { socket.destroy(); } });});