From 89018fe7902d774a9605f51be65b8af1330d834d Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Tue, 3 Nov 2026 09:21:47 +0100
Subject: [PATCH] fix(client): preserve latency deadlines when sharing
 heartbeat pings

---
 packages/client/src/daemon-client.test.ts | 24 ++++++++++++++++++++++-
 packages/client/src/daemon-client.ts      | 16 +++++++++++----
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/packages/client/src/daemon-client.test.ts b/packages/client/src/daemon-client.test.ts
index 4f48553907..a16d2e4be8 100644
--- a/packages/client/src/daemon-client.test.ts
+++ b/packages/client/src/daemon-client.test.ts
@@ -2670,8 +2670,18 @@ test("chunks extend the heartbeat deadline until silence causes two misses", asy
   daemon.openConnection();
   await connection;
   client.ensureConnected({ verify: true });
+  let measurementError: unknown;
+  void client.measureLatency({ timeoutMs: 5_000 }).catch((error: unknown) => {
+    measurementError = error;
+  });
+  await vi.advanceTimersByTimeAsync(5_000);
+  expect(measurementError).toEqual(
+    expect.objectContaining({
+      message: "Latency measurement timed out (5000ms)",
+    }),
+  );
   for (let index = 0; index < 4; index++) {
-    await vi.advanceTimersByTimeAsync(10_000);
+    await vi.advanceTimersByTimeAsync(index === 0 ? 5_000 : 10_000);
     daemon.receiveChunk(index === 0);
     expect(daemon.closesFromClient()).toEqual([]);
   }
@@ -2831,6 +2841,18 @@ test("sends only one ping while a slow pong is still outstanding", async () => {
   expect(session.pingTimestamps()).toEqual(["10s"]);
 });
 
+test("a heartbeat sharing a latency measurement still saves the RTT", async () => {
+  useHeartbeatClock();
+  const session = new DaemonClientSession();
+  session.daemonAnswersPingsAfter("2s");
+  await session.connect();
+  await session.advance(9_000);
+  const measurement = session.measureLatency({ timeoutMs: 5_000 });
+  await session.advance(2_000);
+  await expect(measurement).resolves.toBe(2000);
+  expect(session.lastLivenessRttMs()).toBe(2000);
+});
+
 test("reports the round-trip time of the last successful heartbeat", async () => {
   useHeartbeatClock();
   const session = new DaemonClientSession();
diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts
index 65c10538b8..4bad95d31f 100644
--- a/packages/client/src/daemon-client.ts
+++ b/packages/client/src/daemon-client.ts
@@ -2536,9 +2536,18 @@ export class DaemonClient {
 
   measureLatency(params?: { timeoutMs?: number }): Promise<number> {
     const timeoutMs = Math.max(1, params?.timeoutMs ?? DEFAULT_LIVENESS_TIMEOUT_MS);
-    return this.sendPingAwaitRtt({ timeoutMs, drivesLivenessFailure: false }).catch((error) => {
-      throw toTimeoutError(error, "Latency measurement", timeoutMs);
+    let deadline: ReturnType<typeof setTimeout>;
+    const timeout = new Promise<never>((_resolve, reject) => {
+      deadline = setTimeout(() => reject(new PingTimeoutError(timeoutMs)), timeoutMs);
     });
+    return Promise.race([
+      this.sendPingAwaitRtt({ timeoutMs, drivesLivenessFailure: false }),
+      timeout,
+    ])
+      .catch((error) => {
+        throw toTimeoutError(error, "Latency measurement", timeoutMs);
+      })
+      .finally(() => clearTimeout(deadline));
   }
 
   private async livenessPing(params?: { timeoutMs?: number }): Promise<number> {
@@ -2639,8 +2648,7 @@ export class DaemonClient {
   private tickLivenessHeartbeat(): void {
     this.stopLivenessHeartbeat();
     const transport = this.transport;
-    const tick =
-      this.pingProbe?.promise ?? this.livenessPing({ timeoutMs: LIVENESS_HEARTBEAT_TIMEOUT_MS });
+    const tick = this.livenessPing({ timeoutMs: LIVENESS_HEARTBEAT_TIMEOUT_MS });
     void tick
       .catch(() => {})
       .finally(() => {