Keep a healthy host connection from being replaced on stale latency data (#9116)

On a host with two saved connections, the app replaced a working connection with the other one and showed "Reconnecting to host" for about two seconds each time. The adaptive switch rule needs three consecutive wins of at least 40 ms. It counted the cached alternate sample again on every active-only tick, so one measurement became three wins in 20 seconds. After a liveness-timeout reconnect it also compared the previous socket's heartbeat RTT with the alternate. A win now needs a candidate measured in the current probe cycle and a heartbeat RTT from the socket that is connected now, and the streak resets when the active client leaves connected or a fresh probe of the candidate fails. Closing a client with subscription releases in flight no longer publishes a false transport close or aborts close() midway. The host runtime does this on every legitimate switch. Follow-ups: #9147, #9149.
2026-11-09 09:54:40 +00:00 · 2026-11-09 09:17:44 +07:00
parent def249db99
commit 8cc8e12007
4 changed files with 350 additions and 23 deletions
@@ -430,6 +430,20 @@ function agentPermission(id: string): AgentPermissionRequest {
  return { id, provider: "codex", name: id, kind: "tool", title: id };
 }

+const relayPrimary: HostConnection = {
+  id: "relay:relay.paseo.sh:443",
+  type: "relay",
+  relayEndpoint: "relay.paseo.sh:443",
+  daemonPublicKeyB64: "pk_test",
+};
+
+const relayAlternate: HostConnection = {
+  id: "relay:relay-next.paseo.sh:443",
+  type: "relay",
+  relayEndpoint: "relay-next.paseo.sh:443",
+  daemonPublicKeyB64: "pk_test",
+};
+
 function makeHost(input?: Partial<HostProfile>): HostProfile {
  const direct: HostConnection = {
    id: "direct:lan:6767",
@@ -1566,6 +1580,233 @@ describe("HostRuntimeController", () => {
    expect(switched).toBe(true);
  });

+  it("keeps the live connection while one alternate sample is re-read at active ticks", async () => {
+    // Replays the recorded Android sequence: both relays probed at startup, the alternate
+    // measured once, then only the active connection is re-read every ten seconds.
+    useHostRuntimeClock();
+    const host = makeHost({
+      connections: [relayPrimary, relayAlternate],
+      preferredConnectionId: relayPrimary.id,
+    });
+    const clients: FakeDaemonClient[] = [];
+    const latencies: Record<string, number | Error> = {
+      [relayPrimary.id]: 1053,
+      [relayAlternate.id]: 251,
+    };
+    const controller = new HostRuntimeController({
+      host,
+      deps: makeDeps(latencies, clients),
+    });
+
+    await controller.start();
+    expect(controller.getSnapshot().activeConnectionId).toBe(relayPrimary.id);
+    const activeClient = controller.getSnapshot().client as unknown as FakeDaemonClient;
+    const clientsAfterStart = clients.length;
+
+    await vi.advanceTimersByTimeAsync(10_500);
+    activeClient.heartbeatReportsRtt(376);
+    await vi.advanceTimersByTimeAsync(30_000);
+
+    expect(clients.length).toBe(clientsAfterStart);
+    expect(controller.getSnapshot().activeConnectionId).toBe(relayPrimary.id);
+    expect(controller.getSnapshot().client).toBe(activeClient);
+    expect(activeClient.isDisposed()).toBe(false);
+    await controller.stop();
+  });
+
+  it("does not hold the previous socket's RTT against a reconnected active connection", async () => {
+    // Replays the recorded foreground recovery: the active client reconnects on its own while an
+    // alternate probe is in flight. The recovered socket has no heartbeat RTT yet.
+    useHostRuntimeClock();
+    const host = makeHost({
+      connections: [relayPrimary, relayAlternate],
+      preferredConnectionId: relayPrimary.id,
+    });
+    const clients: FakeDaemonClient[] = [];
+    const latencies: Record<string, number | Error> = {
+      [relayPrimary.id]: 300,
+      [relayAlternate.id]: 392,
+    };
+    let connectGate: Promise<void> | null = null;
+    const baseDeps = makeDeps(latencies, clients);
+    const controller = new HostRuntimeController({
+      host,
+      deps: {
+        ...baseDeps,
+        connectToDaemon: async (input) => {
+          if (connectGate) await connectGate;
+          return baseDeps.connectToDaemon(input);
+        },
+      },
+    });
+
+    await controller.start();
+    expect(controller.getSnapshot().activeConnectionId).toBe(relayPrimary.id);
+    const activeClient = controller.getSnapshot().client as unknown as FakeDaemonClient;
+    activeClient.heartbeatReportsRtt(403);
+    await vi.advanceTimersByTimeAsync(150_500);
+    expect(controller.getSnapshot().activeConnectionId).toBe(relayPrimary.id);
+
+    activeClient.setConnectionState({ status: "disconnected", reason: "liveness_timeout" });
+    activeClient.setConnectionState({ status: "connecting", attempt: 2 });
+    activeClient.heartbeatReportsRtt(null);
+    latencies[relayAlternate.id] = 280;
+    const alternateConnect = createDeferred<void>();
+    connectGate = alternateConnect.promise;
+    await vi.advanceTimersByTimeAsync(2_800);
+    activeClient.setConnectionState({ status: "connected" });
+    await vi.advanceTimersByTimeAsync(500);
+    alternateConnect.resolve();
+    connectGate = null;
+    await vi.advanceTimersByTimeAsync(100);
+    expect(controller.getSnapshot().probeByConnectionId.get(relayAlternate.id)).toEqual({
+      status: "available",
+      latencyMs: 280,
+    });
+    const clientsAfterProbe = clients.length;
+
+    await vi.advanceTimersByTimeAsync(30_000);
+
+    expect(clients.length).toBe(clientsAfterProbe);
+    expect(controller.getSnapshot().activeConnectionId).toBe(relayPrimary.id);
+    expect(controller.getSnapshot().client).toBe(activeClient);
+    expect(activeClient.isDisposed()).toBe(false);
+    await controller.stop();
+  });
+
+  it("starts the switch streak over when the active socket reconnects", async () => {
+    useHostRuntimeClock();
+    const host = makeHost({ preferredConnectionId: "direct:lan:6767" });
+    const clients: FakeDaemonClient[] = [];
+    const latencies: Record<string, number | Error> = {
+      "direct:lan:6767": 15,
+      "relay:relay.paseo.sh:443": 60,
+    };
+    const controller = new HostRuntimeController({
+      host,
+      deps: makeDeps(latencies, clients),
+    });
+
+    await controller.start({ autoProbe: false });
+    expect(controller.getSnapshot().activeConnectionId).toBe("direct:lan:6767");
+    const activeClient = controller.getSnapshot().client as unknown as FakeDaemonClient;
+
+    latencies["relay:relay.paseo.sh:443"] = 30;
+    activeClient.heartbeatReportsRtt(95);
+    for (let index = 0; index < 2; index += 1) {
+      await vi.advanceTimersByTimeAsync(120_000);
+      await controller.runProbeCycleNow();
+      expect(controller.getSnapshot().activeConnectionId).toBe("direct:lan:6767");
+    }
+
+    activeClient.setConnectionState({ status: "disconnected", reason: "network changed" });
+    activeClient.setConnectionState({ status: "connected" });
+    activeClient.heartbeatReportsRtt(95);
+    await vi.advanceTimersByTimeAsync(120_000);
+    await controller.runProbeCycleNow();
+    expect(controller.getSnapshot().activeConnectionId).toBe("direct:lan:6767");
+
+    for (let index = 0; index < 2; index += 1) {
+      await vi.advanceTimersByTimeAsync(120_000);
+      await controller.runProbeCycleNow();
+    }
+    expect(controller.getSnapshot().activeConnectionId).toBe("relay:relay.paseo.sh:443");
+    await controller.stop();
+  });
+
+  it("ends the switch streak when a fresh probe of the alternate fails", async () => {
+    useHostRuntimeClock();
+    const host = makeHost({ preferredConnectionId: "direct:lan:6767" });
+    const clients: FakeDaemonClient[] = [];
+    const latencies: Record<string, number | Error> = {
+      "direct:lan:6767": 15,
+      "relay:relay.paseo.sh:443": 60,
+    };
+    const controller = new HostRuntimeController({
+      host,
+      deps: makeDeps(latencies, clients),
+    });
+
+    await controller.start({ autoProbe: false });
+    expect(controller.getSnapshot().activeConnectionId).toBe("direct:lan:6767");
+    const activeClient = controller.getSnapshot().client as unknown as FakeDaemonClient;
+    activeClient.heartbeatReportsRtt(95);
+
+    latencies["relay:relay.paseo.sh:443"] = 30;
+    for (let index = 0; index < 2; index += 1) {
+      await vi.advanceTimersByTimeAsync(120_000);
+      await controller.runProbeCycleNow();
+      expect(controller.getSnapshot().activeConnectionId).toBe("direct:lan:6767");
+    }
+
+    latencies["relay:relay.paseo.sh:443"] = new Error("relay unreachable");
+    await vi.advanceTimersByTimeAsync(120_000);
+    await controller.runProbeCycleNow();
+    expect(controller.getSnapshot().probeByConnectionId.get("relay:relay.paseo.sh:443")).toEqual({
+      status: "unavailable",
+      latencyMs: null,
+    });
+
+    latencies["relay:relay.paseo.sh:443"] = 30;
+    await vi.advanceTimersByTimeAsync(120_000);
+    await controller.runProbeCycleNow();
+    expect(controller.getSnapshot().activeConnectionId).toBe("direct:lan:6767");
+
+    for (let index = 0; index < 2; index += 1) {
+      await vi.advanceTimersByTimeAsync(120_000);
+      await controller.runProbeCycleNow();
+    }
+    expect(controller.getSnapshot().activeConnectionId).toBe("relay:relay.paseo.sh:443");
+    await controller.stop();
+  });
+
+  it("restarts the streak when a different candidate becomes the fastest fresh sample", async () => {
+    useHostRuntimeClock();
+    const direct: HostConnection = {
+      id: "direct:lan:6767",
+      type: "directTcp",
+      endpoint: "lan:6767",
+    };
+    const host = makeHost({
+      connections: [direct, relayPrimary, relayAlternate],
+      preferredConnectionId: direct.id,
+    });
+    const clients: FakeDaemonClient[] = [];
+    const latencies: Record<string, number | Error> = {
+      [direct.id]: 15,
+      [relayPrimary.id]: 60,
+      [relayAlternate.id]: 70,
+    };
+    const controller = new HostRuntimeController({
+      host,
+      deps: makeDeps(latencies, clients),
+    });
+
+    await controller.start({ autoProbe: false });
+    expect(controller.getSnapshot().activeConnectionId).toBe(direct.id);
+    const activeClient = controller.getSnapshot().client as unknown as FakeDaemonClient;
+    activeClient.heartbeatReportsRtt(95);
+
+    latencies[relayPrimary.id] = 30;
+    for (let index = 0; index < 2; index += 1) {
+      await vi.advanceTimersByTimeAsync(120_000);
+      await controller.runProbeCycleNow();
+      expect(controller.getSnapshot().activeConnectionId).toBe(direct.id);
+    }
+
+    latencies[relayAlternate.id] = 20;
+    await vi.advanceTimersByTimeAsync(120_000);
+    await controller.runProbeCycleNow();
+    expect(controller.getSnapshot().activeConnectionId).toBe(direct.id);
+
+    for (let index = 0; index < 2; index += 1) {
+      await vi.advanceTimersByTimeAsync(120_000);
+      await controller.runProbeCycleNow();
+    }
+    expect(controller.getSnapshot().activeConnectionId).toBe(relayAlternate.id);
+    await controller.stop();
+  });
+
  it("exposes one snapshot with active connection and status from same source", async () => {
    const host = makeHost();
    const clients: FakeDaemonClient[] = [];
@@ -1063,37 +1063,46 @@ export class HostRuntimeController {
            expectedProbeVersion: requestVersion,
          });
        }
-        this.switchCandidateConnectionId = null;
-        this.switchCandidateHitCount = 0;
+        this.resetSwitchCandidate();
        return;
      }

-      if (!activeProbe || activeProbe.status !== "available") {
+      // Adaptive switching counts fresh comparisons only. The probe map keeps the last known
+      // sample per connection for display and failover; a streak hit needs a candidate measured
+      // in this cycle and a heartbeat RTT from the socket that is active right now. Re-reading
+      // the cached map at every ten-second active tick turned one idle-probe sample into three
+      // "consecutive" wins, and a reconnected socket inherited its predecessor's RTT, so a
+      // working relay connection was replaced within seconds of coming online.
+      const sampledCandidates = connectionsToProbe.filter(
+        (connection) => connection.id !== currentActiveConnectionId,
+      );
+      if (sampledCandidates.length === 0) {
+        // An active-only tick carries no evidence about any candidate.
+        return;
+      }
+      const fastest =
+        sampledCandidates
+          .flatMap((connection) => {
+            const probe = probeByConnectionId.get(connection.id);
+            return probe?.status === "available"
+              ? [{ connectionId: connection.id, latencyMs: probe.latencyMs }]
+              : [];
+          })
+          .sort((left, right) => left.latencyMs - right.latencyMs)[0] ?? null;
+      if (!fastest) {
+        // Every candidate sampled this cycle failed: that breaks the streak.
+        this.resetSwitchCandidate();
        return;
      }

-      const available = Array.from(probeByConnectionId.entries())
-        .filter(
-          (entry): entry is [string, Extract<ConnectionProbeState, { status: "available" }>] =>
-            entry[1].status === "available",
-        )
-        .map(([connectionId, probe]) => ({
-          connectionId,
-          latencyMs: probe.latencyMs,
-        }))
-        .sort((left, right) => left.latencyMs - right.latencyMs);
-
-      const fastest = available[0] ?? null;
-      if (!fastest || fastest.connectionId === currentActiveConnectionId) {
-        this.switchCandidateConnectionId = null;
-        this.switchCandidateHitCount = 0;
+      const activeRttMs = this.currentActiveSocketRttMs();
+      if (activeRttMs === null) {
        return;
      }

-      const improvement = activeProbe.latencyMs - fastest.latencyMs;
+      const improvement = activeRttMs - fastest.latencyMs;
      if (improvement < ADAPTIVE_SWITCH_THRESHOLD_MS) {
-        this.switchCandidateConnectionId = null;
-        this.switchCandidateHitCount = 0;
+        this.resetSwitchCandidate();
        return;
      }

@@ -1105,8 +1114,7 @@ export class HostRuntimeController {
      }

      if (this.switchCandidateHitCount >= ADAPTIVE_SWITCH_CONSECUTIVE_PROBES) {
-        this.switchCandidateConnectionId = null;
-        this.switchCandidateHitCount = 0;
+        this.resetSwitchCandidate();
        await this.switchToConnection({
          connectionId: fastest.connectionId,
          expectedProbeVersion: requestVersion,
@@ -1217,6 +1225,19 @@ export class HostRuntimeController {
    });
  }

+  // The client clears its heartbeat RTT when a socket connects, so a connected client's RTT
+  // describes the socket in use right now; a reconnecting client still carries the old one.
+  private currentActiveSocketRttMs(): number | null {
+    const client = this.snapshot.client;
+    if (!client || client.getConnectionState().status !== "connected") return null;
+    return client.getLastLivenessRttMs();
+  }
+
+  private resetSwitchCandidate(): void {
+    this.switchCandidateConnectionId = null;
+    this.switchCandidateHitCount = 0;
+  }
+
  // The active connection is probed through its own client, which may still be
  // reconnecting when the probe starts. A verdict taken then is stale if the client
  // has since connected, so the decision uses the client's current state.
@@ -1506,6 +1527,8 @@ export class HostRuntimeController {
          return;
        }
        if (state.status === "connected") this.learnRelaySecurity(client, connection.id);
+        // Switch evidence belongs to the socket it was measured against.
+        if (state.status !== "connected") this.resetSwitchCandidate();
        this.applyConnectionEvent({ type: "client_state", state, lastError: client.lastError });
        this.updateSnapshot({
          ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
@@ -5225,6 +5225,65 @@ test("logs configured runtime generation in connection transition events", async
  }
 });

+test("closing a client while subscription releases are in flight reports only the disposal", async () => {
+  const mock = createMockTransport();
+  let transportsCreated = 0;
+  const client = new DaemonClient({
+    url: "ws://test",
+    clientId: "clsk_unit_test",
+    logger: createMockLogger(),
+    reconnect: { enabled: true, baseDelayMs: 1, maxDelayMs: 1 },
+    transportFactory: () => {
+      transportsCreated += 1;
+      return mock.transport;
+    },
+  });
+  clients.push(client);
+
+  const connectPromise = client.connect();
+  mock.triggerOpen();
+  await connectPromise;
+
+  const subscriptions = ["/tmp/one", "/tmp/two"].map((cwd) =>
+    client.observeCheckoutDiff(cwd, { mode: "uncommitted" }),
+  );
+  await vi.waitFor(() => expect(mock.sent).toHaveLength(2));
+  for (const [index, frame] of [...mock.sent].entries()) {
+    const request = parseSentFrame(frame);
+    mock.triggerMessage(
+      wrapSessionMessage({
+        type: "subscribe_checkout_diff_response",
+        payload: {
+          subscriptionId: `checkout-sub-${index}`,
+          cwd: request.cwd,
+          files: [],
+          error: null,
+          requestId: request.requestId,
+        },
+      }),
+    );
+  }
+  await Promise.all(subscriptions.map((subscription) => subscription.ready));
+
+  // The host runtime releases its handlers and closes the previous client without waiting.
+  const releases = subscriptions.map((subscription) =>
+    subscription.release().catch(() => undefined),
+  );
+  await vi.waitFor(() => expect(mock.sent).toHaveLength(4));
+
+  const statuses: string[] = [];
+  client.subscribeConnectionStatus((state) => statuses.push(state.status));
+  statuses.length = 0;
+
+  await client.close();
+  await Promise.all(releases);
+  await new Promise((resolve) => setTimeout(resolve, 20));
+
+  expect(statuses).toEqual(["disposed"]);
+  expect(client.lastError).toBeNull();
+  expect(transportsCreated).toBe(1);
+});
+
 test("subscribes to checkout diff updates via RPC handshake", async () => {
  const logger = createMockLogger();
  const mock = createMockTransport();
@@ -1400,12 +1400,16 @@ export class DaemonClient {
      ),
    release: async (subscriptionId) => {
      if (!this.isConnected) return;
+      const transport = this.transport;
      try {
        await this.sendCorrelatedSessionRequest({
          message: { type: "subscription.release.request", subscriptionId },
          responseType: "subscription.release.response",
        });
      } catch (error) {
+        // close() rejects in-flight requests after dropping the transport. That failure says
+        // nothing about a live connection, and rethrowing it would abort close() midway.
+        if (this.transport !== transport) return;
        this.disposeTransport(1001, "Subscription release failed");
        this.scheduleReconnect({ reason: "Subscription release failed" });
        throw error;