From 5a49ecd858fac5e4d826ea723d2b3c4e453a35e9 Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Thu, 17 Sep 2026 13:18:37 +0300 Subject: [PATCH] fix(coding-agent): handle session remap and initial attach during daemon restarts --- .../daemon-agent-connection.ts | 207 ++++++++++-------- .../src/modes/agent-connection/types.ts | 7 +- .../src/modes/daemon/daemon-errors.ts | 43 ++++ .../src/modes/daemon/daemon-mode.ts | 2 +- .../src/modes/interactive/interactive-mode.ts | 6 + .../test/agent-connection-daemon.test.ts | 97 ++++++++ 6 files changed, 274 insertions(+), 88 deletions(-) diff --git a/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts index 2e9b9b085..7477c8d05 100644 --- a/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts +++ b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts @@ -30,7 +30,12 @@ import { isDaemonNotConnectedError, isDaemonSocketTransientError, } from "../daemon/daemon-client.js"; -import { deserializeDaemonError, isSessionStartingOrRecoveringError } from "../daemon/daemon-errors.js"; +import { + DaemonSessionStartingError, + deserializeDaemonError, + isSessionNotReadyError, + isSessionStartingOrRecoveringError, +} from "../daemon/daemon-errors.js"; import { collectDaemonClientEnv, collectDaemonLaunchEnv, @@ -50,6 +55,7 @@ import { DaemonDirectTransportClosedError, DaemonRoutedClient, } from "../daemon/daemon-routed-client.js"; +import { getDaemonRuntimeIdentity } from "../daemon/daemon-runtime-identity.js"; import type { SessionSummary } from "../daemon/daemon-session-list.js"; import { listDaemonHeartbeats } from "../daemon/heartbeat-catalog.js"; import { @@ -144,7 +150,10 @@ function formatErrorSentence(error: unknown): string { return /[.!?]$/.test(message) ? message : `${message}.`; } -function reconnectDaemonTransportAfterUpdate(client: DaemonTransportClient): Promise { +function reconnectDaemonTransportAfterUpdate( + client: DaemonTransportClient, + timeoutMs = UPDATE_RECONNECT_TIMEOUT_MS, +): Promise { const existing = updateTransportReconnects.get(client); if (existing) { return existing; @@ -152,7 +161,7 @@ function reconnectDaemonTransportAfterUpdate(client: DaemonTransportClient): Pro const reconnectPromise = Promise.resolve() .then(async () => { client.disconnectForReconnect("update"); - const deadline = Date.now() + UPDATE_RECONNECT_TIMEOUT_MS; + const deadline = Date.now() + timeoutMs; let lastError: unknown; while (Date.now() < deadline) { try { @@ -249,6 +258,18 @@ export class DaemonAgentConnection implements AgentConnection { private updateRestartPending = false; private updateReconnectFailed = false; private terminalCloseEmitted = false; + private notifiedBuildIdMismatch = false; + + private emitConnectedStatus(): void { + const daemonBuildId = this.client.hello?.runtime?.buildId; + const clientBuildId = getDaemonRuntimeIdentity().buildId; + const buildIdMismatch = + daemonBuildId && clientBuildId && daemonBuildId !== clientBuildId && !this.notifiedBuildIdMismatch + ? { client: clientBuildId, daemon: daemonBuildId } + : undefined; + if (buildIdMismatch) this.notifiedBuildIdMismatch = true; + void this.emit({ type: "connection_status", status: "connected", buildIdMismatch }); + } private updateReconnectPromise?: Promise; private readonly activeSideQuestionIds = new Set(); private readonly snapshotAssemblies = new Map(); @@ -354,36 +375,48 @@ export class DaemonAgentConnection implements AgentConnection { ); const connection = new DaemonAgentConnection(transport, activeSessionId, options); connection.initialAttachPending = true; - try { + const deadline = Date.now() + (options?.reconnectTimeoutMs ?? DAEMON_RECONNECT_TIMEOUT_MS); + let attempt = 0; + let lastError: unknown; + while (Date.now() < deadline && !connection.disposed && !connection.terminalCloseEmitted) { try { - await connection.attach(); - } catch (error) { - if (!(transport instanceof DaemonRoutedClient)) throw error; - transport.fallbackToSupervisor(); try { - // This retry owns its failure; a parked request would pend the attach forever. - await connection.attach({ recoverable: false }); - } catch (retryError) { - // A control-plane close saved during the window is the authoritative cause. - throw connection.initialControlPlaneClose ?? retryError; + await connection.attach(); + } catch (_error) { + if (transport instanceof DaemonRoutedClient) { + transport.fallbackToSupervisor(); + } + try { + // This retry owns its failure; a parked request would pend the attach forever. + await connection.attach({ recoverable: false }); + } catch (retryError) { + // A control-plane close saved during the window is the authoritative cause. + throw connection.initialControlPlaneClose ?? retryError; + } } - } - connection.initialAttachPending = false; - const initialControlPlaneClose = connection.initialControlPlaneClose; - connection.initialControlPlaneClose = undefined; - if (initialControlPlaneClose) { - // No listeners exist yet: a terminal close rejects the attach; the rest replays through the one handler. - if (getDaemonSocketCloseReason(initialControlPlaneClose) === "shutdown") { - throw initialControlPlaneClose; + connection.initialAttachPending = false; + const initialControlPlaneClose = connection.initialControlPlaneClose; + connection.initialControlPlaneClose = undefined; + if (initialControlPlaneClose) { + // No listeners exist yet: a terminal close rejects the attach; the rest replays through the one handler. + if (getDaemonSocketCloseReason(initialControlPlaneClose) === "shutdown") { + throw initialControlPlaneClose; + } + connection.handleTransportClose(initialControlPlaneClose); } - connection.handleTransportClose(initialControlPlaneClose); + return connection; + } catch (error) { + lastError = connection.initialControlPlaneClose ?? error; + if (connection.disposed || connection.terminalCloseEmitted || !isSessionNotReadyError(error)) { + break; + } + const delayMs = Math.min(deadline - Date.now(), 2000, 250 * 2 ** Math.min(attempt++, 3)); + await delay(delayMs); } - return connection; - } catch (error) { - connection.initialAttachPending = false; - await connection.dispose(); - throw error; } + connection.initialAttachPending = false; + await connection.dispose(); + throw lastError; } async attach(options?: { recoverable?: boolean }): Promise { @@ -1868,29 +1901,35 @@ export class DaemonAgentConnection implements AgentConnection { if (this.disposed || this.terminalCloseEmitted) { return; } - if (isSessionStartingOrRecoveringError(attachError)) { - throw attachError; - } - // If attach failed because session was not found (e.g. restored under a new activeSessionId), - // query list to find restored session by file or sessionId. - const listResponse = await this.client - .request({ type: "list" }, 30000, { recoverable: false }) - .catch(() => undefined); - if (listResponse && listResponse.success) { - const sessions = readSessionSummaries(listResponse.data); - const restored = sessions.find( - (s) => - s.activeSessionId !== undefined && - ((this.attachedSessionFile !== undefined && s.sessionFile === this.attachedSessionFile) || - (this.attachedSessionId !== undefined && s.sessionId === this.attachedSessionId)), - ); - if (restored?.activeSessionId && restored.activeSessionId !== this.activeSessionId) { - this.sessionRevision++; - this.activeSessionId = restored.activeSessionId; - this.lastEventSequence = undefined; - this.lastEventCursor = undefined; - this.retiredEventGenerations.clear(); - await this.attach({ recoverable: false }); + const isUnknownSession = + attachError instanceof Error && + attachError.message.toLowerCase().includes("unknown active session"); + if (isUnknownSession) { + const listResponse = await this.client + .request({ type: "list" }, 30000, { recoverable: false }) + .catch(() => undefined); + if (listResponse && listResponse.success) { + const sessions = readSessionSummaries(listResponse.data); + const restored = sessions.find( + (s) => + (this.attachedSessionFile !== undefined && s.sessionFile === this.attachedSessionFile) || + (this.attachedSessionId !== undefined && s.sessionId === this.attachedSessionId), + ); + if (restored?.activeSessionId) { + if (restored.activeSessionId !== this.activeSessionId) { + this.sessionRevision++; + this.activeSessionId = restored.activeSessionId; + this.lastEventSequence = undefined; + this.lastEventCursor = undefined; + this.retiredEventGenerations.clear(); + } + await this.attach({ recoverable: false }); + } else { + throw new DaemonSessionStartingError( + this.activeSessionId, + "Session worker is queued for restore", + ); + } } else { throw attachError; } @@ -1902,7 +1941,7 @@ export class DaemonAgentConnection implements AgentConnection { const snapshot = await this.getInitialSnapshot({ recoverable: false }); if (this.disposed || this.terminalCloseEmitted) return; void this.emit({ type: "session_resynced", snapshot }); - void this.emit({ type: "connection_status", status: "connected" }); + this.emitConnectedStatus(); } return; } catch (error) { @@ -1910,7 +1949,7 @@ export class DaemonAgentConnection implements AgentConnection { if (this.disposed || this.terminalCloseEmitted) { return; } - const isTransientWorkerNotReady = isSessionStartingOrRecoveringError(error); + const isTransientWorkerNotReady = isSessionNotReadyError(error); // A direct-half failure must not tear down a control-plane socket with a completed handshake. const shouldResetControlPlane = !isTransientWorkerNotReady && @@ -1938,11 +1977,9 @@ export class DaemonAgentConnection implements AgentConnection { this.terminalCloseEmitted = true; const closeReason = getDaemonSocketCloseReason(lastError); const errorMsg = - closeReason === "shutdown" - ? this.formatDaemonSessionClosedError("shutdown") - : closeReason === "update" - ? this.formatDaemonSessionClosedError("update") - : `Daemon reconnection failed: ${formatErrorSentence(lastError)} The session transcript remains saved; restart Prime Agent or reopen the session from Agents View. ${this.formatDaemonDiagnosticContext()}`; + this.updateRestartPending || closeReason === "update" + ? this.formatUpdateReconnectError(lastError) + : `Daemon reconnection failed: ${formatErrorSentence(lastError)} The session transcript remains saved; restart Prime Agent or reopen the session from Agents View. ${this.formatDaemonDiagnosticContext()}`; await this.emit({ type: "closed", error: errorMsg }); } })().finally(() => { @@ -2151,6 +2188,11 @@ export class DaemonAgentConnection implements AgentConnection { void this.reconnectAfterUpdate(); return; } + if (message.reason === "shutdown" && this.options.recoverDaemon) { + this.captureDaemonLogPath(); + void this.reconnect(new Error("Daemon session closed for shutdown")); + return; + } this.terminalCloseEmitted = true; const error = this.formatDaemonSessionClosedError(message.reason); this.rejectSnapshotAssemblies(new Error(error)); @@ -2211,11 +2253,14 @@ export class DaemonAgentConnection implements AgentConnection { status: "reconnecting", error: "session is restarting - waiting for the daemon", }); - const reconnectPromise = reconnectDaemonTransportAfterUpdate(this.client) + const reconnectPromise = reconnectDaemonTransportAfterUpdate( + this.client, + this.options.reconnectTimeoutMs ?? UPDATE_RECONNECT_TIMEOUT_MS, + ) .then(() => this.restoreConnectionAfterUpdate()) .then(() => { if (!this.disposed && !this.terminalCloseEmitted) { - void this.emit({ type: "connection_status", status: "connected" }); + this.emitConnectedStatus(); } }) .catch(async (error: unknown) => { @@ -2244,51 +2289,43 @@ export class DaemonAgentConnection implements AgentConnection { if (!sessionId && !sessionFile) { throw new Error("the previous session identity is unavailable"); } - const deadline = Date.now() + UPDATE_RECONNECT_TIMEOUT_MS; + const deadline = Date.now() + (this.options.reconnectTimeoutMs ?? UPDATE_RECONNECT_TIMEOUT_MS); let lastError: unknown; let attempt = 0; while (!this.disposed && !this.terminalCloseEmitted && Date.now() < deadline) { try { await this.client.reconnect(1000); - if (this.disposed || this.terminalCloseEmitted) { - return; - } - // This loop owns the retry: a socket close must reject these instead of parking them behind a hello it can never produce. + if (this.disposed || this.terminalCloseEmitted) return; + const response = await this.client.request({ type: "list" }, 30000, { recoverable: false }); - if (this.disposed || this.terminalCloseEmitted) { - return; - } + if (this.disposed || this.terminalCloseEmitted) return; if (!response.success) { throw deserializeDaemonError(response); } const sessions = readSessionSummaries(response.data); const restored = sessions.find( (summary) => - summary.activeSessionId !== undefined && - ((sessionFile !== undefined && summary.sessionFile === sessionFile) || - (sessionId !== undefined && summary.sessionId === sessionId)), + (sessionFile !== undefined && summary.sessionFile === sessionFile) || + (sessionId !== undefined && summary.sessionId === sessionId), ); if (restored?.activeSessionId) { - if (this.disposed || this.terminalCloseEmitted) { - return; + if (this.disposed || this.terminalCloseEmitted) return; + if (restored.activeSessionId !== this.activeSessionId) { + this.sessionRevision++; + this.activeSessionId = restored.activeSessionId; + this.lastEventSequence = undefined; + this.lastEventCursor = undefined; + this.retiredEventGenerations.clear(); } - this.sessionRevision++; - this.activeSessionId = restored.activeSessionId; - this.lastEventSequence = undefined; - this.lastEventCursor = undefined; - this.retiredEventGenerations.clear(); await this.attach({ recoverable: false }); - if (this.disposed || this.terminalCloseEmitted) { - return; - } + if (this.disposed || this.terminalCloseEmitted) return; const snapshot = await this.getInitialSnapshot({ recoverable: false }); - if (this.disposed || this.terminalCloseEmitted) { - return; - } + if (this.disposed || this.terminalCloseEmitted) return; this.updateRestartPending = false; void this.emit({ type: "session_resynced", snapshot }); return; } + lastError = new DaemonSessionStartingError(this.activeSessionId, "Session worker is queued for restore"); } catch (error) { lastError = error; } @@ -2296,9 +2333,7 @@ export class DaemonAgentConnection implements AgentConnection { const delayMs = Math.min(deadline - Date.now(), 2000, 250 * 2 ** Math.min(attempt++, 3)); await delay(delayMs); } - if (this.disposed || this.terminalCloseEmitted) { - return; - } + if (this.disposed || this.terminalCloseEmitted) return; throw lastError ?? new Error("the restored session did not become available"); } diff --git a/packages/coding-agent/src/modes/agent-connection/types.ts b/packages/coding-agent/src/modes/agent-connection/types.ts index 7c5468004..717add22f 100644 --- a/packages/coding-agent/src/modes/agent-connection/types.ts +++ b/packages/coding-agent/src/modes/agent-connection/types.ts @@ -780,7 +780,12 @@ export type AgentConnectionEvent = | { type: "session_status"; recap?: string } | { type: "extension_ui_request"; request: AgentConnectionExtensionUiRequest } | { type: "extension_error"; extensionPath: string; event: string; error: string } - | { type: "connection_status"; status: "reconnecting" | "connected"; error?: string } + | { + type: "connection_status"; + status: "reconnecting" | "connected"; + error?: string; + buildIdMismatch?: { client: string; daemon: string }; + } | { type: "heartbeats_changed" } | { type: "closed"; error?: string }; diff --git a/packages/coding-agent/src/modes/daemon/daemon-errors.ts b/packages/coding-agent/src/modes/daemon/daemon-errors.ts index 83720d7b8..f0aa4386b 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-errors.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-errors.ts @@ -45,6 +45,49 @@ export function isSessionStartingOrRecoveringError(error: unknown): boolean { return false; } +export function isSessionClosingError(error: unknown): boolean { + if (!error) return false; + const message = error instanceof Error ? error.message : String(error); + return ( + message.includes("is closing") || + message.includes("is stopping") || + message.includes("still initializing") || + message.includes("did not finish initializing") + ); +} + +export function isSessionNotReadyError(error: unknown): boolean { + if (isSessionStartingOrRecoveringError(error)) return true; + if (isSessionClosingError(error)) return true; + if (error instanceof Error) { + const msg = error.message.toLowerCase(); + if ( + msg.includes("unknown active session") || + msg.includes("session worker is queued") || + msg.includes("cannot send daemon command") || + msg.includes("not connected") || + msg.includes("connection closed") || + msg.includes("econnrefused") || + msg.includes("enoent") + ) { + return true; + } + } + if (typeof error === "object" && error !== null && "code" in error) { + const code = (error as { code?: unknown }).code; + if ( + code === "daemon_not_connected" || + code === "ECONNREFUSED" || + code === "ENOENT" || + code === "ECONNRESET" || + code === "EPIPE" + ) { + return true; + } + } + return false; +} + export function serializeDaemonError(error: unknown): DaemonErrorInfo | undefined { if (error instanceof MissingSessionCwdError) { return { code: "missing_session_cwd", issue: error.issue }; diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index 3bf6d9606..8c130aed2 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -8081,7 +8081,7 @@ export class AgentDaemon { } private getShutdownClosingReason(): DaemonClosingReason { - return this.updateRestart?.phase === "publishing" ? "update" : "shutdown"; + return this.updateRestart !== undefined ? "update" : "shutdown"; } private async shutdown(exitCode: number): Promise { diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 6ee5b4c9e..ab276fdca 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -6138,6 +6138,12 @@ export class InteractiveMode { } else if (event.type === "connection_status") { if (event.status === "connected") { this.clearStatus(); + if (event.buildIdMismatch) { + this.showStatus( + `The Prime Agent daemon updated to ${event.buildIdMismatch.daemon}. Reopen this window to use the updated client.`, + "warning", + ); + } await this.refreshHeartbeatCatalog(); } else if (event.status === "reconnecting") { this.showStatus("session is restarting - waiting for the daemon", "warning"); diff --git a/packages/coding-agent/test/agent-connection-daemon.test.ts b/packages/coding-agent/test/agent-connection-daemon.test.ts index c57a150fc..2a07ff54a 100644 --- a/packages/coding-agent/test/agent-connection-daemon.test.ts +++ b/packages/coding-agent/test/agent-connection-daemon.test.ts @@ -2010,6 +2010,103 @@ describe("DaemonAgentConnection", () => { 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("does not emit a restored session after disposal begins", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.updateRestartSessions = [ -- 2.51.2