Keep connections alive through brief daemon slowdowns (#1581)

* Keep connections alive through brief daemon slowdowns

The active connection's liveness check was a side effect of the
candidate-probe loop: a single 5s latency measurement doubled as the
death timeout, so a daemon that was alive but briefly stalled (e.g. a
busy event loop) was misread as dead and the connection was torn down.

Give DaemonClient its own self-scheduling heartbeat that solely owns
liveness (15s timeout, 2 consecutive misses), and make the probe
read-only on the active connection — it reads the heartbeat's last RTT
instead of pinging. Split the shared ping into a pure measurement path
(measureLatency) and a private liveness path, so nothing outside the
client can drive teardown.

* Don't let a candidate latency timeout count as a liveness failure

The heartbeat (livenessPing) and candidate ranking (measureLatency) share
one in-flight ping slot. If a heartbeat tick lands while a candidate
measurement is still in flight, the heartbeat dedupes onto it — and a
measurement timeout was then recorded as a liveness failure, nudging the
connection one step closer to a spurious teardown.

Tag each ping with whether its timeout may drive a liveness failure
(only the heartbeat sets it) and gate recordLivenessFailure on that flag,
so an adopted measurement probe can never contribute to teardown.
This commit is contained in:
Mohamed Boudra
2026-06-18 21:10:08 +08:00
committed by GitHub
parent eb628765d7
commit ea2f768d3f
4 changed files with 667 additions and 343 deletions
+413 -55
View File
@@ -1,6 +1,6 @@
import { afterEach, expect, expectTypeOf, test, vi } from "vitest";
import { z } from "zod";
import { DaemonClient, type DaemonTransport } from "./daemon-client";
import { DaemonClient, type DaemonTransport, type Logger } from "./daemon-client";
import {
decodeFileTransferFrame,
encodeFileTransferFrame,
@@ -42,7 +42,16 @@ function createMockTransport() {
let serverInfoOrdinal = 1;
const transport: DaemonTransport = {
send: (data) => sent.push(data),
send: (data) => {
sent.push(data);
if (typeof data !== "string") {
return;
}
const frame = JSON.parse(data) as { type?: string };
if (frame.type === "ping") {
onMessage(JSON.stringify({ type: "pong" }));
}
},
close: () => {},
onMessage: (handler) => {
onMessage = handler;
@@ -126,8 +135,205 @@ const clients: DaemonClient[] = [];
afterEach(async () => {
await Promise.all(clients.map((client) => client.close()));
clients.length = 0;
vi.useRealTimers();
});
const noopLogger: Logger = {
debug: () => {},
info: () => {},
warn: () => {},
error: () => {},
};
type PongMode = { kind: "answer"; delayMs: number } | { kind: "silent" };
class FakeDaemon {
private onMessage: (data: unknown) => void = () => {};
private onOpen: () => void = () => {};
private onClose: (event?: unknown) => void = () => {};
private onError: (event?: unknown) => void = () => {};
private pongMode: PongMode = { kind: "answer", delayMs: 0 };
private withheldPongs = 0;
private pingsSentAt: number[] = [];
private closeEvents: Array<{ code?: number; reason?: string }> = [];
readonly transport: DaemonTransport = {
send: (data) => {
if (typeof data !== "string") {
return;
}
const frame = JSON.parse(data) as { type?: string };
if (frame.type !== "ping") {
return;
}
this.pingsSentAt.push(performance.now());
if (this.withheldPongs > 0) {
this.withheldPongs -= 1;
return;
}
if (this.pongMode.kind === "silent") {
return;
}
if (this.pongMode.delayMs === 0) {
this.onMessage(JSON.stringify({ type: "pong" }));
return;
}
setTimeout(() => {
this.onMessage(JSON.stringify({ type: "pong" }));
}, this.pongMode.delayMs);
},
close: (code?: number, reason?: string) => {
this.closeEvents.push({ code, reason });
},
onMessage: (handler) => {
this.onMessage = handler;
return () => {};
},
onOpen: (handler) => {
this.onOpen = handler;
return () => {};
},
onClose: (handler) => {
this.onClose = handler;
return () => {};
},
onError: (handler) => {
this.onError = handler;
return () => {};
},
};
openConnection(): void {
this.onOpen();
this.onMessage(
JSON.stringify({
type: "session",
message: {
type: "status",
payload: {
status: "server_info",
serverId: "srv_heartbeat_test",
hostname: null,
version: null,
},
},
}),
);
}
daemonAnswersPingsAfter(delay: "fast" | `${number}s`): void {
this.pongMode = {
kind: "answer",
delayMs: delay === "fast" ? 0 : Number.parseFloat(delay) * 1000,
};
}
daemonGoesSilent(): void {
this.pongMode = { kind: "silent" };
}
daemonWithholdsNextPongThenAnswersFast(): void {
this.withheldPongs += 1;
this.daemonAnswersPingsAfter("fast");
}
daemonClosesWith(reason: string): void {
this.onClose({ reason });
}
triggerError(event?: unknown): void {
this.onError(event);
}
pingTimestamps(): string[] {
return this.pingsSentAt.map((timestamp) => `${timestamp / 1000}s`);
}
teardownCount(): number {
return this.closeEvents.filter((event) => event.reason === "Liveness check timed out").length;
}
closesFromClient(): Array<{ code?: number; reason?: string }> {
return this.closeEvents;
}
}
class DaemonClientSession {
private readonly daemon = new FakeDaemon();
private readonly client: DaemonClient;
constructor() {
this.client = new DaemonClient({
url: "ws://test",
clientId: "clsk_heartbeat_test",
logger: noopLogger,
reconnect: { enabled: false },
transportFactory: () => this.daemon.transport,
});
clients.push(this.client);
}
async connect(): Promise<void> {
const connection = this.client.connect();
this.daemon.openConnection();
await connection;
}
async advance(ms: number): Promise<void> {
await vi.advanceTimersByTimeAsync(ms);
}
daemonAnswersPingsAfter(delay: "fast" | `${number}s`): void {
this.daemon.daemonAnswersPingsAfter(delay);
}
daemonGoesSilent(): void {
this.daemon.daemonGoesSilent();
}
daemonWithholdsNextPongThenAnswersFast(): void {
this.daemon.daemonWithholdsNextPongThenAnswersFast();
}
daemonClosesWith(reason: string): void {
this.daemon.daemonClosesWith(reason);
}
pingTimestamps(): string[] {
return this.daemon.pingTimestamps();
}
state(): ReturnType<DaemonClient["getConnectionState"]> {
return this.client.getConnectionState();
}
lastError(): string | null {
return this.client.lastError;
}
teardownCount(): number {
return this.daemon.teardownCount();
}
closesFromClient(): Array<{ code?: number; reason?: string }> {
return this.daemon.closesFromClient();
}
lastLivenessRttMs(): number | null {
return this.client.getLastLivenessRttMs();
}
measureLatency(input: { timeoutMs: number }): Promise<number> {
return this.client.measureLatency(input);
}
}
function useHeartbeatClock(): void {
vi.useFakeTimers({
toFake: ["Date", "setTimeout", "clearTimeout", "setInterval", "clearInterval", "performance"],
});
}
test("dedupes in-flight checkout status requests per agentId", async () => {
const logger = createMockLogger();
const mock = createMockTransport();
@@ -316,59 +522,205 @@ test("keeps the transport connected when a session RPC ping times out", async ()
expect(client.getConnectionState().status).toBe("connected");
});
test("reconnects after repeated top-level liveness checks time out", async () => {
const logger = createMockLogger();
const mock = createMockTransport();
test("stays online through ten minutes of pongs that arrive five seconds late", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonAnswersPingsAfter("5.5s");
const client = new DaemonClient({
url: "ws://test",
clientId: "clsk_unit_test",
logger,
reconnect: { enabled: false },
transportFactory: () => mock.transport,
});
clients.push(client);
await session.connect();
await session.advance(10 * 60 * 1000);
const connectPromise = client.connect();
mock.triggerOpen();
await connectPromise;
expect(client.getConnectionState().status).toBe("connected");
await expect(client.checkLiveness({ timeoutMs: 1 })).rejects.toThrow("Liveness check timed out");
expect(client.getConnectionState().status).toBe("connected");
await expect(client.checkLiveness({ timeoutMs: 1 })).rejects.toThrow("Liveness check timed out");
expect(client.getConnectionState()).toEqual({
status: "disconnected",
reason: "Liveness check timed out (1ms)",
});
expect(session.state()).toEqual({ status: "connected" });
expect(session.teardownCount()).toBe(0);
expect(session.lastError()).toBeNull();
expect(session.pingTimestamps().length).toBeGreaterThan(0);
expect(session.lastLivenessRttMs()).toBe(5500);
});
test("resets liveness failures after a top-level pong", async () => {
const logger = createMockLogger();
const mock = createMockTransport();
test("tears down and reports a liveness timeout after the daemon goes silent for two cycles", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonGoesSilent();
const client = new DaemonClient({
url: "ws://test",
clientId: "clsk_unit_test",
logger,
reconnect: { enabled: false },
transportFactory: () => mock.transport,
await session.connect();
await session.advance(51_000);
expect(session.state()).toEqual({
status: "disconnected",
reason: "Liveness check timed out (15000ms)",
});
clients.push(client);
expect(session.teardownCount()).toBe(1);
});
const connectPromise = client.connect();
mock.triggerOpen();
await connectPromise;
test("survives a single missed pong when answers resume", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonWithholdsNextPongThenAnswersFast();
await expect(client.checkLiveness({ timeoutMs: 1 })).rejects.toThrow("Liveness check timed out");
await session.connect();
await session.advance(40_000);
const healthyProbe = client.checkLiveness({ timeoutMs: 100 });
mock.triggerMessage(JSON.stringify({ type: "pong" }));
await expect(healthyProbe).resolves.toEqual({ rttMs: expect.any(Number) });
expect(session.state()).toEqual({ status: "connected" });
expect(session.teardownCount()).toBe(0);
});
await expect(client.checkLiveness({ timeoutMs: 1 })).rejects.toThrow("Liveness check timed out");
expect(client.getConnectionState().status).toBe("connected");
test("keeps the connection proven online from heartbeats alone when no other traffic flows", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonAnswersPingsAfter("fast");
await session.connect();
await session.advance(35_000);
expect(session.state()).toEqual({ status: "connected" });
expect(session.pingTimestamps()).toEqual(["10s", "20s", "30s"]);
});
test("starts pinging one interval after connecting and holds the cadence", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonAnswersPingsAfter("fast");
await session.connect();
await session.advance(9_999);
expect(session.pingTimestamps()).toEqual([]);
await session.advance(20_001);
expect(session.pingTimestamps()).toEqual(["10s", "20s", "30s"]);
});
test("stops pinging once the connection is gone", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonAnswersPingsAfter("fast");
await session.connect();
await session.advance(10_000);
session.daemonClosesWith("daemon shutting down");
await session.advance(30_000);
expect(session.pingTimestamps()).toEqual(["10s"]);
});
test("sends only one ping while a slow pong is still outstanding", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonAnswersPingsAfter("15.1s");
await session.connect();
await session.advance(20_000);
expect(session.pingTimestamps()).toEqual(["10s"]);
});
test("reports the round-trip time of the last successful heartbeat", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonAnswersPingsAfter("2.3s");
await session.connect();
await session.advance(12_300);
expect(session.lastLivenessRttMs()).toBe(2300);
});
test("treats a pong just under the timeout as alive and just over as a miss", async () => {
useHeartbeatClock();
const alive = new DaemonClientSession();
alive.daemonAnswersPingsAfter("14.9s");
await alive.connect();
await alive.advance(24_900);
expect(alive.state()).toEqual({ status: "connected" });
const missed = new DaemonClientSession();
missed.daemonAnswersPingsAfter("15.1s");
await missed.connect();
await missed.advance(25_000);
expect(missed.state()).toEqual({ status: "connected" });
expect(missed.lastLivenessRttMs()).toBeNull();
});
test("goes red with the daemon's reason when the connection is closed", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
await session.connect();
session.daemonClosesWith("Control unresponsive");
expect(session.state()).toEqual({
status: "disconnected",
reason: "Control unresponsive",
});
expect(session.lastError()).toBe("Control unresponsive");
expect(session.closesFromClient()).toEqual([]);
});
test("two candidate latency timeouts do not tear down the live connection", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonGoesSilent();
await session.connect();
const firstMeasurement = session.measureLatency({ timeoutMs: 1000 });
const firstMeasurementError = firstMeasurement.then(
() => null,
(error) => error,
);
await session.advance(1000);
await expect(firstMeasurementError).resolves.toEqual(
expect.objectContaining({
message: "Latency measurement timed out (1000ms)",
}),
);
const secondMeasurement = session.measureLatency({ timeoutMs: 1000 });
const secondMeasurementError = secondMeasurement.then(
() => null,
(error) => error,
);
await session.advance(1000);
await expect(secondMeasurementError).resolves.toEqual(
expect.objectContaining({
message: "Latency measurement timed out (1000ms)",
}),
);
expect(session.state()).toEqual({ status: "connected" });
expect(session.teardownCount()).toBe(0);
expect(session.pingTimestamps()).toEqual(["0s", "1s"]);
});
test("a candidate measurement that times out under a heartbeat tick does not count toward teardown", async () => {
useHeartbeatClock();
const session = new DaemonClientSession();
session.daemonGoesSilent();
await session.connect();
// A candidate measurement is still in flight when the +10s heartbeat tick lands,
// so the heartbeat shares the in-flight ping. Let the measurement time out.
await session.advance(9_000);
const measurement = session.measureLatency({ timeoutMs: 5_000 });
const measurementError = measurement.then(
() => null,
(error) => error,
);
await session.advance(5_500);
await expect(measurementError).resolves.toEqual(
expect.objectContaining({ message: "Latency measurement timed out (5000ms)" }),
);
// The measurement timeout must not have been recorded as a liveness failure: a
// single genuine heartbeat miss after it must still leave the connection up.
await session.advance(25_000);
expect(session.state()).toEqual({ status: "connected" });
expect(session.teardownCount()).toBe(0);
});
test("listDirectory sends a list file explorer request and returns directory entries", async () => {
@@ -1389,7 +1741,7 @@ test("restartServer remains restart-only and sends restart_server_request", asyn
});
test("transitions out of connecting when connect timeout elapses", async () => {
vi.useFakeTimers();
useHeartbeatClock();
try {
const logger = createMockLogger();
const mock = createMockTransport();
@@ -1426,7 +1778,7 @@ test("transitions out of connecting when connect timeout elapses", async () => {
});
test("reconnects after relay close with replaced-by-new-connection reason", async () => {
vi.useFakeTimers();
useHeartbeatClock();
try {
const logger = createMockLogger();
const first = createMockTransport();
@@ -2658,7 +3010,7 @@ test("imports an agent by provider handle id", async () => {
});
test("uses server-provided dictation finish timeout budget", async () => {
vi.useFakeTimers();
useHeartbeatClock();
const logger = createMockLogger();
const mock = createMockTransport();
@@ -2748,7 +3100,7 @@ test("resolves dictation finish when final arrives after finish accepted", async
});
test("cancels waiters when send fails (no leaked timeouts)", async () => {
vi.useFakeTimers();
useHeartbeatClock();
const logger = createMockLogger();
const mock = createMockTransport();
let sendCount = 0;
@@ -2783,7 +3135,7 @@ test("cancels waiters when send fails (no leaked timeouts)", async () => {
const internal = client as unknown as { waiters: Set<unknown> };
expect(internal.waiters.size).toBe(0);
vi.runOnlyPendingTimers();
await vi.advanceTimersByTimeAsync(0);
vi.useRealTimers();
});
@@ -3689,7 +4041,7 @@ test("sends close_items_request and resolves close_items_response", async () =>
});
test("waitForFinish with timeout=0 omits timeoutMs and has no client deadline", async () => {
vi.useFakeTimers();
useHeartbeatClock();
try {
const logger = createMockLogger();
const mock = createMockTransport();
@@ -3715,14 +4067,20 @@ test("waitForFinish with timeout=0 omits timeoutMs and has no client deadline",
expect(request.agentId).toBe("agent-wait-zero-timeout");
expect(request).not.toHaveProperty("timeoutMs");
const settled = vi.fn();
let settled: "pending" | "resolved" | "rejected" = "pending";
void waitPromise.then(
() => settled("resolved"),
() => settled("rejected"),
() => {
settled = "resolved";
return null;
},
() => {
settled = "rejected";
return null;
},
);
await vi.advanceTimersByTimeAsync(5 * 60 * 1000);
expect(settled).not.toHaveBeenCalled();
expect(settled).toBe("pending");
mock.triggerMessage(
wrapSessionMessage({