From 6dd6d800e0ff7ba95233943d355839690b700035 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 3 Nov 2026 09:13:16 +0100 Subject: [PATCH] fix(client): keep active transfers alive during foreground checks Foreground used an independent three-second probe that closed live connections while a pong waited behind bulk data. Share the heartbeat clock and refresh its outstanding deadline on validated inbound frames, including incomplete chunk assemblies. --- packages/client/src/daemon-client.test.ts | 88 +++++++++++++++++++++- packages/client/src/daemon-client.ts | 71 ++++++++--------- packages/protocol/src/chunked-transport.ts | 2 + 3 files changed, 120 insertions(+), 41 deletions(-) diff --git a/packages/client/src/daemon-client.test.ts b/packages/client/src/daemon-client.test.ts index b2d7f46651..645bc0c3f1 100644 --- a/packages/client/src/daemon-client.test.ts +++ b/packages/client/src/daemon-client.test.ts @@ -657,6 +657,23 @@ class FakeDaemon { }, }; + finishTransfer(): void { + // Four MiB of JSON whitespace precedes this valid application message. + const bytes = new TextEncoder().encode('{"type":"pong"}'); + const frame = new Uint8Array(bytes.length + 4); + frame.set([0x31, 1, 0, 2]); + frame.set(bytes, 4); + this.onMessage(frame); + this.daemonAnswersPingsAfter("fast"); + } + + receiveChunk(start: boolean): void { + const frame = new Uint8Array(1024 * 1024 + 4); + frame.set([0x31, 1, 0, start ? 1 : 0]); + frame.fill(32, 4); + this.onMessage(frame); + } + openConnection(): void { this.onOpen(); this.onMessage( @@ -667,6 +684,7 @@ class FakeDaemon { payload: { status: "server_info", serverId: "srv_heartbeat_test", + features: { binaryChunkingV1: true }, hostname: null, version: null, }, @@ -2571,7 +2589,7 @@ test("keeps default connect timeout shorter than session RPC waiters", async () } }); -test("foreground verification replaces a silently broken socket without waiting for heartbeats", async () => { +test("foreground checks a silently broken socket through two normal heartbeat misses", async () => { useHeartbeatClock(); const first = new FakeDaemon(); const second = new FakeDaemon(); @@ -2589,7 +2607,11 @@ test("foreground verification replaces a silently broken socket without waiting await connection; client.ensureConnected({ verify: true }); - await vi.advanceTimersByTimeAsync(3_000); + await vi.advanceTimersByTimeAsync(15_000); + expect(first.closesFromClient()).toEqual([]); + await vi.advanceTimersByTimeAsync(25_000); + expect(first.closesFromClient()).toEqual([{ code: 1001, reason: "Liveness check timed out" }]); + client.ensureConnected(); expect(attempts).toBe(2); second.openConnection(); expect(client.getConnectionState()).toEqual({ status: "connected" }); @@ -2617,7 +2639,67 @@ test("foreground verification preserves a healthy socket and deduplicates simult expect(daemon.closesFromClient()).toEqual([]); }); -test("an obsolete foreground probe cannot close a replacement connection", async () => { +test("chunks extend the heartbeat deadline until silence causes two misses", async () => { + useHeartbeatClock(); + const daemon = new FakeDaemon(); + daemon.daemonGoesSilent(); + const client = new DaemonClient({ + url: "ws://test", + clientId: "bulk-resume", + logger: noopLogger, + reconnect: { enabled: false }, + transportFactory: () => daemon.transport, + }); + clients.push(client); + const connection = client.connect(); + daemon.openConnection(); + await connection; + client.ensureConnected({ verify: true }); + for (let index = 0; index < 4; index++) { + await vi.advanceTimersByTimeAsync(10_000); + daemon.receiveChunk(index === 0); + expect(daemon.closesFromClient()).toEqual([]); + } + // No complete message or pong has arrived, but chunks keep the ping alive. + expect(daemon.pingTimestamps()).toEqual(["0s"]); + await vi.advanceTimersByTimeAsync(15_000); + expect(daemon.closesFromClient()).toEqual([]); + await vi.advanceTimersByTimeAsync(25_000); + expect(daemon.closesFromClient()).toEqual([{ code: 1001, reason: "Liveness check timed out" }]); +}); + +test("foreground preserves a multi-MB transfer with its pong queued behind chunks", async () => { + useHeartbeatClock(); + const daemon = new FakeDaemon(); + daemon.daemonGoesSilent(); + const client = new DaemonClient({ + url: "ws://test", + clientId: "bulk-resume", + logger: noopLogger, + reconnect: { enabled: false }, + transportFactory: () => daemon.transport, + }); + clients.push(client); + const connection = client.connect(); + daemon.openConnection(); + await connection; + client.ensureConnected({ verify: true }); + await vi.advanceTimersByTimeAsync(3_001); + expect(daemon.closesFromClient()).toEqual([]); + for (let index = 0; index < 4; index++) { + await vi.advanceTimersByTimeAsync(10_000); + daemon.receiveChunk(index === 0); + expect(daemon.closesFromClient()).toEqual([]); + } + // No complete message or pong has arrived, but chunks keep the ping alive. + expect(daemon.pingTimestamps()).toEqual(["0s"]); + daemon.finishTransfer(); + await vi.advanceTimersByTimeAsync(30_000); + expect(client.getConnectionState()).toEqual({ status: "connected" }); + expect(daemon.closesFromClient()).toEqual([]); +}); + +test("an obsolete heartbeat cannot close a replacement connection", async () => { useHeartbeatClock(); const first = new FakeDaemon(); const second = new FakeDaemon(); diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index bfc97f5674..65c10538b8 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -1295,6 +1295,7 @@ interface PingProbe { reject: (error: Error) => void; timeoutHandle: ReturnType; startedAt: number; + extendDeadline: () => void; // Whether a timeout on this ping should be recorded as a liveness failure. Only the // heartbeat sets this; a latency measurement never drives teardown, even when a // heartbeat tick shares (dedupes onto) an in-flight measurement ping. @@ -1438,7 +1439,6 @@ export class DaemonClient { private runtimeMetricsInterval: ReturnType | null = null; private runtimeMetrics: DaemonClientRuntimeMetrics | null = null; private pingProbe: PingProbe | null = null; - private connectionVerification: DaemonTransport | null = null; private livenessHeartbeatTimer: ReturnType | null = null; private lastLivenessRttMs: number | null = null; private consecutiveLivenessFailures = 0; @@ -1588,6 +1588,7 @@ export class DaemonClient { send: (frame) => transport.send(frame), onFailure: () => transport.close(4003, "Transport send failed"), deliver: (frame) => this.handleTransportMessage(frame), + onActivity: () => this.recordTransportActivity(), }); this.lastServerInfoMessage = null; this.authFailureReasonValue = null; @@ -1784,7 +1785,7 @@ export class DaemonClient { this.shouldReconnect = true; } if (this.connectionState.status === "connected") { - if (options?.verify) this.verifyConnection(); + if (options?.verify) this.tickLivenessHeartbeat(); return; } if (this.connectionState.status === "connecting") return; @@ -1797,26 +1798,9 @@ export class DaemonClient { ); } - private verifyConnection(): void { - const transport = this.transport; - if (!transport || this.connectionVerification === transport) return; - this.connectionVerification = transport; - // A session probe has its own deadline, independent of a heartbeat that the OS - // may have suspended. A successful response also proves the session can serve RPCs. - void this.ping({ timeoutMs: 3_000 }) - .catch((error: unknown) => { - if (this.transport !== transport || this.connectionState.status !== "connected") return; - this.disposeTransport(1001, "Connection verification failed"); - this.scheduleReconnect({ - reason: error instanceof Error ? error.message : String(error), - event: "CONNECTION_VERIFICATION_FAILED", - reasonCode: "liveness_timeout", - }); - this.ensureConnected(); - }) - .finally(() => { - if (this.connectionVerification === transport) this.connectionVerification = null; - }); + private recordTransportActivity(): void { + this.consecutiveLivenessFailures = 0; + if (this.pingProbe?.drivesLivenessFailure) this.pingProbe.extendDeadline(); } requireConnectedHosts(): void { @@ -2590,21 +2574,24 @@ export class DaemonClient { resolveProbe = resolve; rejectProbe = reject; }); + const expire = () => { + if (this.pingProbe !== probe) return; + this.pingProbe = null; + const error = new PingTimeoutError(timeoutMs); + probe.reject(error); + if (probe.drivesLivenessFailure) { + this.recordLivenessFailure(toTimeoutError(error, "Liveness check", timeoutMs)); + } + }; const probe: PingProbe = { promise, resolve: (value) => resolveProbe?.(value), reject: (error) => rejectProbe?.(error), - timeoutHandle: setTimeout(() => { - if (this.pingProbe !== probe) { - return; - } - this.pingProbe = null; - const error = new PingTimeoutError(timeoutMs); - probe.reject(error); - if (probe.drivesLivenessFailure) { - this.recordLivenessFailure(toTimeoutError(error, "Liveness check", timeoutMs)); - } - }, timeoutMs), + timeoutHandle: setTimeout(expire, timeoutMs), + extendDeadline: () => { + clearTimeout(probe.timeoutHandle); + probe.timeoutHandle = setTimeout(expire, timeoutMs); + }, startedAt, drivesLivenessFailure: params.drivesLivenessFailure, }; @@ -2645,14 +2632,22 @@ export class DaemonClient { } this.livenessHeartbeatTimer = setTimeout(() => { this.livenessHeartbeatTimer = null; - this.livenessPing({ timeoutMs: LIVENESS_HEARTBEAT_TIMEOUT_MS }) - .catch(() => {}) - .finally(() => { - this.scheduleNextLivenessHeartbeat(); - }); + this.tickLivenessHeartbeat(); }, LIVENESS_HEARTBEAT_INTERVAL_MS); } + private tickLivenessHeartbeat(): void { + this.stopLivenessHeartbeat(); + const transport = this.transport; + const tick = + this.pingProbe?.promise ?? this.livenessPing({ timeoutMs: LIVENESS_HEARTBEAT_TIMEOUT_MS }); + void tick + .catch(() => {}) + .finally(() => { + if (this.transport === transport) this.scheduleNextLivenessHeartbeat(); + }); + } + // ============================================================================ // Agent RPCs (requestId-correlated) // ============================================================================ diff --git a/packages/protocol/src/chunked-transport.ts b/packages/protocol/src/chunked-transport.ts index a32d68424f..8faa22cd96 100644 --- a/packages/protocol/src/chunked-transport.ts +++ b/packages/protocol/src/chunked-transport.ts @@ -25,6 +25,7 @@ export class ChunkedTransport { supportsDrain?: boolean; send: (frame: Frame, waitForDrain: boolean) => void | Promise; deliver: (frame: Frame) => void; + onActivity?: () => void; onFailure?: (error: unknown) => void; }, ) { @@ -134,6 +135,7 @@ export class ChunkedTransport { { cause: error }, ); } + this.io.onActivity?.(); if (message !== undefined) this.io.deliver(message); }