jb/paseo
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Recover stale daemon connections with liveness probes (#1124)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
packages/server/src/client/daemon-client.test.ts
+78
View File
@@ -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();
packages/server/src/client/daemon-client.ts
+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);