Recover stale daemon connections with liveness probes (#1124)

Use top-level WebSocket ping/pong for client liveness so session RPC timeouts stay scoped to the operation that timed out. Back off inactive connection probes while a host already has a healthy active connection.
This commit is contained in:
Mohamed Boudra
2026-05-21 11:43:22 +08:00
committed by GitHub
parent 28415b3542
commit dee8f485de
5 changed files with 285 additions and 2 deletions
+2
View File
@@ -146,6 +146,8 @@ There is no dedicated welcome message; the server emits a `status` session messa
**Top-level WS envelopes** are `hello`, `recording_state`, `ping`/`pong`, and `session` (which wraps the rich union of session messages).
Client liveness checks use the top-level JSON `ping`/`pong` envelope, not a session RPC and not RFC6455 protocol ping. The app runs through browser and React Native WebSocket APIs, which do not expose protocol ping, so this envelope is the portable way to test the direct or relay data path. Session RPC timeouts are operation failures and must not be treated as proof that the socket is dead.
New session RPCs use dotted names with `.request` and `.response` suffixes, such as `checkout.github.set_auto_merge.request` and `checkout.github.set_auto_merge.response`. See [rpc-namespacing.md](rpc-namespacing.md) for the convention and migration rules for older flat RPC names.
**Notable session message types:**
@@ -76,6 +76,10 @@ class FakeDaemonClient {
return { rttMs: 0 };
}
async checkLiveness(): Promise<{ rttMs: number }> {
return this.ping();
}
setReconnectEnabled(_enabled: boolean): void {}
setConnectionState(next: ConnectionState): void {
@@ -620,6 +624,85 @@ describe("HostRuntimeController", () => {
expect((initialClient as unknown as FakeDaemonClient | null)?.closeCalls).toBe(1);
});
it("uses liveness probes instead of session RPC timeouts for active connection health", async () => {
const relay: HostConnection = {
id: "relay:relay.paseo.sh:443",
type: "relay",
relayEndpoint: "relay.paseo.sh:443",
useTls: true,
daemonPublicKeyB64: "pk_test",
};
const host = makeHost({
connections: [relay],
preferredConnectionId: relay.id,
});
const clients: FakeDaemonClient[] = [];
const controller = new HostRuntimeController({
host,
deps: makeDeps({ [relay.id]: 12 }, clients),
});
await controller.start({ autoProbe: false });
expect(controller.getSnapshot().connectionStatus).toBe("online");
const activeClient = controller.getSnapshot().client as unknown as FakeDaemonClient;
activeClient.ping = async () => {
throw new Error("Timeout waiting for message (5000ms)");
};
activeClient.checkLiveness = async () => ({ rttMs: 11 });
clearProbeBackoff(controller);
await controller.runProbeCycleNow();
const snapshot = controller.getSnapshot();
expect(snapshot.probeByConnectionId.get(relay.id)).toEqual({
status: "available",
latencyMs: 11,
});
expect(snapshot.connectionStatus).toBe("online");
});
it("backs off inactive connection probes while a host is online", async () => {
const host = makeHost({ preferredConnectionId: "direct:lan:6767" });
const clients: FakeDaemonClient[] = [];
const latencies: Record<string, number | Error> = {
"direct:lan:6767": 10,
"relay:relay.paseo.sh:443": 50,
};
const controller = new HostRuntimeController({
host,
deps: makeDeps(latencies, clients),
});
await controller.start({ autoProbe: false });
expect(controller.getSnapshot().activeConnectionId).toBe("direct:lan:6767");
const initialClientCount = clients.length;
const initialRelayProbe = controller
.getSnapshot()
.probeByConnectionId.get("relay:relay.paseo.sh:443");
latencies["direct:lan:6767"] = 12;
latencies["relay:relay.paseo.sh:443"] = 25;
const lastProbedAt = (
controller as unknown as {
connectionLastProbedAt: Map<string, number>;
}
).connectionLastProbedAt;
const now = performance.now();
lastProbedAt.set("direct:lan:6767", now - 60_000);
lastProbedAt.set("relay:relay.paseo.sh:443", now - 60_000);
await controller.runProbeCycleNow();
const snapshot = controller.getSnapshot();
expect(clients.length).toBe(initialClientCount);
expect(snapshot.probeByConnectionId.get("direct:lan:6767")).toEqual({
status: "available",
latencyMs: 12,
});
expect(snapshot.probeByConnectionId.get("relay:relay.paseo.sh:443")).toEqual(initialRelayProbe);
});
it("switches only after the faster alternative wins consecutive probes", async () => {
const host = makeHost({ preferredConnectionId: "direct:lan:6767" });
const clients: FakeDaemonClient[] = [];
+13 -2
View File
@@ -142,6 +142,7 @@ export interface HostRuntimeStartOptions {
const PROBE_TICK_MS = 2_000;
const PROBE_STEADY_MS = 10_000;
const PROBE_MAX_BACKOFF_MS = 30_000;
const PROBE_INACTIVE_WHILE_ONLINE_MS = 120_000;
const ADAPTIVE_SWITCH_THRESHOLD_MS = 40;
const ADAPTIVE_SWITCH_CONSECUTIVE_PROBES = 3;
const DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT = 200;
@@ -434,11 +435,15 @@ function findConnectionById(host: HostProfile, connectionId: string | null): Hos
function probeIntervalForConnection(
firstSeenAt: number,
isActiveOnline: boolean,
hasActiveOnlineConnection: boolean,
now: number,
): number {
if (isActiveOnline) {
return PROBE_STEADY_MS;
}
if (hasActiveOnlineConnection) {
return PROBE_INACTIVE_WHILE_ONLINE_MS;
}
const age = now - firstSeenAt;
if (age < 10_000) return 2_000;
if (age < 30_000) return 5_000;
@@ -702,6 +707,7 @@ export class HostRuntimeController {
const now = performance.now();
const isOnline = this.snapshot.connectionStatus === "online";
const activeConnectionId = this.snapshot.activeConnectionId;
const hasActiveOnlineConnection = isOnline && activeConnectionId !== null;
const connectionsToProbe = this.host.connections.filter((connection) => {
const lastProbed = this.connectionLastProbedAt.get(connection.id);
@@ -710,7 +716,12 @@ export class HostRuntimeController {
}
const firstSeen = this.connectionFirstSeenAt.get(connection.id) ?? now;
const isActiveOnline = isOnline && connection.id === activeConnectionId;
const interval = probeIntervalForConnection(firstSeen, isActiveOnline, now);
const interval = probeIntervalForConnection(
firstSeen,
isActiveOnline,
hasActiveOnlineConnection,
now,
);
return now - lastProbed >= interval;
});
@@ -899,7 +910,7 @@ export class HostRuntimeController {
const activated = await maybeActivateFirstAvailable(connection.id, connectedClient);
shouldCloseClient = shouldCloseClient && !activated;
const { rttMs } = await connectedClient.ping({ timeoutMs: 5000 });
const { rttMs } = await connectedClient.checkLiveness({ timeoutMs: 5000 });
if (!this.isCurrentProbeRequest(requestVersion)) {
return;
}
@@ -281,6 +281,84 @@ test("does not reconnect after close when ensureConnected is called", async () =
expect(client.getConnectionState().status).toBe("disposed");
});
test("keeps the transport connected when a session RPC ping times out", async () => {
const logger = createMockLogger();
const mock = createMockTransport();
const client = new DaemonClient({
url: "ws://test",
clientId: "clsk_unit_test",
logger,
reconnect: { enabled: false },
transportFactory: () => mock.transport,
});
clients.push(client);
const connectPromise = client.connect();
mock.triggerOpen();
await connectPromise;
expect(client.getConnectionState().status).toBe("connected");
await expect(client.ping({ timeoutMs: 1 })).rejects.toThrow("Timeout waiting for message");
expect(client.getConnectionState().status).toBe("connected");
});
test("reconnects after repeated top-level liveness checks time out", async () => {
const logger = createMockLogger();
const mock = createMockTransport();
const client = new DaemonClient({
url: "ws://test",
clientId: "clsk_unit_test",
logger,
reconnect: { enabled: false },
transportFactory: () => mock.transport,
});
clients.push(client);
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)",
});
});
test("resets liveness failures after a top-level pong", async () => {
const logger = createMockLogger();
const mock = createMockTransport();
const client = new DaemonClient({
url: "ws://test",
clientId: "clsk_unit_test",
logger,
reconnect: { enabled: false },
transportFactory: () => mock.transport,
});
clients.push(client);
const connectPromise = client.connect();
mock.triggerOpen();
await connectPromise;
await expect(client.checkLiveness({ timeoutMs: 1 })).rejects.toThrow("Liveness check timed out");
const healthyProbe = client.checkLiveness({ timeoutMs: 100 });
mock.triggerMessage(JSON.stringify({ type: "pong" }));
await expect(healthyProbe).resolves.toEqual({ rttMs: expect.any(Number) });
await expect(client.checkLiveness({ timeoutMs: 1 })).rejects.toThrow("Liveness check timed out");
expect(client.getConnectionState().status).toBe("connected");
});
test("listDirectory sends a list file explorer request and returns directory entries", async () => {
const logger = createMockLogger();
const mock = createMockTransport();
+109
View File
@@ -721,6 +721,8 @@ class DaemonRpcError extends Error {
const DEFAULT_RECONNECT_BASE_DELAY_MS = 1500;
const DEFAULT_RECONNECT_MAX_DELAY_MS = 30000;
const DEFAULT_CONNECT_TIMEOUT_MS = 15000;
const DEFAULT_LIVENESS_TIMEOUT_MS = 5000;
const LIVENESS_FAILURE_RECONNECT_THRESHOLD = 2;
/** Default timeout for waiting for connection before sending queued messages */
const DEFAULT_SEND_QUEUE_TIMEOUT_MS = 10000;
@@ -827,6 +829,14 @@ interface PendingSend {
timeoutHandle: ReturnType<typeof setTimeout>;
}
interface LivenessProbe {
promise: Promise<{ rttMs: number }>;
resolve: (value: { rttMs: number }) => void;
reject: (error: Error) => void;
timeoutHandle: ReturnType<typeof setTimeout>;
startedAt: number;
}
export class DaemonClient {
private transport: DaemonTransport | null = null;
private transportCleanup: Array<() => void> = [];
@@ -870,6 +880,8 @@ export class DaemonClient {
private lastServerInfoMessage: ServerInfoStatusPayload | null = null;
private runtimeMetricsInterval: ReturnType<typeof setInterval> | null = null;
private runtimeMetrics: DaemonClientRuntimeMetrics | null = null;
private livenessProbe: LivenessProbe | null = null;
private consecutiveLivenessFailures = 0;
constructor(private config: DaemonClientConfig) {
this.logger = config.logger ?? consoleLogger;
@@ -1131,6 +1143,7 @@ export class DaemonClient {
this.disposeTransport(1000, "Client closed");
this.clearWaiters(new Error("Daemon client closed"));
this.rejectPendingSendQueue(new Error("Daemon client closed"));
this.rejectLivenessProbe(new Error("Daemon client closed"));
this.terminalStreams.clearSlots();
this.lastServerInfoMessage = null;
if (this.runtimeMetricsInterval) {
@@ -1566,6 +1579,54 @@ export class DaemonClient {
};
}
checkLiveness(params?: { timeoutMs?: number }): Promise<{ rttMs: number }> {
if (this.connectionState.status !== "connected" || !this.transport) {
return Promise.reject(
new Error(`Transport not connected (status: ${this.connectionState.status})`),
);
}
if (this.livenessProbe) {
return this.livenessProbe.promise;
}
const startedAt = perfNow();
const timeoutMs = Math.max(1, params?.timeoutMs ?? DEFAULT_LIVENESS_TIMEOUT_MS);
let resolveProbe: ((value: { rttMs: number }) => void) | null = null;
let rejectProbe: ((error: Error) => void) | null = null;
const promise = new Promise<{ rttMs: number }>((resolve, reject) => {
resolveProbe = resolve;
rejectProbe = reject;
});
const probe: LivenessProbe = {
promise,
resolve: (value) => resolveProbe?.(value),
reject: (error) => rejectProbe?.(error),
timeoutHandle: setTimeout(() => {
if (this.livenessProbe !== probe) {
return;
}
this.livenessProbe = null;
const error = new Error(`Liveness check timed out (${timeoutMs}ms)`);
probe.reject(error);
this.recordLivenessFailure(error);
}, timeoutMs),
startedAt,
};
this.livenessProbe = probe;
try {
this.transport.send(JSON.stringify({ type: "ping" }));
} catch (error) {
this.clearLivenessProbe();
const err = error instanceof Error ? error : new Error(String(error));
this.recordLivenessFailure(err);
return Promise.reject(err);
}
return promise;
}
// ============================================================================
// Agent RPCs (requestId-correlated)
// ============================================================================
@@ -4267,7 +4328,10 @@ export class DaemonClient {
return;
}
this.consecutiveLivenessFailures = 0;
if (parsed.data.type === "pong") {
this.resolveLivenessProbe();
this.runtimeMetrics?.recordMessage("pong", bytes, perfNow() - startMs);
return;
}
@@ -4416,6 +4480,7 @@ export class DaemonClient {
// and responses from the previous connection will never arrive.
this.clearWaiters(new Error(reason ?? "Connection lost"));
this.rejectPendingSendQueue(new Error(reason ?? "Connection lost"));
this.rejectLivenessProbe(new Error(reason ?? "Connection lost"));
this.terminalStreams.clearSlots();
this.lastServerInfoMessage = null;
@@ -4464,6 +4529,50 @@ export class DaemonClient {
}, delay);
}
private resolveLivenessProbe(): void {
const probe = this.livenessProbe;
if (!probe) {
return;
}
this.livenessProbe = null;
clearTimeout(probe.timeoutHandle);
probe.resolve({ rttMs: perfNow() - probe.startedAt });
}
private clearLivenessProbe(): void {
const probe = this.livenessProbe;
if (!probe) {
return;
}
this.livenessProbe = null;
clearTimeout(probe.timeoutHandle);
}
private rejectLivenessProbe(error: Error): void {
const probe = this.livenessProbe;
if (!probe) {
return;
}
this.livenessProbe = null;
clearTimeout(probe.timeoutHandle);
probe.reject(error);
}
private recordLivenessFailure(error: Error): void {
this.consecutiveLivenessFailures += 1;
if (this.consecutiveLivenessFailures < LIVENESS_FAILURE_RECONNECT_THRESHOLD) {
return;
}
this.consecutiveLivenessFailures = 0;
this.lastErrorValue = error.message;
this.disposeTransport(1001, "Liveness check timed out");
this.scheduleReconnect({
reason: error.message,
event: "LIVENESS_TIMEOUT",
reasonCode: "liveness_timeout",
});
}
private handleSessionMessage(msg: SessionOutboundMessage): void {
if (msg.type === "status") {
const serverInfo = parseServerInfoStatusPayload(msg.payload);