import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; 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, getModel } from "@earendil-works/pi-ai"; import { describe, expect, it, vi } from "vitest"; import { createManualContinuationMessage, MANUAL_CONTINUATION_PROMPT } from "../src/core/messages.js"; import { DaemonAgentConnection } from "../src/modes/agent-connection/daemon-agent-connection.js"; import type { AgentConnectionEvent, AgentConnectionRlmChildAgentSnapshot, AgentConnectionSessionTreeEntry, AgentConnectionState, } from "../src/modes/agent-connection/types.js"; import { DaemonCapabilityUnavailableError, type DaemonClientCloseListener, type DaemonClientMessageListener, type DaemonClientRequestOptions, type DaemonHello, DaemonSocketClosedError, type DaemonTransportClient, } from "../src/modes/daemon/daemon-client.js"; import { DaemonSessionStartingError } from "../src/modes/daemon/daemon-errors.js"; import { DAEMON_PROTOCOL_INFO, DAEMON_SCHEMA_REVISION, type DaemonAttachResult, type DaemonCommand, type DaemonOutbound, type DaemonResponse, } from "../src/modes/daemon/daemon-protocol.js"; import { DaemonRoutedClient } from "../src/modes/daemon/daemon-routed-client.js"; import type { DaemonWorkerClient } from "../src/modes/daemon/daemon-worker-client.js"; import { InteractiveMode } from "../src/modes/interactive/interactive-mode.js"; import { attachJsonlLineReader, serializeJsonLine } from "../src/modes/rpc/jsonl.js"; class FakeDaemonClient { readonly requests: DaemonCommand[] = []; readonly requestTimeouts: number[] = []; attachResultFactory: ((command: Extract) => DaemonAttachResult) | undefined; restoredAttachGate: Promise | undefined; restoredAttachCompleted = 0; closeCount = 0; emitCloseOnClose = false; connected = true; reconnectCount = 0; resetTransportCount = 0; reconnectError: Error | undefined; attachFailures = 0; attachError: Error | undefined; connectionStateGate: Promise | undefined; connectionStateFactory: ((activeSessionId: string) => AgentConnectionState) | undefined; rlmChildren: AgentConnectionRlmChildAgentSnapshot[] = []; rlmChildrenEventSequence = 12; rlmChildrenGate: Promise | undefined; abortBashUnknownCommand = false; abortAndClearQueueUnknownCommand = false; inputPauseAcquireGate: Promise | undefined; cronAddGate: Promise | undefined; promptGate: Promise | undefined; promptError: Error | undefined; promptResponseError: string | undefined; cancelPromptAdmissionStatus: "cancelled" | "owned" | "unknown" = "owned"; serverCapabilities = new Set(); updateRestartSessions: Array> = []; createdActiveSessionId = "active-created"; createError: Error | undefined; createGate: Promise | undefined; hello: DaemonHello | undefined = { type: "daemon_hello", socketPath: "/tmp/fake.sock", protocol: DAEMON_PROTOCOL_INFO, schemaRevision: DAEMON_SCHEMA_REVISION, clientId: "fake-client", serverCapabilities: ["prompt_admission_cancellation", "session_input_admission"], }; private readonly messageListeners = new Set(); private readonly closeListeners = new Set(); async request( command: DaemonCommand, timeoutMs = 30000, options: DaemonClientRequestOptions = {}, ): Promise { this.requests.push(command); this.requestTimeouts.push(timeoutMs); switch (command.type) { case "prompt": if (this.promptGate) await this.promptGate; if (this.promptError) throw this.promptError; if (this.promptResponseError) { return { type: "response", command: command.type, success: false, error: this.promptResponseError }; } return { type: "response", command: command.type, success: true }; case "prompt_and_wait": if (this.promptGate) await this.promptGate; if (this.promptError) throw this.promptError; return { type: "response", command: command.type, success: true }; case "cancel_prompt_admission": return { type: "response", command: command.type, success: true, data: { status: this.cancelPromptAdmissionStatus }, }; case "list": return { type: "response", command: command.type, success: true, data: { sessions: this.updateRestartSessions }, }; case "create": { if (this.createGate) await this.createGate; if (this.createError) throw this.createError; return { type: "response", command: command.type, success: true, data: { id: this.createdActiveSessionId, activeSessionId: this.createdActiveSessionId }, }; } case "attach": { if (this.attachFailures > 0) { this.attachFailures--; throw this.attachError ?? new Error("attach failed"); } if (command.activeSessionId === "active-restored" && this.restoredAttachGate) { await this.restoredAttachGate; this.restoredAttachCompleted++; } if (command.activeSessionId === "missing") { return { type: "response", command: command.type, success: false, error: "Unknown active session: missing", }; } const response: DaemonResponse = { type: "response", command: command.type, success: true, data: this.attachResultFactory?.(command) ?? createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }; options.onResponse?.(response); return response; } case "get_queue": return { type: "response", command: command.type, success: true, data: { steering: ["steer"], followUp: ["follow"] }, }; case "get_connection_state": await this.connectionStateGate; return { type: "response", command: command.type, success: true, data: this.connectionStateFactory?.(command.activeSessionId) ?? createConnectionState(command.activeSessionId, "session-current"), }; case "get_messages": return { type: "response", command: command.type, success: true, data: { messages: [{ role: "user", content: "current prompt", timestamp: 4 }] }, }; case "get_rlm_children": await this.rlmChildrenGate; return { type: "response", command: command.type, success: true, data: { children: this.rlmChildren, eventSequence: this.rlmChildrenEventSequence }, }; case "get_commands": return { type: "response", command: command.type, success: true, data: { commands: [ { name: "route", description: "Select a routing policy", argumentCompletions: [{ value: "policy ", label: "policy", description: "Follow a policy" }], source: "extension", sourceInfo: { path: "/tmp/extensions/route.ts", source: "local", scope: "user", origin: "top-level", }, }, ], }, }; case "get_resource_snapshot": return { type: "response", command: command.type, success: true, data: { contextFiles: [{ path: "/tmp/AGENTS.md" }], skills: [ { name: "demo-skill", description: "Demo skill", filePath: "/tmp/skills/demo-skill/SKILL.md", sourceInfo: { path: "/tmp/skills/demo-skill/SKILL.md", source: "local", scope: "project", origin: "top-level", baseDir: "/tmp/skills", }, }, ], prompts: [], extensions: [], themes: [], diagnostics: { skills: [], prompts: [], extensions: [], themes: [], }, }, }; case "replace_acp_mcp_servers": return { type: "response", command: command.type, success: true }; case "get_model_catalog": return { type: "response", command: command.type, success: true, data: { models: [getModel("openai", "gpt-5.1")], configuredProviders: ["openai"], }, }; case "get_available_models": return { type: "response", command: command.type, success: true, data: { models: [getModel("openai", "gpt-5.1")] }, }; case "get_session_context": return { type: "response", command: command.type, success: true, data: { context: { messages: [{ role: "user", content: "context prompt", timestamp: 3 }], thinkingLevel: "medium", model: { provider: "anthropic", modelId: "claude-sonnet-4-5" }, }, }, }; case "get_session_tree": return { type: "response", command: command.type, success: true, data: command.projection === "tree" ? { entries: [{ entry: projectedUserEntry("user-1", null) }], leafId: "user-1" } : { flatNodes: [ { entry: { type: "message", id: "user-1", parentId: null, timestamp: "2026-01-01T00:00:00.000Z", message: { role: "user", content: "hello", timestamp: 1 }, }, }, ], leafId: "user-1", }, }; case "get_tool_definition": return { type: "response", command: command.type, success: true, data: { toolDefinition: { name: command.name, label: command.name, description: `${command.name} description`, promptSnippet: `${command.name} prompt`, promptGuidelines: [`Use ${command.name}`], parameters: { type: "object" }, renderShell: "self", }, }, }; case "clear_queue": return { type: "response", command: command.type, success: true, data: { steering: ["cleared"], followUp: [] }, }; case "abort_and_clear_queue": if (this.abortAndClearQueueUnknownCommand) { return { type: "response", command: command.type, success: false, error: "Unknown daemon command: abort_and_clear_queue", }; } return { type: "response", command: command.type, success: true, data: { steering: ["aborted"], followUp: ["cleared"] }, }; case "acquire_session_input_pause": if (this.inputPauseAcquireGate) await this.inputPauseAcquireGate; return { type: "response", command: command.type, success: true, data: { pauseId: "pause-1" }, }; case "release_session_input_pause": return { type: "response", command: command.type, success: true }; case "heartbeats_list": return this.serverCapabilities.has("heartbeat_catalog") ? { type: "response", command: command.type, success: true, data: { heartbeats: [] } } : { type: "response", command: command.type, success: false, error: "Unknown daemon command: heartbeats_list", }; case "heartbeat_manage": return this.serverCapabilities.has("heartbeat_management") ? { type: "response", command: command.type, success: true, data: { heartbeat: { id: command.jobId } }, } : { type: "response", command: command.type, success: false, error: "Unknown daemon command: heartbeat_manage", }; case "list_saved_sessions": { const activeSessionId = "activeSessionId" in command ? command.activeSessionId : undefined; options.onProgress?.({ id: "daemon_test", type: "session_list_progress", command: "list_saved_sessions", ...(activeSessionId ? { activeSessionId } : {}), loaded: 1, total: 2, }); options.onProgress?.({ id: "daemon_test", type: "session_list_item", command: "list_saved_sessions", ...(activeSessionId ? { activeSessionId } : {}), session: { path: "/tmp/session-a.jsonl", id: "session-a", cwd: "/tmp", name: "Saved session", created: "2026-01-01T00:00:00.000Z", modified: "2026-01-02T00:00:00.000Z", messageCount: 2, firstMessage: "hello", allMessagesText: "hello world", }, }); options.onProgress?.({ id: "daemon_test", type: "session_list_progress", command: "list_saved_sessions", ...(activeSessionId ? { activeSessionId } : {}), loaded: 2, total: 2, }); const sessions = [ { path: "/tmp/session-a.jsonl", id: "session-a", cwd: "/tmp", name: "Saved session", created: "2026-01-01T00:00:00.000Z", modified: "2026-01-02T00:00:00.000Z", messageCount: 2, firstMessage: "hello", allMessagesText: "hello world", }, ]; return { type: "response", command: command.type, success: true, data: { sessions: command.streamOnly === true ? [] : sessions, }, }; } case "wait_for_headless_completion": return { type: "response", command: command.type, success: true, data: { enabled: false, continuationsUsed: 0, turnsUsed: 0, tokensUsed: 0, limits: { maxContinuations: 0 }, }, }; case "roster_subscribe": return { type: "response", command: command.type, success: true, data: { roster: [] } }; case "wait_for_idle": case "set_scoped_models": case "rename_saved_session": case "extension_ui_response": case "detach": return { type: "response", command: command.type, success: true }; case "cancel_rlm_child": if (command.childId === "stale-daemon") { return { type: "response", command: command.type, success: false, error: "Unknown daemon command: cancel_rlm_child", }; } return { type: "response", command: command.type, success: true, data: { cancelled: command.childId === "child-1" }, }; case "execute_bash": if (command.command === "stale-daemon") { return { type: "response", command: command.type, success: false, error: "Unknown daemon command: execute_bash", }; } return { type: "response", command: command.type, success: true }; case "abort_bash": if (this.abortBashUnknownCommand) { return { type: "response", command: command.type, success: false, error: "Unknown daemon command: abort_bash", }; } return { type: "response", command: command.type, success: true }; case "start_side_question": return { type: "response", command: command.type, success: true }; case "delete_saved_session": return { type: "response", command: command.type, success: true, data: { ok: true, method: "trash" }, }; case "refine": return { type: "response", command: command.type, success: true, data: { id: "refine_daemon", summary: "Daemon refinement", rationale: "Test daemon refine timeout", expectedOutcome: "Refine request completes", appliedEdits: [], harnessStatePath: "/tmp/harness_state.json", }, }; case "cron_add": await this.cronAddGate; return { type: "response", command: command.type, success: true, data: { job: { id: `cron-${this.requests.filter((request) => request.type === "cron_add").length}`, status: "active", source: "cron", activeSessionId: command.activeSessionId, sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", cwd: "/tmp/project", prompt: command.prompt, schedule: { kind: "cron", expression: command.schedule }, createdAt: "2026-01-01T00:00:00.000Z", updatedAt: "2026-01-01T00:00:00.000Z", runCount: 0, }, }, }; case "switch_session": return { type: "response", command: command.type, success: false, error: "Stored session working directory does not exist: /tmp/missing\nSession file: /tmp/session.jsonl\nCurrent working directory: /tmp/current", errorInfo: { code: "missing_session_cwd", issue: { sessionFile: "/tmp/session.jsonl", sessionCwd: "/tmp/missing", fallbackCwd: "/tmp/current", }, }, }; case "import_jsonl": return { type: "response", command: command.type, success: false, error: "File not found: /tmp/not-found.jsonl", errorInfo: { code: "session_import_file_not_found", filePath: "/tmp/not-found.jsonl", }, }; default: throw new Error(`Unexpected command: ${command.type}`); } } supportsServerCapability(capability: string): boolean { return this.serverCapabilities.has(capability); } onMessage(listener: DaemonClientMessageListener): () => void { this.messageListeners.add(listener); return () => { this.messageListeners.delete(listener); }; } onClose(listener: DaemonClientCloseListener): () => void { this.closeListeners.add(listener); return () => { this.closeListeners.delete(listener); }; } emitMessage(message: DaemonOutbound): void { for (const listener of [...this.messageListeners]) { listener(message); } } emitClose(error: Error): void { for (const listener of [...this.closeListeners]) { listener(error); } } enableRequestRecovery(): void {} async connect(): Promise { if (this.connected) { throw new Error("Prime Agent daemon client is already connected"); } this.reconnectCount++; if (this.reconnectError) { throw this.reconnectError; } this.connected = true; } async waitForHello(): Promise { return this.hello!; } resetTransportForReconnect(): void { this.resetTransportCount++; this.connected = false; } async reconnect(): Promise { if (this.connected) { return; } this.reconnectCount++; if (this.reconnectError) { throw this.reconnectError; } this.connected = true; } get isConnected(): boolean { return this.connected; } getMessageListenerCount(): number { return this.messageListeners.size; } getCloseListenerCount(): number { return this.closeListeners.size; } close(): void { this.closeCount++; this.connected = false; if (this.emitCloseOnClose) { this.emitClose(new Error("Daemon socket closed")); } } disconnectForReconnect(reason: "shutdown" | "update"): void { this.closeCount++; this.connected = false; this.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", reason)); } } function asDaemonClient(client: FakeDaemonClient): DaemonTransportClient { return client as unknown as DaemonTransportClient; } /** The projected tree entry the daemon sends for a user message. */ function projectedUserEntry(id: string, parentId: string | null): AgentConnectionSessionTreeEntry { return { type: "message", id, parentId, timestamp: "2026-01-01T00:00:00.000Z", role: "user", text: "hello", hasText: true, }; } /** * Push a snapshot-staling session event so the next read goes back to the * daemon instead of the cached attach snapshot. */ function emitStaleSessionActionUpdate(fakeClient: FakeDaemonClient): void { fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [] }, }, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, emittedAt: "2026-01-01T00:00:00.000Z", }, }); } function createConnectionState(activeSessionId: string, sessionId: string): AgentConnectionState { return { activeSessionId, cwd: "/tmp/project", model: undefined, thinkingLevel: "medium", serviceTier: "default", availableThinkingLevels: ["minimal", "low", "medium", "high", "xhigh"], isStreaming: false, isCompacting: false, isBashRunning: false, retryAttempt: 0, steeringMode: "all", followUpMode: "one-at-a-time", sessionFile: `/tmp/${sessionId}.jsonl`, sessionId, sessionName: `${sessionId} name`, sessionDir: "/tmp/sessions", leafId: `${sessionId}-leaf`, autoCompactionEnabled: true, messageCount: 1, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, compactionCount: 0, goal: { active: false, status: "idle", tokensUsed: 0, timeUsedSeconds: 0, continuationsUsed: 0, }, scopedModels: [], activeToolNames: ["ipython"], contextUsage: undefined, }; } interface CreateAttachResultOptions { state?: AgentConnectionState; messages?: AgentMessage[]; streamingMessage?: AgentMessage; sessionContext?: DaemonAttachResult["snapshot"]["sessionContext"]; omitSessionContext?: boolean; sessionTree?: DaemonAttachResult["snapshot"]["sessionTree"]; parent?: DaemonAttachResult["snapshot"]["parent"]; children?: DaemonAttachResult["snapshot"]["children"]; replay?: DaemonAttachResult["replay"]; } function createAttachResult( activeSessionId: string, clientId: string | undefined, capabilities: readonly string[] | undefined, lastEventSequence: number, options: CreateAttachResultOptions = {}, ): DaemonAttachResult { const state = options.state ?? createConnectionState(activeSessionId, "session-current"); const messages = options.messages ?? []; const lastEventCursor = { generation: `generation-${activeSessionId}`, sequence: lastEventSequence }; const summary = { id: activeSessionId, activeSessionId, lifecycle: "live" as const, activity: "idle" as const, isSessionActive: state.isStreaming, sessionId: state.sessionId, cwd: "/tmp/project", isStreaming: state.isStreaming, isCompacting: false, attachedClients: 1, messageCount: messages.length, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, ...(options.streamingMessage ? { streamingMessage: options.streamingMessage } : {}), }; // Slim shape: the daemon omits top-level state/messages for clients with the // "slim_attach" capability, which DaemonAgentConnection always advertises. return { protocol: DAEMON_PROTOCOL_INFO, activeSessionId, snapshot: { activeSessionId, summary, state, messages, ...(options.omitSessionContext ? {} : { sessionContext: options.sessionContext ?? ({ messages, thinkingLevel: state.thinkingLevel, serviceTier: state.serviceTier, model: state.model ? { provider: state.model.provider, modelId: state.model.id } : null, } satisfies NonNullable), }), sessionTree: options.sessionTree ?? { tree: [], leafId: state.leafId }, lastEventSequence, lastEventCursor, ...(options.parent ? { parent: options.parent } : {}), ...(options.children ? { children: options.children } : {}), }, replay: options.replay ?? { status: "complete", toSequence: lastEventSequence, toCursor: lastEventCursor, }, lastEventSequence, lastEventCursor, client: { id: clientId ?? "client-1", capabilities: (capabilities ?? ["attach_snapshot", "event_sequence"]).filter( (capability): capability is DaemonAttachResult["client"]["capabilities"][number] => capability === "attach_snapshot" || capability === "event_sequence" || capability === "extension_ui" || capability === "slim_attach" || capability === "chunked_snapshot", ), }, }; } function emitRlmChildUpdate( client: FakeDaemonClient, activeSessionId: string, sequence: number, child: AgentConnectionRlmChildAgentSnapshot, ): void { client.emitMessage({ type: "session_event", activeSessionId, event: { type: "rlm_child_update", child }, meta: { id: `${activeSessionId}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId, sequence, cursor: { generation: `generation-${activeSessionId}`, sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); } function emitSequencedQueueUpdate(client: FakeDaemonClient, activeSessionId: string, sequence: number): void { client.emitMessage({ type: "session_event", activeSessionId, event: { type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [] }, }, meta: { id: `${activeSessionId}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId, sequence, cursor: { generation: `generation-${activeSessionId}`, sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); } describe("DaemonAgentConnection", () => { it("falls back to the supervisor when the direct socket closes during initial attach", async () => { const supervisor = new FakeDaemonClient(); const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async () => { directConnected = false; const error = new Error("direct attach socket closed"); for (const listener of [...closeListeners]) listener(error); throw error; }, close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ state: { activeSessionId: "active-1" }, }); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(1); await connection.dispose(); }); it("keeps serving the session on the direct link and reconnects a lost supervisor socket", async () => { const supervisor = new FakeDaemonClient(); const directRequests: DaemonCommand["type"][] = []; let closeSent = false; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: DaemonCommand) => { directRequests.push(command.type); if (!closeSent) { closeSent = true; supervisor.connected = false; supervisor.emitClose(new Error("supervisor closed during direct attach")); } return { type: "response" as const, command: command.type, success: true as const, data: command.type === "attach" ? createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12) : undefined, }; }, close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); await vi.waitFor(() => expect(supervisor.reconnectCount).toBe(1)); expect(routed.hasDirectTransport).toBe(true); // Held-direct recovery is control-plane only: no re-attach crosses either socket. expect(directRequests.filter((type) => type === "attach")).toHaveLength(1); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(0); await connection.dispose(); }); it("rejects the fallback attach with the authoritative shutdown instead of parking it", async () => { const supervisor = new FakeDaemonClient(); const attachOptions: (DaemonClientRequestOptions | undefined)[] = []; const originalRequest = supervisor.request.bind(supervisor); supervisor.request = async (command: DaemonCommand, timeoutMs?: number, options?: DaemonClientRequestOptions) => { if (command.type === "attach") { attachOptions.push(options); throw new Error("supervisor socket is gone"); } return originalRequest(command, timeoutMs ?? 30000, options ?? {}); }; const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async () => { // The authoritative stop lands while the direct link is still up; then the direct attach dies. supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown")); directConnected = false; const error = new Error("direct attach socket closed"); for (const listener of [...closeListeners]) listener(error); throw error; }, close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); await expect(DaemonAgentConnection.attach(routed, "active-1")).rejects.toThrow("Reason: shutdown"); expect(attachOptions).toEqual([expect.objectContaining({ recoverable: false })]); }); it("absorbs a direct loss inside the held roster re-attach into a session-plane reattach", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("agent_roster"); const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; let rosterSubscribes = 0; const originalRequest = supervisor.request.bind(supervisor); supervisor.request = async (command: DaemonCommand, timeoutMs?: number, options?: DaemonClientRequestOptions) => { if (command.type === "roster_subscribe") { rosterSubscribes++; if (rosterSubscribes === 2) { directConnected = false; for (const listener of [...closeListeners]) listener(new Error("direct worker socket died")); } } return originalRequest(command, timeoutMs ?? 30000, options ?? {}); }; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async (command: Extract) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); await connection.subscribeAgentRoster(() => {}); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); supervisor.hello = { ...supervisor.hello! }; supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket lost")); await vi.waitFor(() => expect(events.filter((event) => event.type === "session_resynced")).toHaveLength(1)); expect(events.some((event) => event.type === "closed")).toBe(false); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(1); await connection.dispose(); }); it("stands down for update restoration when the update close lands inside the held roster re-attach", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("agent_roster"); supervisor.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; let rosterSubscribes = 0; const originalRequest = supervisor.request.bind(supervisor); supervisor.request = async (command: DaemonCommand, timeoutMs?: number, options?: DaemonClientRequestOptions) => { if (command.type === "roster_subscribe") { rosterSubscribes++; if (rosterSubscribes === 2) { supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); } } return originalRequest(command, timeoutMs ?? 30000, options ?? {}); }; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); // A short deadline would abort the restore if the held loop wrongly stayed its owner. const connection = await DaemonAgentConnection.attach(routed, "active-1", { reconnectTimeoutMs: 30 }); await connection.subscribeAgentRoster(() => {}); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); supervisor.hello = { ...supervisor.hello! }; supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket lost")); await vi.waitFor(() => expect(events.some((event) => event.type === "session_resynced")).toBe(true), { timeout: 5000, }); expect(events.find((event) => event.type === "session_resynced")).toMatchObject({ snapshot: { state: { activeSessionId: "active-restored" } }, }); expect(events.some((event) => event.type === "closed")).toBe(false); await connection.dispose(); }); it("rebinds the roster subscription onto the recovered supervisor socket while the direct link holds", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("agent_roster"); const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); await connection.subscribeAgentRoster(() => {}); // A reconnect delivers a fresh hello object; the roster store keys its subscription on it. supervisor.hello = { ...supervisor.hello! }; supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket lost")); await vi.waitFor(() => expect(supervisor.requests.filter((request) => request.type === "roster_subscribe")).toHaveLength(2), ); expect(routed.hasDirectTransport).toBe(true); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(0); await connection.dispose(); }); it("resets a half-open supervisor handshake without closing the healthy direct socket", async () => { const supervisor = new FakeDaemonClient(); const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); let helloAttempts = 0; supervisor.waitForHello = vi.fn(async () => { helloAttempts++; if (helloAttempts === 1) throw new Error("hello timed out on half-open socket"); return supervisor.hello!; }); supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket closed")); await vi.waitFor(() => expect(events.some((event) => event.type === "connection_status" && event.status === "connected")).toBe(true), ); expect(supervisor.reconnectCount).toBe(2); expect(supervisor.resetTransportCount).toBe(1); expect(routed.hasDirectTransport).toBe(true); await connection.dispose(); }); it("routes an update close through update restoration and drops the stale direct link", async () => { const supervisor = new FakeDaemonClient(); supervisor.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); const restored = new Promise((resolve) => { connection.subscribe(async (event) => { if (event.type === "session_resynced") resolve(event); }); }); supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await expect(restored).resolves.toMatchObject({ type: "session_resynced" }); expect(routed.hasDirectTransport).toBe(false); await connection.dispose(); }); it("falls back to the supervisor when the direct socket dies mid-session without recoverDaemon", async () => { const supervisor = new FakeDaemonClient(); const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async (command: Extract) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); supervisor.serverCapabilities.add("session_input_pause"); const pause = await connection.acquireSessionInputPause("lease-1"); directConnected = false; for (const listener of [...closeListeners]) listener(new Error("direct worker socket died")); await vi.waitFor(() => expect(events.filter((event) => event.type === "session_resynced")).toHaveLength(1)); expect(events.some((event) => event.type === "closed")).toBe(false); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(1); // The fence died with the direct link; its holder learns on release while the session lives on. await expect(pause.release()).rejects.toThrow("invalidated by a daemon reconnect"); await connection.dispose(); }); it("takes update restoration when the direct link closes for an update while a pause is held", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("session_input_pause"); supervisor.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async (command: Extract) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); await connection.acquireSessionInputPause("lease-1"); directConnected = false; for (const listener of [...closeListeners]) { listener(new DaemonSocketClosedError("/tmp/worker.sock", "update")); } await vi.waitFor(() => expect(events.some((event) => event.type === "session_resynced")).toBe(true)); expect(events.some((event) => event.type === "closed")).toBe(false); await connection.dispose(); }); it("outlives the reconnect deadline while the direct link streams, then bounds recovery once it dies", async () => { const supervisor = new FakeDaemonClient(); const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async (command: Extract) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1", { reconnectTimeoutMs: 60 }); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); supervisor.connected = false; supervisor.reconnectError = new Error("supervisor is down"); supervisor.emitClose(new Error("supervisor socket lost")); // Three failed control-plane attempts span well past the 60ms deadline. await vi.waitFor(() => expect(supervisor.resetTransportCount).toBeGreaterThanOrEqual(3)); expect(events.some((event) => event.type === "closed")).toBe(false); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(0); directConnected = false; for (const listener of [...closeListeners]) listener(new Error("direct worker socket died")); await vi.waitFor(() => expect(events.some((event) => event.type === "closed")).toBe(true), { timeout: 8000 }); expect(events.find((event) => event.type === "closed")).toMatchObject({ error: expect.stringContaining("Daemon reconnection failed"), }); expect(events.filter((event) => event.type === "session_resynced")).toHaveLength(0); await connection.dispose(); }, 10_000); it("keeps the source direct socket when a cross-worker reattach is rejected", async () => { const supervisor = new FakeDaemonClient(); supervisor.request = vi.fn(async (command: DaemonCommand) => ({ type: "response" as const, command: command.type, success: false as const, error: "target disappeared", })); const close = vi.fn(); const direct = { isConnected: true, onMessage: () => () => {}, onClose: () => () => {}, close, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const response = await routed.request({ type: "reattach", activeSessionId: "source-active", targetActiveSessionId: "missing-target", }); expect(response.success).toBe(false); expect(close).not.toHaveBeenCalled(); expect(routed.hasDirectTransport).toBe(true); routed.close(); }); it("carries an opt-out-only telemetry policy on attach", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { telemetryDisabled: true, }); await connection.attach(); expect(fakeClient.requests[0]).toMatchObject({ type: "attach", activeSessionId: "active-1", telemetryDisabled: true, }); }); it.each([true, false])("capability-gates owned-session recovery context: %s", async (supported) => { const fakeClient = new FakeDaemonClient(); if (supported) fakeClient.serverCapabilities.add("owned_session_recovery_context"); const recoveryConfig = { cwd: "/tmp/fresh-owner" }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-owned", { ownedSession: true, ownedSessionRecoveryConfig: recoveryConfig, }); await connection.attach(); const request = fakeClient.requests[0]; if (supported) expect(request).toMatchObject({ recoveryConfig }); else expect(request).not.toHaveProperty("recoveryConfig"); }); it("forwards queueIfBusy for prompt admission", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.prompt("queued input", { streamingBehavior: "followUp", queueIfBusy: true }); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "prompt", activeSessionId: "active-1", message: "queued input", streamingBehavior: "followUp", queueIfBusy: true, }); }); it("forwards hidden custom prompt metadata", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const customMessage = createManualContinuationMessage(1); await connection.prompt(MANUAL_CONTINUATION_PROMPT, { customMessage, expandPromptTemplates: false, }); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "prompt", activeSessionId: "active-1", message: MANUAL_CONTINUATION_PROMPT, customMessage, expandPromptTemplates: false, }); }); it("forwards signal-backed prompts with a unique cancellable admission id", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController(); const customMessage = createManualContinuationMessage(1); await connection.prompt("startup", { signal: abort.signal, queueIfBusy: true, customMessage, expandPromptTemplates: false, }); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "prompt", activeSessionId: "active-1", message: "startup", queueIfBusy: true, customMessage, expandPromptTemplates: false, admissionId: expect.stringMatching(/^prompt-admission:/), }); }); it("cancels rejected signal-backed admission and removes its abort listener", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptError = new Error("transport failed"); fakeClient.cancelPromptAdmissionStatus = "cancelled"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController(); const add = vi.spyOn(abort.signal, "addEventListener"); const remove = vi.spyOn(abort.signal, "removeEventListener"); await expect(connection.prompt("startup", { signal: abort.signal })).rejects.toMatchObject({ message: "transport failed", cancelled: true, }); expect(fakeClient.requests.map((request) => request.type)).toEqual(["prompt", "cancel_prompt_admission"]); expect(add).toHaveBeenCalledOnce(); expect(remove).toHaveBeenCalledOnce(); expect(remove.mock.calls[0]?.[1]).toBe(add.mock.calls[0]?.[1]); }); it("preserves a definitive prompt rejection when the signal remains live", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptResponseError = "session rejected prompt"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await expect(connection.prompt("startup", { signal: new AbortController().signal })).rejects.toEqual( new Error("session rejected prompt"), ); expect(fakeClient.requests.map((request) => request.type)).toEqual(["prompt"]); }); it("sends zero requests for a pre-aborted prompt", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController(); abort.abort(); await expect(connection.prompt("startup", { signal: abort.signal })).rejects.toMatchObject({ status: "cancelled", }); expect(fakeClient.requests).toEqual([]); }); it("accepts a successful prompt response when cancellation reports unknown", async () => { const fakeClient = new FakeDaemonClient(); let releasePrompt = () => {}; fakeClient.promptGate = new Promise((resolve) => { releasePrompt = resolve; }); fakeClient.cancelPromptAdmissionStatus = "unknown"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController(); const prompt = connection.prompt("startup", { signal: abort.signal }); abort.abort(); await vi.waitFor(() => expect(fakeClient.requests.map((request) => request.type)).toContain("cancel_prompt_admission"), ); expect(fakeClient.requests.find((request) => request.type === "cancel_prompt_admission")).not.toHaveProperty( "cancelOwned", ); releasePrompt(); await expect(prompt).resolves.toBeUndefined(); }); it("requests owned prompt cancellation when the daemon advertises it", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("owned_prompt_cancellation"); let releasePrompt = () => {}; fakeClient.promptGate = new Promise((resolve) => { releasePrompt = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController(); const prompt = connection.prompt("startup", { signal: abort.signal }); abort.abort(); await vi.waitFor(() => expect(fakeClient.requests.map((request) => request.type)).toContain("cancel_prompt_admission"), ); expect(fakeClient.requests.find((request) => request.type === "cancel_prompt_admission")).toMatchObject({ cancelOwned: true, }); releasePrompt(); await expect(prompt).resolves.toBeUndefined(); }); it("preserves a definitive prompt rejection when cancellation reports owned", async () => { const fakeClient = new FakeDaemonClient(); let releasePrompt = () => {}; fakeClient.promptGate = new Promise((resolve) => { releasePrompt = resolve; }); fakeClient.promptResponseError = "post-ownership validation failed"; fakeClient.cancelPromptAdmissionStatus = "owned"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController(); const prompt = connection.prompt("startup", { signal: abort.signal }); abort.abort(); await vi.waitFor(() => expect(fakeClient.requests.map((request) => request.type)).toContain("cancel_prompt_admission"), ); releasePrompt(); await expect(prompt).rejects.toThrow("post-ownership validation failed"); }); it("treats a lost prompt response plus unknown cancellation as uncertain", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptError = new Error("lost response"); fakeClient.cancelPromptAdmissionStatus = "unknown"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await expect(connection.prompt("startup", { signal: new AbortController().signal })).rejects.toMatchObject({ message: "lost response", status: "unknown", cancelled: false, }); }); it("translates unsupported admission cancellation into an admission error", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptError = new DaemonCapabilityUnavailableError("prompt", "prompt_admission_cancellation"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await expect(connection.prompt("startup", { signal: new AbortController().signal })).rejects.toMatchObject({ status: "unsupported", }); expect(fakeClient.requests.map((request) => request.type)).toEqual(["prompt"]); }); it("uses cancellable admission for promptAndWait", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptError = new Error("lost response"); fakeClient.cancelPromptAdmissionStatus = "cancelled"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await expect(connection.promptAndWait("wait", { signal: new AbortController().signal })).rejects.toMatchObject({ status: "cancelled", }); expect(fakeClient.requests.map((request) => request.type)).toEqual([ "prompt_and_wait", "cancel_prompt_admission", ]); }); it("uses fleet heartbeat scope for residents and session scope for owned workers", async () => { const residentClient = new FakeDaemonClient(); residentClient.serverCapabilities.add("heartbeat_catalog"); const resident = new DaemonAgentConnection(asDaemonClient(residentClient), "resident-1"); await resident.listHeartbeats(); const ownedClient = new FakeDaemonClient(); ownedClient.serverCapabilities.add("heartbeat_catalog"); const owned = new DaemonAgentConnection(asDaemonClient(ownedClient), "owned-1", { ownedSession: true }); await owned.listHeartbeats(); expect(residentClient.requests.at(-1)).toEqual(expect.objectContaining({ type: "heartbeats_list" })); expect(residentClient.requests.at(-1)).not.toHaveProperty("activeSessionId"); expect(ownedClient.requests.at(-1)).toEqual( expect.objectContaining({ type: "heartbeats_list", activeSessionId: "owned-1" }), ); }); it("serializes concurrent owned-session promotion commands", async () => { const fakeClient = new FakeDaemonClient(); let releaseFirst = () => {}; fakeClient.cronAddGate = new Promise((resolve) => { releaseFirst = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { ownedSession: true, }); const first = connection.addCronJob("0 * * * *", "first"); const second = connection.addCronJob("30 * * * *", "second"); await vi.waitFor(() => { expect(fakeClient.requests.filter((request) => request.type === "cron_add")).toHaveLength(1); }); releaseFirst(); await Promise.all([first, second]); const commands = fakeClient.requests.filter( (command): command is Extract => command.type === "cron_add", ); expect(commands.map((command) => command.promoteOwnedSession)).toEqual([true, false]); }); it("gates side-question follow-up transcripts on the daemon capability", async () => { const oldDaemonClient = new FakeDaemonClient(); const oldConnection = new DaemonAgentConnection(asDaemonClient(oldDaemonClient), "active-original"); // First questions carry no transcript and must keep working on old daemons. await oldConnection.startSideQuestion("turn-1", "What changed?"); expect(oldDaemonClient.requests.map((request) => request.type)).toEqual(["start_side_question"]); // Follow-ups would silently lose their side context on an old daemon. await expect( oldConnection.startSideQuestion("turn-2", "And then?", [{ question: "What changed?", answer: "The parser." }]), ).rejects.toThrow("older build without side-conversation follow-ups"); expect(oldDaemonClient.requests).toHaveLength(1); const newDaemonClient = new FakeDaemonClient(); newDaemonClient.serverCapabilities.add("side_question_transcript"); const newConnection = new DaemonAgentConnection(asDaemonClient(newDaemonClient), "active-original"); await newConnection.startSideQuestion("turn-2", "And then?", [ { question: "What changed?", answer: "The parser." }, ]); const sent = newDaemonClient.requests.find( (command): command is Extract => command.type === "start_side_question", ); expect(sent?.previousTurns).toEqual([{ question: "What changed?", answer: "The parser." }]); }); it("gates transient bash on the daemon capability", async () => { const oldDaemonClient = new FakeDaemonClient(); const oldConnection = new DaemonAgentConnection(asDaemonClient(oldDaemonClient), "active-original"); // Regular bash keeps working on old daemons. await oldConnection.executeBash("ls"); expect(oldDaemonClient.requests.map((request) => request.type)).toEqual(["execute_bash"]); // A transient run on an old daemon would be recorded into the session. await expect(oldConnection.executeBash("ls", { transient: true })).rejects.toThrow( "older build without side-conversation bash", ); expect(oldDaemonClient.requests).toHaveLength(1); const newDaemonClient = new FakeDaemonClient(); newDaemonClient.serverCapabilities.add("transient_bash"); const newConnection = new DaemonAgentConnection(asDaemonClient(newDaemonClient), "active-original"); await newConnection.executeBash("ls", { excludeFromContext: true, transient: true, runId: "side-run-1" }); const sent = newDaemonClient.requests.find( (command): command is Extract => command.type === "execute_bash", ); expect(sent).toMatchObject({ command: "ls", excludeFromContext: true, transient: true, runId: "side-run-1" }); }); it("degrades an unavailable heartbeat catalog without sending an unsupported command", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); await expect(connection.listHeartbeats()).resolves.toEqual([]); expect(fakeClient.requests).toEqual([]); await expect(connection.manageHeartbeat("active-original", "job-1", "pause")).rejects.toThrow( "requires a newer Prime Agent daemon", ); expect(fakeClient.requests).toEqual([]); }); it.each([false, true])("reattaches after an update restart (deferring=%s)", async (deferSessionEvents) => { const fakeClient = new FakeDaemonClient(); fakeClient.emitCloseOnClose = true; const restoredMessages: AgentMessage[] = [{ role: "user", content: "restored prompt", timestamp: 2 }]; fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: createConnectionState(command.activeSessionId, "session-current"), messages: command.activeSessionId === "active-restored" ? restoredMessages : [], }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { deferSessionEvents, }); const events: AgentConnectionEvent[] = []; const restored = new Promise((resolve) => { connection.subscribe((event) => { events.push(event); if (event.type === "session_resynced") { resolve(event); } }); }); // The reconnect settles without a timer: wait on the status sequence itself. let sawReconnecting = false; const connectedAgain = new Promise((resolve) => { connection.subscribe((event) => { if (event.type !== "connection_status") return; if (event.status === "reconnecting") { sawReconnecting = true; } else if (sawReconnecting && event.status === "connected") { resolve(); } }); }); await connection.attach(); if (deferSessionEvents) emitSequencedSessionEvent(fakeClient, "active-original", 100); fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "update", }); await restored; expect(fakeClient.closeCount).toBe(1); expect(fakeClient.reconnectCount).toBe(1); await expect(restored).resolves.toMatchObject({ type: "session_resynced", snapshot: { state: { activeSessionId: "active-restored", sessionId: "session-current" }, messages: restoredMessages, }, }); await connection.flushBufferedSessionEvents(); expect(fakeClient.reconnectCount).toBe(1); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "list", "attach"]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-restored", resumeCursor: undefined, }); await connectedAgain; expect(events).toEqual([ expect.objectContaining({ type: "connection_status", status: "reconnecting" }), expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ activeSessionId: "active-restored" }), }), }), { type: "connection_status", status: "connected" }, ]); }); it("reattaches when an update socket close arrives before the session notice", async () => { const fakeClient = new FakeDaemonClient(); const restoredMessages: AgentMessage[] = [{ role: "user", content: "restored prompt", timestamp: 2 }]; fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: createConnectionState(command.activeSessionId, "session-current"), messages: command.activeSessionId === "active-restored" ? restoredMessages : [], }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const restored = new Promise((resolve) => { connection.subscribe((event) => { if (event.type === "session_resynced") { resolve(event); } }); }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await expect(restored).resolves.toMatchObject({ type: "session_resynced", snapshot: { state: { activeSessionId: "active-restored", sessionId: "session-current" }, messages: restoredMessages, }, }); expect(fakeClient.reconnectCount).toBe(1); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "list", "attach"]); }); it("coordinates one transport reconnect across connections sharing a daemon client", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.emitCloseOnClose = true; fakeClient.updateRestartSessions = [ { id: "restored-a", activeSessionId: "restored-a", sessionId: "session-a", sessionFile: "/tmp/session-a.jsonl", }, { id: "restored-b", activeSessionId: "restored-b", sessionId: "session-b", sessionFile: "/tmp/session-b.jsonl", }, ]; const sessionIds: Record = { "active-a": "session-a", "active-b": "session-b", "restored-a": "session-a", "restored-b": "session-b", }; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: createConnectionState(command.activeSessionId, sessionIds[command.activeSessionId]!), }); const connectionA = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-a"); const connectionB = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-b"); await connectionA.attach(); await connectionB.attach(); const restoredA = new Promise((resolve) => { connectionA.subscribe((event) => { if (event.type === "session_resynced") { resolve(event); } }); }); const restoredB = new Promise((resolve) => { connectionB.subscribe((event) => { if (event.type === "session_resynced") { resolve(event); } }); }); fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-a", reason: "update" }); await expect(Promise.all([restoredA, restoredB])).resolves.toEqual([ expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ sessionId: "session-a" }), }), }), expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ sessionId: "session-b" }), }), }), ]); expect(fakeClient.closeCount).toBe(1); expect(fakeClient.reconnectCount).toBe(1); }); it("does not reconnect after an explicit shutdown session close", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach(); fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "shutdown", }); fakeClient.emitClose(new Error("Daemon socket closed")); await Promise.resolve(); expect(fakeClient.reconnectCount).toBe(0); expect(closedEvents).toHaveLength(1); expect(closedEvents[0]).toMatchObject({ type: "closed", error: expect.stringContaining("The Prime Agent daemon shut down while this window was attached."), }); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain("Session ID: session-current."); expect(closedError).toContain("Session file: /tmp/session-current.jsonl."); expect(closedError).toContain("Diagnostic log:"); }); it("does not infer an update when shutdown closes the socket before the session notice", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown")); await Promise.resolve(); expect(fakeClient.reconnectCount).toBe(0); expect(closedEvents).toHaveLength(1); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain("The Prime Agent daemon shut down while this window was attached."); }); it("re-attaches and accepts prompt when connection drops with recoverDaemon", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 5000, }); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach(); fakeClient.connected = false; fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown")); await connection.prompt("hello after restart"); expect(events.some((e) => e.type === "connection_status" && e.status === "reconnecting")).toBe(true); expect(events.some((e) => e.type === "session_resynced")).toBe(true); expect(events.some((e) => e.type === "connection_status" && e.status === "connected")).toBe(true); expect( fakeClient.requests.some( (r) => r.type === "prompt" && "message" in r && (r as { message: unknown }).message === "hello after restart", ), ).toBe(true); }); it("retries and succeeds when attach is refused with Session worker is starting", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 5000, }); const resynced = new Promise((resolve) => { connection.subscribe((event) => { if (event.type === "session_resynced") resolve(event); }); }); await connection.attach(); fakeClient.attachFailures = 1; fakeClient.attachError = new DaemonSessionStartingError(undefined, "Session worker is starting"); fakeClient.connected = false; fakeClient.emitClose(new Error("socket closed")); const event = await resynced; expect(event.type).toBe("session_resynced"); expect(fakeClient.resetTransportCount).toBe(0); }); it("emits a clear give-up error with transcript details after reconnect timeout expires", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.reconnectError = new Error("connection refused"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 50, }); const closed = new Promise((resolve) => { connection.subscribe((event) => { if (event.type === "closed") resolve(event.error ?? ""); }); }); await connection.attach(); fakeClient.connected = false; fakeClient.emitClose(new Error("socket died")); const error = await closed; expect(error).toContain("Daemon reconnection failed"); expect(error).toContain("restart Prime Agent or reopen the session from Agents View"); expect(error).toContain("Session ID: session-current."); }); it("retries initial attach when worker is starting", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.attachFailures = 1; fakeClient.attachError = new DaemonSessionStartingError(undefined, "Session worker is starting"); const connection = await DaemonAgentConnection.attach(asDaemonClient(fakeClient), "active-original", { reconnectTimeoutMs: 5000, }); expect(connection).toBeDefined(); await connection.dispose(); }); it("waits for queued restore and remaps activeSessionId through list", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 5000, }); const resynced = new Promise((resolve) => { connection.subscribe((event) => { if (event.type === "session_resynced") resolve(event); }); }); await connection.attach(); let listQueries = 0; const originalRequest = fakeClient.request.bind(fakeClient); fakeClient.request = async (command, ...args) => { if ( command.type === "attach" && "activeSessionId" in command && (command as any).activeSessionId === "active-original" ) { return { type: "response", command: command.type, success: false, error: "Unknown active session: active-original", }; } if (command.type === "list") { listQueries++; if (listQueries === 1) { return { type: "response", command: command.type, success: true, data: { sessions: [{ sessionId: "session-current", sessionFile: "/tmp/session.jsonl" }] }, }; } return { type: "response", command: command.type, success: true, data: { sessions: [ { activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session.jsonl", }, ], }, }; } return originalRequest(command, ...args); }; fakeClient.connected = false; fakeClient.emitClose(new Error("socket closed")); const event = await resynced; expect(event.type).toBe("session_resynced"); expect(listQueries).toBeGreaterThanOrEqual(2); expect((event as any).snapshot.state.activeSessionId).toBe("active-restored"); }); it("emits accurate update recovery error after update recovery timeout", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.reconnectError = new Error("connection refused"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 10, }); const closed = new Promise((resolve) => { connection.subscribe((event) => { if (event.type === "closed") resolve(event.error ?? ""); }); }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); const error = await closed; expect(error).toContain("restarted for an update, but this window could not reconnect"); expect(error).not.toContain("daemon shut down while this window was attached"); expect(error).toContain("Session ID: session-current."); }); it.each([false, true])("bounds the reattach wait (saved transcript=%s)", async (transcriptSaved) => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); const directory = mkdtempSync(join(tmpdir(), "prime-unrestorable-")); const transcript = join(directory, "session-current.jsonl"); if (transcriptSaved) writeFileSync(transcript, ""); try { const fakeClient = new FakeDaemonClient(); fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: { ...createConnectionState(command.activeSessionId, "session-current"), sessionFile: transcript }, }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { sessionSearchTimeoutMs: 50, }); const closed = new Promise((resolve) => { connection.subscribe((event) => { if (event.type === "closed") resolve(event.error ?? ""); }); }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.advanceTimersByTimeAsync(300); // The bounded search, not the 120s update window, produced the verdict. const error = await closed; expect(error).toContain("is not in the restarted daemon's session list"); expect(error).toContain("Session ID: session-current."); expect(error).toContain(transcriptSaved ? "/resume session-current" : "start a fresh session with /new"); const readsAtVerdict = fakeClient.requests.filter((request) => request.type === "list").length; await vi.advanceTimersByTimeAsync(120_000); expect(fakeClient.requests.filter((request) => request.type === "list")).toHaveLength(readsAtVerdict); await expect(connection.prompt("after the restart")).rejects.toThrow("restarted daemon's session list"); } finally { rmSync(directory, { recursive: true, force: true }); vi.useRealTimers(); } }); it("does not carry a resolved roster miss into the next restart's first miss", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); const events: AgentConnectionEvent[] = []; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { sessionSearchTimeoutMs: 50, }); connection.subscribe((event) => void events.push(event)); await connection.attach(); // First restart: the roster misses the session for a moment, then lists it again. fakeClient.updateRestartSessions = []; fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.advanceTimersByTimeAsync(20); expect(events.filter((event) => event.type === "closed")).toHaveLength(0); fakeClient.updateRestartSessions = [ { id: "active-original", sessionId: "session-current", activeSessionId: "active-restored" }, ]; await vi.advanceTimersByTimeAsync(120); expect(events.filter((event) => event.type === "closed")).toHaveLength(0); // Much later, a second restart misses again: the fresh miss must wait its own bound. await vi.advanceTimersByTimeAsync(120_000); fakeClient.updateRestartSessions = []; fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.advanceTimersByTimeAsync(5); expect(events.filter((event) => event.type === "closed")).toHaveLength(0); await vi.advanceTimersByTimeAsync(400); expect(events.filter((event) => event.type === "closed")).toHaveLength(1); } finally { vi.useRealTimers(); } }); it("keeps waiting for a listed session and re-attaches it once its worker is routed", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.updateRestartSessions = [{ id: "active-original", sessionId: "session-current" }]; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { sessionSearchTimeoutMs: 50, }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); // The daemon lists the session while its worker is still queued: the bound must not fire. await vi.advanceTimersByTimeAsync(600); fakeClient.updateRestartSessions = [ { id: "active-original", sessionId: "session-current", activeSessionId: "active-restored" }, ]; await vi.advanceTimersByTimeAsync(600); const attaches = fakeClient.requests.filter((request) => request.type === "attach"); expect(attaches.at(-1)).toMatchObject({ activeSessionId: "active-restored" }); await expect(connection.prompt("after the restore")).resolves.toBeUndefined(); } finally { vi.useRealTimers(); } }); async function attachLostConnection(fakeClient: FakeDaemonClient): Promise<{ connection: DaemonAgentConnection; events: AgentConnectionEvent[]; }> { const events: AgentConnectionEvent[] = []; const client = asDaemonClient(fakeClient); const connection = new DaemonAgentConnection(client, "active-original", { sessionSearchTimeoutMs: 50 }); connection.subscribe((event) => void events.push(event)); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.advanceTimersByTimeAsync(300); return { connection, events }; } it("starts a fresh session when /new is requested after the unrestorable verdict", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.createdActiveSessionId = "active-fresh"; const { connection, events } = await attachLostConnection(fakeClient); await expect(connection.newSession()).resolves.toEqual({ cancelled: false }); expect(fakeClient.requests.find((request) => request.type === "create")).toMatchObject({ config: { cwd: "/tmp/project" }, }); expect(fakeClient.requests.filter((request) => request.type === "attach").at(-1)).toMatchObject({ activeSessionId: "active-fresh", }); expect(events.filter((event) => event.type === "session_replaced").at(-1)).toMatchObject({ state: { activeSessionId: "active-fresh" }, }); await expect(connection.prompt("after the recovery")).resolves.toBeUndefined(); } finally { vi.useRealTimers(); } }); it.each([undefined, "/tmp/elsewhere"])( "resumes a transcript after the verdict (cwdOverride=%s)", async (cwdOverride) => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); const { connection } = await attachLostConnection(fakeClient); await expect(connection.switchSession("/tmp/saved.jsonl", { cwdOverride })).resolves.toEqual({ cancelled: false, }); const create = fakeClient.requests.find((request) => request.type === "create"); expect(create).toMatchObject({ sessionPath: "/tmp/saved.jsonl" }); expect(create?.type === "create" ? create.config?.cwd : undefined).toBe(cwdOverride); } finally { vi.useRealTimers(); } }, ); it("reports an actionable error and stays retryable when the replacement cannot be created", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, sessionSearchTimeoutMs: 50, }); await connection.attach(); // This verdict arrives after a socket loss, which must leave the transport revivable. fakeClient.attachFailures = 99; fakeClient.attachError = new Error("Unknown active session: active-original"); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock")); await vi.advanceTimersByTimeAsync(2000); expect(fakeClient.closeCount).toBe(0); fakeClient.attachFailures = 0; fakeClient.createError = new Error("connect ECONNREFUSED"); await expect(connection.newSession()).rejects.toThrow(/ECONNREFUSED.*relaunch this window/); // The verdict survives: a retry reaches the daemon again instead of a poisoned connection. await expect(connection.newSession()).rejects.toThrow("relaunch this window"); expect(fakeClient.requests.filter((request) => request.type === "create")).toHaveLength(2); await expect(connection.prompt("x")).rejects.toThrow("restarted daemon's session list"); } finally { vi.useRealTimers(); } }); it("reports a failed session command instead of exiting the client after the terminal close", async () => { const shown: string[] = []; const listeners = new Set<(event: AgentConnectionEvent) => void>(); const mode = Object.create(InteractiveMode.prototype) as { connectionClosed: boolean; handleClearCommand(): Promise; subscribeToAgent(): void; }; Object.assign(mode, { stop: () => {}, stopWorkingLoader: () => {}, clearStatus: () => {}, agentConnection: { subscribe: (listener: (event: AgentConnectionEvent) => void) => { listeners.add(listener); return () => listeners.delete(listener); }, newSession: async () => { throw new Error("The daemon could not start a replacement session: ECONNREFUSED. Relaunch this window."); }, }, showError: (message: string) => { shown.push(message); }, }); const exit = vi.spyOn(process, "exit").mockImplementation((code) => { throw new Error(`process.exit(${String(code)})`); }); mode.subscribeToAgent(); for (const listener of listeners) await listener({ type: "closed", error: "restarted daemon's session list" }); expect(mode.connectionClosed).toBe(true); await expect(mode.handleClearCommand()).resolves.toBeUndefined(); expect(exit).not.toHaveBeenCalled(); expect(shown.join(" ")).toContain("Relaunch this window"); }); it("remaps activeSessionId and succeeds on request failing with Unknown active session", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 5000, }); await connection.attach(); fakeClient.request = async (command) => { if ( "activeSessionId" in command && (command as { activeSessionId: unknown }).activeSessionId === "active-original" ) { return { type: "response", command: command.type, success: false, error: "Unknown active session: active-original", }; } if (command.type === "list") { return { type: "response", command: command.type, success: true, data: { sessions: [ { activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ], }, }; } if (command.type === "get_session_stats") { return { type: "response", command: command.type, success: true, data: { sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", userMessages: 1, assistantMessages: 1, toolCalls: 0, toolResults: 0, totalMessages: 2, }, }; } return { type: "response", command: command.type, success: true, data: createAttachResult("active-restored", "fake-client", [], 12), }; }; const stats = await connection.getSessionStats(); expect(stats.sessionId).toBe("session-current"); }); it("survives a request when remap fails without unhandled crash", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 10, }); await connection.attach(); fakeClient.request = async (command) => ({ type: "response", command: command.type, success: false, error: "Unknown active session: active-original", }); await expect(connection.getSessionStats()).rejects.toThrow(); }); it("bounds remap wait on unlisted session and recovers via /new", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.promptResponseError = "Unknown active session: active-original"; const events: AgentConnectionEvent[] = []; const client = asDaemonClient(fakeClient); const connection = new DaemonAgentConnection(client, "active-original", { sessionSearchTimeoutMs: 50 }); connection.subscribe((e) => void events.push(e)); await connection.attach(); fakeClient.attachFailures = 999; fakeClient.attachError = new Error("Unknown active session: active-original"); const promptRejection = expect(connection.prompt("hi")).rejects.toThrow("restarted daemon's session list"); await vi.advanceTimersByTimeAsync(300); await promptRejection; expect(events.filter((e) => e.type === "closed")).toHaveLength(1); fakeClient.attachFailures = 0; fakeClient.createdActiveSessionId = "active-fresh"; await expect(connection.newSession()).resolves.toEqual({ cancelled: false }); expect(events.filter((e) => e.type === "session_replaced")).toHaveLength(1); } finally { vi.useRealTimers(); } }); it("single-flights concurrent /new in the verdict state", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.createdActiveSessionId = "active-fresh"; const { connection } = await attachLostConnection(fakeClient); const results = await Promise.all([connection.newSession(), connection.newSession()]); expect(results).toEqual([{ cancelled: false }, { cancelled: false }]); expect(fakeClient.requests.filter((r) => r.type === "create")).toHaveLength(1); expect(fakeClient.requests.filter((r) => r.type === "new_session")).toHaveLength(0); } finally { vi.useRealTimers(); } }); it("releases created session when disposed mid-replacement without unhandled error", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.createdActiveSessionId = "active-fresh"; const { connection } = await attachLostConnection(fakeClient); let release!: () => void; fakeClient.createGate = new Promise((r) => { release = r; }); const [replace, dispose] = [connection.newSession(), connection.dispose()]; release(); await Promise.all([replace, dispose]); const releases = fakeClient.requests.filter((r) => r.type === "detach" || r.type === "complete_owned_session"); expect(releases.at(-1)).toMatchObject({ activeSessionId: "active-fresh" }); } finally { vi.useRealTimers(); } }); it("rejects queued prompt with unrestorable verdict when reconnect produces verdict", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); const client = asDaemonClient(fakeClient); const connection = new DaemonAgentConnection(client, "active-original", { sessionSearchTimeoutMs: 50 }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); const promptRejection = expect(connection.prompt("queued")).rejects.toThrow("restarted daemon's session list"); await vi.advanceTimersByTimeAsync(300); await promptRejection; } finally { vi.useRealTimers(); } }); it("does not emit a restored session after disposal begins", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; let releaseRestoredAttach: (() => void) | undefined; fakeClient.restoredAttachGate = new Promise((resolve) => { releaseRestoredAttach = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.waitFor(() => { expect( fakeClient.requests.some( (request) => request.type === "attach" && request.activeSessionId === "active-restored", ), ).toBe(true); }); await connection.dispose(); releaseRestoredAttach?.(); await vi.waitFor(() => { expect(fakeClient.restoredAttachCompleted).toBe(1); }); for (let flush = 0; flush < 5; flush++) { await Promise.resolve(); } expect(events).toEqual([expect.objectContaining({ type: "connection_status", status: "reconnecting" })]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "detach", activeSessionId: "active-restored" }); }); it("returns to normal close handling after update restoration times out", async () => { vi.useFakeTimers(); try { const fakeClient = new FakeDaemonClient(); fakeClient.reconnectError = new Error("daemon unavailable"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach(); fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "update", }); await vi.advanceTimersByTimeAsync(120100); expect(closedEvents).toHaveLength(1); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain( "The Prime Agent daemon restarted for an update, but this window could not reconnect", ); expect(closedError).toContain("Last error: daemon unavailable"); expect(closedError).toContain("restart Prime Agent and reopen it from Agents View"); expect(closedError).toContain("Session ID: session-current."); expect(closedError).toContain("Session file: /tmp/session-current.jsonl."); expect(closedError).toContain("Diagnostic log:"); const reconnectCountAfterFailure = fakeClient.reconnectCount; fakeClient.emitClose(new Error("Daemon socket closed")); await Promise.resolve(); expect(fakeClient.reconnectCount).toBe(reconnectCountAfterFailure); expect(closedEvents).toHaveLength(1); await connection.dispose(); } finally { vi.useRealTimers(); } }); it("reattaches using the replacement session identity after a session switch", async () => { const fakeClient = new FakeDaemonClient(); const restoredMessages: AgentMessage[] = [{ role: "user", content: "switched prompt", timestamp: 3 }]; fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-next", sessionFile: "/tmp/session-next.jsonl", }, ]; fakeClient.attachResultFactory = (command) => { const sessionId = command.activeSessionId === "active-restored" ? "session-next" : "session-current"; return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: createConnectionState(command.activeSessionId, sessionId), messages: command.activeSessionId === "active-restored" ? restoredMessages : [], }); }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); await connection.attach(); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-original", state: createConnectionState("active-original", "session-next"), messages: [{ role: "user", content: "switched prompt", timestamp: 2 }], }); const restored = new Promise((resolve) => { connection.subscribe((event) => { if (event.type === "session_resynced") { resolve(event); } }); }); fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "update", }); fakeClient.emitClose(new Error("Daemon socket closed")); await expect(restored).resolves.toMatchObject({ type: "session_resynced", snapshot: { state: { activeSessionId: "active-restored", sessionId: "session-next" }, messages: restoredMessages, }, }); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "list", "attach"]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-restored", }); }); it("forwards catch-up snapshots as non-destructive resync events", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach(); const messages: AgentMessage[] = [{ role: "user", content: "caught up", timestamp: 2 }]; const streamingMessage = { role: "assistant", content: [{ type: "thinking", thinking: "Still reasoning" }], } as AgentMessage; const snapshot = createAttachResult("active-1", "client-1", undefined, 13, { state: { ...createConnectionState("active-1", "session-current"), isStreaming: true }, messages, streamingMessage, }).snapshot; fakeClient.emitMessage({ type: "session_resynced", activeSessionId: "active-1", snapshot, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, cursor: { generation: "generation-active-1", sequence: 13 }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); expect(events).toEqual([ { type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ sessionId: "session-current" }), messages, streamingMessage, lastEventSequence: 13, }), }, ]); }); it("exposes attach snapshots as the initial connection snapshot", async () => { const fakeClient = new FakeDaemonClient(); const snapshotMessage: AgentMessage = { role: "user", content: "snapshot prompt", timestamp: 1 }; const messages: AgentMessage[] = [snapshotMessage]; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 15, { state: createConnectionState(command.activeSessionId, "session-snapshot"), messages, sessionContext: { messages, thinkingLevel: "medium", serviceTier: "default", model: null, }, sessionTree: { tree: [{ entry: projectedUserEntry("user-1", null), children: [] }], leafId: "user-1", }, parent: { activeSessionId: "parent-active", sessionId: "parent-session", nodeId: "parent-node", childId: "child-1", }, }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ state: { activeSessionId: "active-1", sessionId: "session-snapshot", }, messages, sessionContext: { messages, thinkingLevel: "medium", model: null, }, sessionTree: { leafId: "user-1", }, parent: { activeSessionId: "parent-active", sessionId: "parent-session", nodeId: "parent-node", childId: "child-1", }, lastEventSequence: 15, replay: { status: "complete", toSequence: 15, }, }); await expect(connection.getState()).resolves.toMatchObject({ sessionId: "session-snapshot", }); await expect(connection.getMessages()).resolves.toEqual(messages); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach"]); }); it("times out an attach whose streamed snapshot never completes", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.attachResultFactory = (command) => { const result = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12); return { ...result, snapshotStream: { id: "snapshot-stalled", messageCount: 0, targetChunkBytes: 512 * 1024 }, }; }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { snapshotTimeoutMs: 10, }); await expect(connection.attach()).rejects.toThrow("Timed out waiting for snapshot snapshot-stalled"); fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "snapshot-stalled", snapshot: createAttachResult("active-1", "client-1", undefined, 12).snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "snapshot-stalled", chunkCount: 0, lastEventSequence: 12, }); expect((connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size).toBe(0); }); it("rejects one failed snapshot without interrupting another session on the shared client", async () => { const fakeClient = new FakeDaemonClient(); const sibling = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-2"); await sibling.attach(); const siblingEvents: AgentConnectionEvent[] = []; sibling.subscribe((event) => { siblingEvents.push(event); }); fakeClient.attachResultFactory = (command) => { const result = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12); if (command.activeSessionId !== "active-1") { return result; } const { messages: _messages, ...snapshot } = result.snapshot; queueMicrotask(() => { fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: command.activeSessionId, snapshotId: "snapshot-failed", snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, }); fakeClient.emitMessage({ type: "session_snapshot_failed", activeSessionId: command.activeSessionId, snapshotId: "snapshot-failed", error: "snapshot encoder failed", }); }); return { ...result, snapshot: { ...result.snapshot, messages: [] }, snapshotStream: { id: "snapshot-failed", messageCount: 0, targetChunkBytes: 512 * 1024 }, }; }; const failed = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await expect(failed.attach()).rejects.toThrow("snapshot encoder failed"); emitSequencedQueueUpdate(fakeClient, "active-2", 13); await vi.waitFor(() => expect(siblingEvents).toHaveLength(1)); expect(siblingEvents[0]).toMatchObject({ type: "session_event", event: { type: "session_action_update" } }); expect(fakeClient.closeCount).toBe(0); await failed.dispose(); await sibling.dispose(); }); it("assembles chunked attach snapshots even when chunks arrive before the attach response continuation", async () => { const fakeClient = new FakeDaemonClient(); const messages: AgentMessage[] = [ { role: "user", content: "first", timestamp: 1 }, { role: "user", content: "second", timestamp: 2 }, ]; fakeClient.attachResultFactory = (command) => { const full = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 23, { state: createConnectionState(command.activeSessionId, "session-streamed"), messages, }); const { messages: _messages, ...snapshotHeader } = full.snapshot; queueMicrotask(() => { fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: command.activeSessionId, snapshotId: "snapshot-streamed", snapshot: snapshotHeader, messageCount: messages.length, targetChunkBytes: 512 * 1024, }); fakeClient.emitMessage({ type: "session_snapshot_chunk", activeSessionId: command.activeSessionId, snapshotId: "snapshot-streamed", index: 0, messages: [messages[0]!], }); fakeClient.emitMessage({ type: "session_snapshot_chunk", activeSessionId: command.activeSessionId, snapshotId: "snapshot-streamed", index: 1, messages: [messages[1]!], }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: command.activeSessionId, snapshotId: "snapshot-streamed", chunkCount: 2, lastEventSequence: 23, }); }); return { ...full, snapshot: { ...full.snapshot, messages: [] }, snapshotStream: { id: "snapshot-streamed", messageCount: messages.length, targetChunkBytes: 512 * 1024, }, }; }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach(); await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ state: { sessionId: "session-streamed" }, messages, lastEventSequence: 23, }); expect(events).toEqual([]); }); it.each([ ["headless", "message_end"], ["headless", "message_update"], ["reconnect", "message_end"], ["reconnect", "message_update"], ] as const)("preserves %s %s coalesced with attach snapshot completion", async (mode, eventType) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: mode === "reconnect", recoverDaemon: async () => {}, }); const input = new PassThrough(); const detachReader = attachJsonlLineReader(input, (line) => fakeClient.emitMessage(JSON.parse(line))); try { if (mode === "reconnect") { await connection.attach(); await connection.flushBufferedSessionEvents(); } const message = fauxAssistantMessage("new response"); fakeClient.attachResultFactory = (command) => { const full = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { state: { ...createConnectionState("active-1", "session-current"), isStreaming: true }, streamingMessage: fauxAssistantMessage("old partial response"), }); const { messages: _messages, ...snapshot } = full.snapshot; const records: DaemonOutbound[] = [ { type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "coalesced", snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, }, { type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "coalesced", chunkCount: 0, lastEventSequence: 12, lastEventCursor: full.lastEventCursor, }, { type: "session_event", activeSessionId: "active-1", event: eventType === "message_end" ? { type: eventType, message } : { type: eventType, message, assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: "response", partial: message, }, }, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, cursor: { generation: "generation-active-1", sequence: 13 }, emittedAt: "2026-01-01T00:00:00.000Z", }, }, ]; // All records dispatch before the attach response's promise continuation. queueMicrotask(() => input.write(records.map(serializeJsonLine).join(""))); return { ...full, snapshotStream: { id: "coalesced", messageCount: 0, targetChunkBytes: 512 * 1024 } }; }; const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); if (mode === "reconnect") { fakeClient.connected = false; fakeClient.emitClose(new Error("socket closed")); await vi.waitFor(() => expect(events.some((event) => event.type === "session_resynced")).toBe(true)); } else { await connection.attach(); } const snapshot = await connection.getInitialSnapshot(); expect(snapshot.streamingMessage).toEqual(eventType === "message_end" ? undefined : message); expect(snapshot.lastEventCursor).toEqual({ generation: "generation-active-1", sequence: 13 }); expect(fakeClient.requests.map((request) => request.type)).toContain("get_messages"); expect(events).toContainEqual( expect.objectContaining({ type: "session_event", event: expect.objectContaining({ type: eventType }) }), ); for (const event of events) { if (event.type === "session_resynced") expect(event.snapshot.streamingMessage).toEqual(snapshot.streamingMessage); } } finally { detachReader(); input.destroy(); await connection.dispose(); } }); it("distinguishes chunked catch-up snapshots from runtime replacements", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); const emitSnapshot = ( purpose: "replacement" | "resync", snapshotId: string, sequence: number, sessionId: string, ) => { const messages: AgentMessage[] = [{ role: "user", content: purpose, timestamp: sequence }]; const full = createAttachResult("active-1", "client-1", undefined, sequence, { state: createConnectionState("active-1", sessionId), messages, }); const { messages: _messages, ...snapshot } = full.snapshot; fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId, snapshot, messageCount: messages.length, targetChunkBytes: 512 * 1024, purpose, }); fakeClient.emitMessage({ type: "session_snapshot_chunk", activeSessionId: "active-1", snapshotId, index: 0, messages, }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId, chunkCount: 1, lastEventSequence: sequence, lastEventCursor: { generation: "generation-active-1", sequence }, }); }; emitSnapshot("resync", "snapshot-resync", 13, "session-current"); emitSnapshot("replacement", "snapshot-replacement", 14, "session-next"); await vi.waitFor(() => expect(events).toHaveLength(2)); expect(events).toEqual([ expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ sessionId: "session-current" }), }), }), expect.objectContaining({ type: "session_replaced", state: expect.objectContaining({ sessionId: "session-next" }), }), ]); expect((connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size).toBe(0); }); it.each(["replacement", "resync"] as const)( "recovers a failed chunked %s snapshot without interrupting a sibling session", async (purpose) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const sibling = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-2"); await Promise.all([connection.attach(), sibling.attach()]); const events: AgentConnectionEvent[] = []; const siblingEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); sibling.subscribe((event) => { siblingEvents.push(event); }); const recoveredSessionId = purpose === "replacement" ? "session-next" : "session-current"; fakeClient.connectionStateFactory = (activeSessionId) => createConnectionState( activeSessionId, activeSessionId === "active-1" ? recoveredSessionId : "session-sibling", ); fakeClient.requests.length = 0; const snapshotId = `snapshot-failed-${purpose}`; const full = createAttachResult("active-1", "client-1", undefined, 13, { state: createConnectionState("active-1", recoveredSessionId), }); const { messages: _messages, ...snapshot } = full.snapshot; if (purpose === "replacement") { fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", recoveredSessionId), messages: [], snapshotFollows: true, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, cursor: { generation: "generation-active-1", sequence: 13 }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); } fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId, snapshot, messageCount: 1, targetChunkBytes: 512 * 1024, purpose, }); fakeClient.emitMessage({ type: "session_snapshot_failed", activeSessionId: "active-1", snapshotId, error: `${purpose} snapshot failed`, }); await vi.waitFor(() => expect(events).toHaveLength(1)); if (purpose === "replacement") { expect(events[0]).toMatchObject({ type: "session_replaced", state: { sessionId: recoveredSessionId }, messages: [{ role: "user", content: "current prompt", timestamp: 4 }], }); } else { expect(events[0]).toMatchObject({ type: "session_resynced", snapshot: { state: { sessionId: recoveredSessionId }, messages: [{ role: "user", content: "current prompt", timestamp: 4 }], }, }); } expect(fakeClient.requests.map((request) => request.type)).toEqual([ "get_connection_state", "get_messages", "get_session_context", ]); emitSequencedQueueUpdate(fakeClient, "active-2", 13); await vi.waitFor(() => expect(siblingEvents).toHaveLength(1)); expect(siblingEvents[0]).toMatchObject({ type: "session_event", event: { type: "session_action_update" } }); expect((connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size).toBe( 0, ); await connection.dispose(); await sibling.dispose(); }, ); it("keeps attach snapshots usable when the daemon omits duplicate session context", async () => { const fakeClient = new FakeDaemonClient(); const messages: AgentMessage[] = [{ role: "user", content: "snapshot prompt", timestamp: 1 }]; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 15, { state: createConnectionState(command.activeSessionId, "session-snapshot"), messages, omitSessionContext: true, }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const snapshot = await connection.getInitialSnapshot(); expect(snapshot).toMatchObject({ state: { activeSessionId: "active-1", sessionId: "session-snapshot", }, messages, lastEventSequence: 15, }); expect(snapshot.sessionContext).toBeUndefined(); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach"]); await expect(connection.getSessionContext()).resolves.toMatchObject({ messages: [{ role: "user", content: "context prompt", timestamp: 3 }], }); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "get_session_context"]); }); it("keeps reconnect usable from attach snapshots when replay is unavailable", async () => { const fakeClient = new FakeDaemonClient(); const reconnectedMessages: AgentMessage[] = [{ role: "user", content: "reconnected prompt", timestamp: 2 }]; let attachCount = 0; fakeClient.attachResultFactory = (command) => { attachCount++; if (attachCount === 1) { return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { state: createConnectionState(command.activeSessionId, "session-initial"), }); } return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 20, { state: createConnectionState(command.activeSessionId, "session-reconnected"), messages: reconnectedMessages, sessionContext: { messages: reconnectedMessages, thinkingLevel: "medium", serviceTier: "default", model: null, }, replay: { status: "unavailable", fromSequence: 14, toSequence: 20, reason: "event_replay_not_available", }, }); }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); emitSequencedQueueUpdate(fakeClient, "active-1", 14); await connection.attach(); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-1", capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], resumeCursor: { activeSessionId: "active-1", generation: "generation-active-1", sequence: 14, }, }); await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ state: { sessionId: "session-reconnected", }, messages: reconnectedMessages, replay: { status: "unavailable", fromSequence: 14, toSequence: 20, reason: "event_replay_not_available", }, lastEventSequence: 20, }); await expect(connection.getState()).resolves.toMatchObject({ sessionId: "session-reconnected", }); await expect(connection.getMessages()).resolves.toEqual(reconnectedMessages); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "attach"]); }); it("resets a connected transport when reattach fails during supervisor recovery", async () => { const fakeClient = new FakeDaemonClient(); const recoverDaemon = vi.fn(async () => undefined); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { recoverDaemon, reconnectTimeoutMs: 2000, }); const statuses: string[] = []; let resyncs = 0; connection.subscribe((event) => { if (event.type === "connection_status") { statuses.push(event.status); } if (event.type === "session_resynced") { resyncs++; } }); await connection.attach(); fakeClient.attachFailures = 1; fakeClient.connected = false; fakeClient.emitClose(new Error("Daemon socket closed")); await vi.waitFor(() => expect(statuses).toEqual(["reconnecting", "connected"])); expect(recoverDaemon).toHaveBeenCalledTimes(2); expect(fakeClient.reconnectCount).toBe(2); expect(fakeClient.resetTransportCount).toBe(1); expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(3); expect(resyncs).toBe(1); }); it.each([ ["reconnect", "attach"], ["reconnect", "snapshot"], ["reconnect", "backoff"], ["update", "attach"], ["update", "snapshot"], ["update", "backoff"], ] as const)("stops %s after a terminal close during %s", async (recovery, stage) => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 1000, }); let release!: () => void; const gate = new Promise((resolve) => { release = resolve; }); try { await connection.attach(); fakeClient.updateRestartSessions = [{ activeSessionId: "active-1", sessionId: "session-current" }]; const request = fakeClient.request.bind(fakeClient); vi.spyOn(fakeClient, "request").mockImplementation(async (...args) => { const response = await request(...args); if (args[0].type === "attach") { if (stage === "backoff") throw new Error("attach temporarily unavailable"); if (stage === "attach") await gate; } return response; }); const getInitialSnapshot = connection.getInitialSnapshot.bind(connection); const snapshotRead = vi.spyOn(connection, "getInitialSnapshot").mockImplementation(async (...args) => { const snapshot = await getInitialSnapshot(...args); if (stage === "snapshot") await gate; return snapshot; }); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); fakeClient.connected = false; fakeClient.emitClose( recovery === "update" ? new DaemonSocketClosedError("/tmp/fake.sock", "update") : new Error("Daemon socket closed"), ); await vi.advanceTimersByTimeAsync(0); expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(2); if (stage === "snapshot") expect(snapshotRead).toHaveBeenCalledOnce(); const requestCount = fakeClient.requests.length; const reconnectCount = fakeClient.reconnectCount; const resetCount = fakeClient.resetTransportCount; const closeCount = fakeClient.closeCount; fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-1", reason: "killed" }); release(); await vi.advanceTimersByTimeAsync(120100); expect(fakeClient.requests).toHaveLength(requestCount); expect(fakeClient.reconnectCount).toBe(reconnectCount); expect(fakeClient.resetTransportCount).toBe(resetCount); expect(fakeClient.closeCount).toBe(closeCount); expect(events).toEqual([ expect.objectContaining({ type: "connection_status", status: "reconnecting" }), expect.objectContaining({ type: "closed", error: expect.stringContaining("The daemon stopped this agent session."), }), ]); } finally { release(); await connection.dispose(); vi.useRealTimers(); } }); it("does not reconnect after disposal while daemon recovery is pending", async () => { const fakeClient = new FakeDaemonClient(); let finishRecovery: (() => void) | undefined; const recoverDaemon = vi.fn( () => new Promise((resolve) => { finishRecovery = resolve; }), ); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { recoverDaemon, reconnectTimeoutMs: 2000, }); await connection.attach(); fakeClient.connected = false; fakeClient.emitClose(new Error("Daemon socket closed")); await vi.waitFor(() => expect(recoverDaemon).toHaveBeenCalledOnce()); await connection.dispose(); finishRecovery?.(); for (let flush = 0; flush < 5; flush++) { await Promise.resolve(); } expect(fakeClient.reconnectCount).toBe(0); }); it("isolates subscriber failures during transport recovery", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 2000, }); const statuses: string[] = []; connection.subscribe(async () => { throw new Error("broken subscriber"); }); connection.subscribe((event) => { if (event.type === "connection_status") { statuses.push(event.status); } }); await connection.attach(); fakeClient.connected = false; fakeClient.emitClose(new Error("Daemon socket closed")); await vi.waitFor(() => expect(statuses).toEqual(["reconnecting", "connected"])); expect(fakeClient.reconnectCount).toBe(1); }); it("does not let a stalled subscriber block update recovery", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const statuses: string[] = []; connection.subscribe(() => new Promise(() => undefined)); connection.subscribe((event) => { if (event.type === "connection_status") { statuses.push(event.status); } }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.waitFor(() => expect(statuses).toEqual(["reconnecting", "connected"])); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-restored", }); }); it("preserves a newer live child update over an in-flight roster read", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("authoritative_child_roster"); let releaseRoster!: () => void; fakeClient.rlmChildrenGate = new Promise((resolve) => { releaseRoster = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const children = connection.getRlmChildSnapshots(); emitRlmChildUpdate(fakeClient, "active-1", 13, { id: "child-live", label: "live child", status: "running", sessionDir: "/tmp/child-live", }); releaseRoster(); await expect(children).resolves.toEqual([expect.objectContaining({ id: "child-live", status: "running" })]); }); it("refreshes initial snapshots after live events make the cached snapshot stale", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { state: createConnectionState(command.activeSessionId, "session-attached"), messages: [{ role: "user", content: "attached prompt", timestamp: 1 }], }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); emitSequencedQueueUpdate(fakeClient, "active-1", 13); const snapshot = await connection.getInitialSnapshot(); expect(snapshot).toMatchObject({ state: { sessionId: "session-current", }, messages: [{ role: "user", content: "current prompt", timestamp: 4 }], sessionContext: { messages: [{ role: "user", content: "context prompt", timestamp: 3 }], }, }); // The session tree is fetched lazily (only when the tree/branch selector is // opened), so refreshing the initial snapshot must not request it. expect(snapshot.sessionTree).toBeUndefined(); expect(fakeClient.requests.map((request) => request.type)).toEqual([ "attach", "get_connection_state", "get_messages", "get_session_context", ]); }); it("ignores older sequenced events after an attach snapshot", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach(); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "session-old"), messages: [{ role: "user", content: "old prompt", timestamp: 1 }], meta: { id: "active-1:10", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 10, emittedAt: "2026-01-01T00:00:00.000Z", }, }); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "session-new"), messages: [{ role: "user", content: "new prompt", timestamp: 2 }], meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, emittedAt: "2026-01-01T00:00:00.000Z", }, }); expect(events).toEqual([ { type: "session_replaced", state: expect.objectContaining({ sessionId: "session-new", }), messages: [{ role: "user", content: "new prompt", timestamp: 2 }], }, ]); await expect(connection.getState()).resolves.toMatchObject({ sessionId: "session-new", }); }); it("fails closed when a daemon disconnect invalidates an input pause", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("session_input_pause"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); const pause = await connection.acquireSessionInputPause("lease-1"); fakeClient.disconnectForReconnect("shutdown"); await expect(pause.release()).rejects.toThrow("invalidated by a daemon reconnect"); await expect(connection.acquireSessionInputPause("lease-1")).rejects.toThrow("connection is closed"); expect(events).toContainEqual( expect.objectContaining({ type: "closed", error: expect.stringContaining("fence was invalidated"), }), ); }); it("releases an input pause whose acquisition resolves after disconnect", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("session_input_pause"); let releaseAcquire!: () => void; fakeClient.inputPauseAcquireGate = new Promise((resolve) => { releaseAcquire = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const acquisition = connection.acquireSessionInputPause("lease-1"); await vi.waitFor(() => expect(fakeClient.requests.some((request) => request.type === "acquire_session_input_pause")).toBe(true), ); fakeClient.disconnectForReconnect("shutdown"); releaseAcquire(); await expect(acquisition).rejects.toThrow("acquisition was invalidated by a daemon reconnect"); expect(fakeClient.requests.filter((request) => request.type === "release_session_input_pause")).toHaveLength(1); }); it("asks for the projected session tree when the daemon supports it", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("session_tree_projection"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); emitStaleSessionActionUpdate(fakeClient); await expect(connection.getSessionTree()).resolves.toEqual({ tree: [{ entry: projectedUserEntry("user-1", null), children: [] }], leafId: "user-1", }); expect(fakeClient.requests[1]).toMatchObject({ type: "get_session_tree", activeSessionId: "active-1", projection: "tree", }); }); it("projects the session tree locally when the daemon lacks the projection", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); emitStaleSessionActionUpdate(fakeClient); await expect(connection.getSessionTree()).resolves.toEqual({ tree: [{ entry: projectedUserEntry("user-1", null), children: [] }], leafId: "user-1", }); expect(fakeClient.requests[1]).toMatchObject({ type: "get_session_tree", activeSessionId: "active-1" }); expect(fakeClient.requests[1]).not.toHaveProperty("projection"); }); it("capability-gates published extension command argument completions", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); // An older daemon may still send the field; the client must not complete from it. const withoutCapability = await connection.getCommands(); expect(withoutCapability[0]?.name).toBe("route"); expect(withoutCapability[0]?.argumentCompletions).toBeUndefined(); fakeClient.serverCapabilities.add("command_argument_completions"); const withCapability = await connection.getCommands(); expect(withCapability[0]?.argumentCompletions).toEqual([ { value: "policy ", label: "policy", description: "Follow a policy" }, ]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "get_commands", activeSessionId: "active-1" }); }); it("ignores delayed events from a retired daemon generation", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const steeringUpdates: Array = []; connection.subscribe((event) => { if (event.type === "session_event" && event.event.type === "session_action_update") { steeringUpdates.push(event.event.actions.steering); } }); await connection.attach(); const emitQueue = (generation: string, sequence: number, steering: string) => { fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "session_action_update", actions: { queuedCount: 1, steering: [steering], followUps: [] } }, meta: { id: `${generation}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence, cursor: { generation, sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); }; emitQueue("generation-new", 1, "new"); emitQueue("generation-active-1", 13, "old"); await vi.waitFor(() => expect(steeringUpdates).toEqual([["new"]])); await connection.attach(); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", resumeCursor: { generation: "generation-new", sequence: 1 }, }); }); }); const DEFERRAL_EVENT_BASE_SEQUENCE = 13; function emitSequencedSessionEvent(client: FakeDaemonClient, activeSessionId: string, sequence: number): void { client.emitMessage({ type: "session_event", activeSessionId, event: { type: "session_info_changed", name: String(sequence) }, meta: { id: `${activeSessionId}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId, sequence, cursor: { generation: `generation-${activeSessionId}`, sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, } satisfies DaemonOutbound); } async function nextMessageLoopTurn(): Promise { await new Promise((resolve) => setImmediate(resolve)); await new Promise((resolve) => setImmediate(resolve)); } describe("DaemonAgentConnection deferred session events", () => { it("defers events between attach and flush so the attach snapshot stays usable", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE); await nextMessageLoopTurn(); expect(delivered).toEqual([]); fakeClient.requests.length = 0; const snapshot = await connection.getInitialSnapshot(); expect(fakeClient.requests.map((request) => request.type)).not.toContain("get_messages"); expect(snapshot.state.activeSessionId).toBe("active-1"); await connection.flushBufferedSessionEvents(); expect(delivered).toHaveLength(1); expect(delivered[0]).toMatchObject({ type: "session_event" }); emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE + 1); await nextMessageLoopTurn(); expect(delivered).toHaveLength(2); } finally { await connection.dispose(); } }); it("delivers events live and refetches when deferral is not opted in", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); try { await connection.attach(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE); await nextMessageLoopTurn(); expect(delivered).toHaveLength(1); fakeClient.requests.length = 0; await connection.getInitialSnapshot(); expect(fakeClient.requests.map((request) => request.type)).toContain("get_messages"); } finally { await connection.dispose(); } }); it("drops deferred events a resync snapshot already contains", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); const resyncSequence = DEFERRAL_EVENT_BASE_SEQUENCE + 2; emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE); emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE + 1); fakeClient.emitMessage({ type: "session_resynced", activeSessionId: "active-1", snapshot: createAttachResult("active-1", undefined, undefined, resyncSequence).snapshot, meta: { id: `active-1:${resyncSequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: resyncSequence, cursor: { generation: "generation-active-1", sequence: resyncSequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, } satisfies DaemonOutbound); await nextMessageLoopTurn(); expect(delivered.filter((event) => event.type === "session_resynced")).toHaveLength(1); expect(delivered.filter((event) => event.type === "session_event")).toEqual([]); await connection.flushBufferedSessionEvents(); expect(delivered.filter((event) => event.type === "session_event")).toEqual([]); emitSequencedSessionEvent(fakeClient, "active-1", resyncSequence + 1); await connection.flushBufferedSessionEvents(); expect(delivered.filter((event) => event.type === "session_event")).toHaveLength(1); } finally { await connection.dispose(); } }); it("resyncs an already rendered snapshot when deferral overflows", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); const renderedSnapshot = await connection.getInitialSnapshot(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); // The initial renderer already has its snapshot when the buffer overflows. for (let index = 0; index <= 1000; index++) { emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE + index); } await nextMessageLoopTurn(); expect(delivered).toHaveLength(0); const recoveredMessages: AgentMessage[] = [{ role: "user", content: "recovered", timestamp: 1 }]; fakeClient.attachResultFactory = (command) => { emitSequencedSessionEvent(fakeClient, "active-1", 1014); return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1013, { messages: recoveredMessages, }); }; await connection.flushBufferedSessionEvents(); expect(renderedSnapshot.messages).toEqual([]); expect(delivered).toEqual([ expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ messages: recoveredMessages }), }), { type: "session_event", event: { type: "session_info_changed", name: "1014" } }, ]); emitSequencedSessionEvent(fakeClient, "active-1", 1015); await nextMessageLoopTurn(); expect(delivered).toHaveLength(3); } finally { await connection.dispose(); } }); it.each([true, false])( "finishes overflow recovery under continuous traffic (snapshot covers overflow=%s)", async (covered) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); let release!: () => void; const rendering = new Promise((resolve) => { release = resolve; }); try { await connection.attach(); let sequence = 12; for (let index = 0; index <= 1000; index++) emitSequencedSessionEvent(fakeClient, "active-1", ++sequence); let attaches = 0; fakeClient.attachResultFactory = (command) => { if (++attaches > 3) throw new Error("overflow recovery did not terminate"); const before = sequence; if (covered || attaches === 1) { for (let index = 0; index <= 1000; index++) emitSequencedSessionEvent(fakeClient, "active-1", ++sequence); } return createAttachResult( command.activeSessionId, command.clientId, command.capabilities, !covered && attaches === 1 ? before : sequence, ); }; const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); if (event.type === "session_resynced") return rendering; }); const flush = connection.flushBufferedSessionEvents(); await nextMessageLoopTurn(); expect(attaches).toBe(covered ? 1 : 2); expect(delivered.at(-1)).toMatchObject({ type: "session_resynced", snapshot: { lastEventSequence: sequence }, }); const beforeLive = delivered.length; for (let index = 0; index < 2000; index++) emitSequencedSessionEvent(fakeClient, "active-1", ++sequence); expect(delivered.slice(beforeLive)).toHaveLength(2000); release(); await flush; expect(attaches).toBe(covered ? 1 : 2); expect(delivered.some((event) => event.type === "closed")).toBe(false); } finally { release(); await connection.dispose(); } }, ); it.each([ ["replacement", false], ["replacement", true], ["streamed replacement", true], ["reattach", false], ["reattach", true], ] as const)("discards overflow recovery across %s (recovery started=%s)", async (change, recoveryStarted) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); let release!: () => void; const gate = new Promise((resolve) => { release = resolve; }); try { await connection.attach(); for (let sequence = 13; sequence <= 1013; sequence++) emitSequencedSessionEvent(fakeClient, "active-1", sequence); const request = fakeClient.request.bind(fakeClient); vi.spyOn(fakeClient, "request").mockImplementation(async (command, ...options) => { if (command.type === "attach" && change === "streamed replacement") { const result = createAttachResult("active-1", undefined, undefined, 1013); result.snapshotStream = { id: "old-recovery", messageCount: 0, targetChunkBytes: 512 * 1024 }; fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "old-recovery", snapshot: result.snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, purpose: "attach", }); return { type: "response", command: command.type, success: true, data: result }; } if (command.type === "attach") await gate; if (command.type === "switch_session") return { type: "response", command: command.type, success: false, error: "Session already active", errorInfo: { code: "session_already_active", activeSessionId: "active-2", sessionPath: "/tmp/target.jsonl", }, }; if (command.type === "reattach") return { type: "response", command: command.type, success: true, data: createAttachResult("active-2", undefined, undefined, 1), }; return request(command, ...options); }); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); const flush = recoveryStarted ? connection.flushBufferedSessionEvents() : undefined; await nextMessageLoopTurn(); if (change === "reattach") { await connection.switchSession("/tmp/target.jsonl"); } else { fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "replacement"), messages: [], }); } const activeSessionId = change === "reattach" ? "active-2" : "active-1"; const sequence = change === "reattach" ? 2 : 1014; emitSequencedSessionEvent(fakeClient, activeSessionId, sequence); if (change === "streamed replacement") { fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "old-recovery", chunkCount: 0, lastEventSequence: 1013, lastEventCursor: { generation: "generation-active-1", sequence: 1013 }, }); } release(); await (flush ?? connection.flushBufferedSessionEvents()); expect(delivered.map((event) => event.type)).toEqual(["session_replaced", "session_event"]); expect(delivered[1]).toEqual({ type: "session_event", event: { type: "session_info_changed", name: String(sequence) }, }); if (change !== "reattach") expect( (connection as unknown as { latestSnapshot: { state: AgentConnectionState } }).latestSnapshot.state .sessionId, ).toBe("replacement"); expect((await connection.getState()).activeSessionId).toBe(activeSessionId); } finally { release(); await connection.dispose(); } }); it("does not advance event cursors for an attach superseded before its snapshot finishes", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); fakeClient.attachResultFactory = (command) => ({ ...createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 100), snapshotStream: { id: "superseded", messageCount: 0, targetChunkBytes: 512 * 1024 }, }); const pendingAttach = connection.attach(); await nextMessageLoopTurn(); fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "superseded", snapshot: createAttachResult("active-1", undefined, undefined, 100).snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, purpose: "attach", }); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "replacement"), messages: [], }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "superseded", chunkCount: 0, lastEventSequence: 100, lastEventCursor: { generation: "generation-active-1", sequence: 100 }, }); await pendingAttach; const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); emitSequencedSessionEvent(fakeClient, "active-1", 13); await connection.flushBufferedSessionEvents(); expect(delivered).toEqual([{ type: "session_event", event: { type: "session_info_changed", name: "13" } }]); const snapshot = await connection.getInitialSnapshot(); expect(snapshot.lastEventCursor).toEqual({ generation: "generation-active-1", sequence: 13 }); } finally { await connection.dispose(); } }); it.each([0, 2, 3])( "releases a superseded attach with %i snapshot frames received before its response", async (frameCount) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); let release!: () => void; const responseGate = new Promise((resolve) => { release = resolve; }); try { await connection.attach(); const result = createAttachResult("active-1", undefined, undefined, 100); fakeClient.attachResultFactory = () => ({ ...result, snapshotStream: { id: "abandoned", messageCount: 1, targetChunkBytes: 512 * 1024 }, }); const request = fakeClient.request.bind(fakeClient); vi.spyOn(fakeClient, "request").mockImplementation(async (...args) => { const response = await request(...args); if (args[0].type === "attach") await responseGate; return response; }); const frames: DaemonOutbound[] = [ { type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "abandoned", snapshot: result.snapshot, messageCount: 1, targetChunkBytes: 512 * 1024, purpose: "attach", }, { type: "session_snapshot_chunk", activeSessionId: "active-1", snapshotId: "abandoned", index: 0, messages: [{ role: "user", content: "old transcript", timestamp: 1 }], }, { type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "abandoned", chunkCount: 1, lastEventSequence: 100, }, ]; const pendingAttach = connection.attach(); for (const frame of frames.slice(0, frameCount)) fakeClient.emitMessage(frame); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "replacement"), messages: [], }); release(); await pendingAttach; expect( (connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size, ).toBe(0); for (const frame of frames.slice(frameCount)) fakeClient.emitMessage(frame); await nextMessageLoopTurn(); expect( (connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size, ).toBe(0); expect((await connection.getInitialSnapshot()).state.sessionId).toBe("replacement"); } finally { release(); await connection.dispose(); } }, ); it.each(["before response", "during stream", "after end"] as const)( "rejects a closed attach immediately %s", async (timing) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const result = createAttachResult("active-1", undefined, undefined, 12); fakeClient.attachResultFactory = () => ({ ...result, snapshotStream: { id: "closed", messageCount: 0, targetChunkBytes: 512 * 1024 }, }); const rejected = vi.fn(); const resolved = vi.fn(); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); const pendingAttach = connection.attach().then(resolved, rejected); try { if (timing !== "before response") await nextMessageLoopTurn(); fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "closed", snapshot: result.snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, purpose: "attach", }); if (timing === "after end") { fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "closed", chunkCount: 0, lastEventSequence: 12, }); } fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-1", reason: "killed" }); await nextMessageLoopTurn(); expect(rejected).toHaveBeenCalledExactlyOnceWith(expect.any(Error)); expect(resolved).not.toHaveBeenCalled(); expect( (connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size, ).toBe(0); for (const purpose of ["attach", "replacement", "resync"] as const) { fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: purpose, snapshot: result.snapshot, messageCount: 1, targetChunkBytes: 512 * 1024, purpose, }); fakeClient.emitMessage({ type: "session_snapshot_chunk", activeSessionId: "active-1", snapshotId: purpose, index: 0, messages: [{ role: "user", content: "late transcript", timestamp: 1 }], }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: purpose, chunkCount: 1, lastEventSequence: 12, }); } fakeClient.emitMessage({ type: "session_snapshot_failed", activeSessionId: "active-1", snapshotId: "failed-after-close", error: "stream closed", }); expect( (connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size, ).toBe(0); fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-1", reason: "killed" }); await nextMessageLoopTurn(); expect( (connection as unknown as { snapshotAssemblies: Map }).snapshotAssemblies.size, ).toBe(0); expect(events).toEqual([expect.objectContaining({ type: "closed" })]); const requestCount = fakeClient.requests.length; await expect(connection.attach()).rejects.toThrow("closed"); expect(fakeClient.requests).toHaveLength(requestCount); } finally { await connection.dispose(); await pendingAttach; } }, ); it.each(["session_replaced", "session_resynced"] as const)( "drops a %s superseded during the initial render", async (type) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); let release!: () => void; const initialRenderPromise = new Promise((resolve) => { release = resolve; }); const ui = { agentConnection: connection, sessionEventQueue: Promise.resolve(), sessionEventGeneration: 0, initialRenderPromise, refreshCommandCatalogForCurrentSession: vi.fn(async () => {}), renderResyncedSession: vi.fn(async () => {}), resetSideQuestion: vi.fn(), resetExtensionUI: vi.fn(), applyConnectionStateSnapshot: vi.fn(), resetCurrentSessionRenderState: vi.fn(), rebindCurrentSession: vi.fn(async () => {}), renderInitialMessages: vi.fn(async () => {}), ui: { requestRender: vi.fn() }, showError: vi.fn(), }; try { await connection.attach(); (InteractiveMode.prototype as unknown as { subscribeToAgent(this: typeof ui): void }).subscribeToAgent.call( ui, ); const snapshot = createAttachResult("active-1", undefined, undefined, 12).snapshot; fakeClient.emitMessage( type === "session_replaced" ? { type, activeSessionId: "active-1", state: snapshot.state, messages: [] } : { type, activeSessionId: "active-1", snapshot }, ); await nextMessageLoopTurn(); const state = createConnectionState("active-1", "latest"); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state, messages: [] }); release(); await ui.sessionEventQueue; expect(ui.renderResyncedSession).not.toHaveBeenCalled(); expect(ui.resetSideQuestion).toHaveBeenCalledOnce(); expect(ui.resetExtensionUI).toHaveBeenCalledOnce(); expect(ui.applyConnectionStateSnapshot).toHaveBeenCalledExactlyOnceWith(state); expect(ui.resetCurrentSessionRenderState).toHaveBeenCalledOnce(); expect(ui.rebindCurrentSession).toHaveBeenCalledOnce(); expect(ui.renderInitialMessages).toHaveBeenCalledOnce(); expect(ui.ui.requestRender).toHaveBeenCalledOnce(); expect(ui.showError).not.toHaveBeenCalled(); } finally { release(); await connection.dispose(); } }, ); it("keeps live events behind the whole replay and shares concurrent flushes", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); let release!: () => void; const gate = new Promise((resolve) => { release = resolve; }); try { await connection.attach(); const delivered: string[] = []; const ui = { agentConnection: connection, sessionEventQueue: Promise.resolve(), sessionEventGeneration: 0, handleEvent: async (event: { type: string; name?: string }) => { delivered.push(event.name!); if (delivered.length === 1) await gate; }, showError: vi.fn(), }; (InteractiveMode.prototype as unknown as { subscribeToAgent(this: typeof ui): void }).subscribeToAgent.call( ui, ); emitSequencedSessionEvent(fakeClient, "active-1", 13); emitSequencedSessionEvent(fakeClient, "active-1", 14); const flush = connection.flushBufferedSessionEvents(); await nextMessageLoopTurn(); emitSequencedSessionEvent(fakeClient, "active-1", 15); expect(connection.flushBufferedSessionEvents()).toBe(flush); expect(delivered).toEqual(["13"]); release(); await flush; await ui.sessionEventQueue; expect(delivered).toEqual(["13", "14", "15"]); expect(ui.showError).not.toHaveBeenCalled(); } finally { release(); await connection.dispose(); } }); it("leaves streaming state in the attach snapshot unchanged until replay", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); const message = fauxAssistantMessage("streaming response"); fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "message_start", message }, }); const snapshot = await connection.getInitialSnapshot(); expect(snapshot.streamingMessage).toBeUndefined(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); await connection.flushBufferedSessionEvents(); expect(delivered).toEqual([{ type: "session_event", event: { type: "message_start", message } }]); expect((await connection.getInitialSnapshot()).streamingMessage).toEqual(message); } finally { await connection.dispose(); } }); it.each([true, false])( "delivers attach-time extension prompts without skipping transcript events (subscribed=%s)", async (subscribed) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { const delivered: AgentConnectionEvent[] = []; const listener = (event: AgentConnectionEvent) => { delivered.push(event); }; if (subscribed) connection.subscribe(listener); fakeClient.attachResultFactory = (command) => { fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "editor-request", method: "editor", payload: { title: "Edit" }, meta: { id: "active-1:14", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 14, cursor: { generation: "generation-active-1", sequence: 14 }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12); }; await connection.attach(); // The transport releases earlier transcript events only after the snapshot. emitSequencedSessionEvent(fakeClient, "active-1", 13); if (!subscribed) { expect(delivered).toEqual([]); connection.subscribe(listener); } expect(delivered).toEqual([ { type: "extension_ui_request", request: { id: "editor-request", method: "editor", payload: { title: "Edit" } }, }, ]); await connection.flushBufferedSessionEvents(); expect(delivered.at(-1)).toEqual({ type: "session_event", event: { type: "session_info_changed", name: "13" }, }); const additionalListener = vi.fn(); connection.subscribe(additionalListener); expect(additionalListener).not.toHaveBeenCalled(); } finally { await connection.dispose(); } }, ); it.each(["notify", "setStatus", "setWidget"])( "bounds attach-time %s updates without dropping queued dialogs", async (method) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); try { await connection.attach(); fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "dialog", method: "editor", payload: { title: "Edit" }, }); for (let index = 0; index < 2000; index++) { fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: String(index), method, payload: { message: String(index), statusKey: "progress", statusText: String(index), widgetKey: "progress", widgetLines: [String(index)], }, }); } const listener = vi.fn(); connection.subscribe(listener); expect(listener).toHaveBeenCalledTimes(128); expect(listener).toHaveBeenNthCalledWith(1, { type: "extension_ui_request", request: { id: "dialog", method: "editor", payload: { title: "Edit" } }, }); expect(listener).toHaveBeenLastCalledWith({ type: "extension_ui_request", request: expect.objectContaining({ id: "1999", method }), }); expect(fakeClient.requests.filter((request) => request.type === "extension_ui_response")).toEqual([]); } finally { await connection.dispose(); } }, ); it.each(["disposed", "replaced", "restarted", "update-session", "update-transport"])( "discards queued extension prompts when %s before subscription", async (boundary) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); try { await connection.attach(); fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "old", method: "editor", payload: { title: "Old" }, meta: { id: "old:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, cursor: { generation: "generation-active-1", sequence: 13 }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); if (boundary === "disposed") await connection.dispose(); else if (boundary === "replaced") fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "replacement"), messages: [], }); else if (boundary.startsWith("update-")) { fakeClient.updateRestartSessions = [{ activeSessionId: "active-1", sessionId: "session-current" }]; if (boundary === "update-session") { fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-1", reason: "update" }); } else { fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); } } else { const snapshot = createAttachResult("active-1", undefined, undefined, 1).snapshot; snapshot.lastEventCursor = { generation: "restarted", sequence: 1 }; fakeClient.emitMessage({ type: "session_resynced", activeSessionId: "active-1", snapshot }); } let resolveConnected!: () => void; const connected = new Promise((resolve) => { resolveConnected = resolve; }); const listener = vi.fn((event: AgentConnectionEvent) => { if (event.type === "connection_status" && event.status === "connected") resolveConnected(); }); connection.subscribe(listener); expect(listener).not.toHaveBeenCalled(); if (boundary.startsWith("update-")) { await connected; expect(listener).toHaveBeenCalledWith({ type: "connection_status", status: "connected" }); expect(listener).not.toHaveBeenCalledWith(expect.objectContaining({ type: "extension_ui_request" })); fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "new", method: "editor", payload: { title: "New" }, }); expect(listener).toHaveBeenLastCalledWith({ type: "extension_ui_request", request: { id: "new", method: "editor", payload: { title: "New" } }, }); } } finally { await connection.dispose(); } }, ); it.each([1, 1001])( "drops %i buffered events and overflow superseded by a restarted worker snapshot", async (count) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); for (let index = 0; index < count; index++) emitSequencedSessionEvent(fakeClient, "active-1", 100 + index); const snapshot = createAttachResult("active-1", undefined, undefined, 1).snapshot; snapshot.lastEventCursor = { generation: "restarted", sequence: 1 }; fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "restart", snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, purpose: "resync", }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "restart", chunkCount: 0, lastEventSequence: 1, lastEventCursor: snapshot.lastEventCursor, }); await nextMessageLoopTurn(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); await connection.flushBufferedSessionEvents(); expect(delivered).toEqual([]); expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(1); expect((await connection.getInitialSnapshot()).lastEventSequence).toBe(1); } finally { await connection.dispose(); } }, ); it("replays only newer target-session events when reattaching with an inline snapshot", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); emitSequencedSessionEvent(fakeClient, "active-1", 13); const request = fakeClient.request.bind(fakeClient); vi.spyOn(fakeClient, "request").mockImplementation(async (command, ...options) => { if (command.type === "switch_session") return { type: "response", command: command.type, success: false, error: "Session already active", errorInfo: { code: "session_already_active", activeSessionId: "active-2", sessionPath: "/tmp/target.jsonl", }, }; if (command.type === "reattach") { emitSequencedSessionEvent(fakeClient, "active-2", 2); return { type: "response", command: command.type, success: true, data: createAttachResult("active-2", undefined, undefined, 1), }; } return request(command, ...options); }); await connection.switchSession("/tmp/target.jsonl"); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); await connection.flushBufferedSessionEvents(); expect(delivered).toEqual([{ type: "session_event", event: { type: "session_info_changed", name: "2" } }]); } finally { await connection.dispose(); } }); });