diff --git a/apps/kimi-code/src/cli/sub/web/remote-control.ts b/apps/kimi-code/src/cli/sub/web/remote-control.ts index 73e74f46f96..bf0f2aac449 100644 --- a/apps/kimi-code/src/cli/sub/web/remote-control.ts +++ b/apps/kimi-code/src/cli/sub/web/remote-control.ts @@ -36,6 +36,8 @@ const MAX_HTTP_REQUEST_BYTES = 10 * 1024 * 1024; const HTTP_REQUEST_TIMEOUT_MS = 30_000; const REGISTER_TIMEOUT_MS = 10_000; const MAX_RECONNECT_DELAY_MS = 30_000; +const RELAY_PING_INTERVAL_MS = 30_000; +const RELAY_SILENCE_TIMEOUT_MS = 300_000; const BLOCKED_REQUEST_HEADERS = new Set([ 'authorization', 'cookie', @@ -95,6 +97,8 @@ export interface RemoteControlOptions { readonly relayOrigin?: string; readonly stderr?: Pick; readonly onStatus?: (status: RemoteControlStatus) => void; + readonly pingIntervalMs?: number; + readonly silenceTimeoutMs?: number; } export interface RemoteControlHandle { @@ -348,6 +352,8 @@ class RemoteControlClient { private pendingHttpBytes = 0; private reconnectAttempt = 0; private reconnectImmediately = false; + private readonly pingIntervalMs: number; + private readonly silenceTimeoutMs: number; private stopped = false; private connected = false; private relayOnline = false; @@ -369,6 +375,8 @@ class RemoteControlClient { this.refreshToken = options.refreshToken; this.stderr = options.stderr ?? process.stderr; this.onStatus = options.onStatus ?? (() => {}); + this.pingIntervalMs = options.pingIntervalMs ?? RELAY_PING_INTERVAL_MS; + this.silenceTimeoutMs = options.silenceTimeoutMs ?? RELAY_SILENCE_TIMEOUT_MS; } async start(): Promise { @@ -403,12 +411,13 @@ class RemoteControlClient { await this.serveCycle(); } catch (error) { if (error instanceof RegistrationError) { - if (!this.connected) this.rejectInitial(error); - else this.stderr.write(`${error.message}\n`); - this.stopped = true; - return; - } - if (!this.stopped && !this.reconnectImmediately) { + if (!this.connected) { + this.rejectInitial(error); + this.stopped = true; + return; + } + this.stderr.write(`${error.message}\n`); + } else if (!this.stopped && !this.reconnectImmediately) { this.stderr.write(`Remote Control disconnected: ${errorMessage(error)}\n`); } } finally { @@ -434,6 +443,7 @@ class RemoteControlClient { private async serveCycle(): Promise { const management = await this.connectRelay('/v1/remote/create'); this.management = management; + this.watchSocket(management, 'management'); management.send( JSON.stringify({ type: 'register', @@ -461,6 +471,7 @@ class RemoteControlClient { `/v1/remote/http?device_id=${encodeURIComponent(this.deviceId)}`, ); this.http = http; + this.watchSocket(http, 'http'); if (management.readyState !== WebSocket.OPEN) { throw new Error('management connection closed'); } @@ -485,6 +496,32 @@ class RemoteControlClient { return connectWebSocket(relayWebSocketUrl(this.relayOrigin, path), this.refreshToken); } + private watchSocket(socket: WebSocket, label: string): void { + const pingTimer = setInterval(() => { + if (socket.readyState === WebSocket.OPEN) socket.ping(); + }, this.pingIntervalMs); + pingTimer.unref(); + let silenceTimer: NodeJS.Timeout | undefined; + const armSilenceTimer = (): void => { + if (silenceTimer !== undefined) clearTimeout(silenceTimer); + silenceTimer = setTimeout(() => { + this.stderr.write( + `Remote Control ${label} connection silent for ${Math.round(this.silenceTimeoutMs / 1000)}s; reconnecting…\n`, + ); + socket.terminate(); + }, this.silenceTimeoutMs); + silenceTimer.unref(); + }; + armSilenceTimer(); + socket.on('message', armSilenceTimer); + socket.on('ping', armSilenceTimer); + socket.on('pong', armSilenceTimer); + socket.once('close', () => { + clearInterval(pingTimer); + if (silenceTimer !== undefined) clearTimeout(silenceTimer); + }); + } + private rejectInitial(error: Error): void { this.initialReject?.(error); this.initialReject = undefined; diff --git a/apps/kimi-code/test/cli/web/remote-control.test.ts b/apps/kimi-code/test/cli/web/remote-control.test.ts index cbad20c4c2f..133f2a5680d 100644 --- a/apps/kimi-code/test/cli/web/remote-control.test.ts +++ b/apps/kimi-code/test/cli/web/remote-control.test.ts @@ -476,6 +476,61 @@ describe('Remote Control tunnel', () => { ), ); }); + + it('reconnects when the relay goes silent without closing the sockets', async () => { + const homeDir = await createRemoteControlHome(TOKEN.refreshToken); + const relay = await startAuthRelay(); + let handle: RemoteControlHandle | undefined; + cleanups.push(async () => handle?.close()); + let logs = ''; + + handle = await startRemoteControl({ + homeDir, + localOrigin: 'http://127.0.0.1:1', + localServerToken: 'local-server-token', + relayOrigin: `http://127.0.0.1:${relay.port}/coding-relay`, + stderr: { write: (text) => ((logs += String(text)), true) }, + pingIntervalMs: 50, + silenceTimeoutMs: 300, + }); + + expect(relay.registrations).toHaveLength(1); + relay.managementSockets[0]!.pause(); + relay.httpSockets[0]!.pause(); + + await waitFor(() => relay.registrations.length === 2, 10_000); + expect(logs).toContain('silent'); + relay.managementSockets[0]!.terminate(); + relay.httpSockets[0]!.terminate(); + }, 15_000); + + it('retries when registration is rejected after a reconnect', async () => { + const homeDir = await createRemoteControlHome(TOKEN.refreshToken); + const relay = await startAuthRelay({ nakRegistrationsAfterFirst: 1 }); + let handle: RemoteControlHandle | undefined; + cleanups.push(async () => handle?.close()); + let logs = ''; + + handle = await startRemoteControl({ + homeDir, + localOrigin: 'http://127.0.0.1:1', + localServerToken: 'local-server-token', + relayOrigin: `http://127.0.0.1:${relay.port}/coding-relay`, + stderr: { write: (text) => ((logs += String(text)), true) }, + }); + + expect(relay.registrations).toHaveLength(1); + relay.managementSockets[0]!.terminate(); + relay.httpSockets[0]!.terminate(); + + await waitFor(() => relay.registrations.length >= 3, 10_000); + await waitFor( + () => relay.managementSockets.some((socket) => socket.readyState === 1) && + relay.httpSockets.some((socket) => socket.readyState === 1), + ); + expect(logs).toContain('DEPLOYING'); + expect(handle.url).toContain('?rc=1&from=kimi_code_cli'); + }, 15_000); }); describe('Remote Control single-instance lock', () => { @@ -611,25 +666,49 @@ async function startAuthRelay( echoProtocol?: boolean; rejectUpgrades?: number; closeManagementDuringFirstHttpHandshake?: boolean; + nakRegistrationsAfterFirst?: number; } = {}, ): Promise<{ port: number; requests: Array<{ authorization?: string; protocol?: string }>; + registrations: unknown[]; + managementSockets: WebSocket[]; + httpSockets: WebSocket[]; }> { const handleProtocols = options.echoProtocol === false ? (): false => false : undefined; const managementServer = new WebSocketServer({ noServer: true, handleProtocols }); const httpTunnelServer = new WebSocketServer({ noServer: true, handleProtocols }); const relayServer = createServer(); const requests: Array<{ authorization?: string; protocol?: string }> = []; + const registrations: unknown[] = []; + const managementSockets: WebSocket[] = []; + const httpSockets: WebSocket[] = []; let remainingRejections = options.rejectUpgrades ?? 0; let closeManagement = options.closeManagementDuringFirstHttpHandshake === true; let delayHttpUpgrade = closeManagement; + let pendingNaks = options.nakRegistrationsAfterFirst ?? 0; managementServer.on('connection', (ws) => { + managementSockets.push(ws); ws.on('error', () => {}); ws.on('message', (data) => { const message = JSON.parse(rawDataText(data)) as { type?: string }; if (message.type === 'register') { + const isReconnectRegistration = registrations.length > 0; + registrations.push(message); + if (isReconnectRegistration && pendingNaks > 0) { + pendingNaks -= 1; + ws.send( + JSON.stringify({ + type: 'register_nak', + payload: { + error_code: 'DEPLOYING', + error_message: 'relay is restarting', + }, + }), + ); + return; + } ws.send(JSON.stringify({ type: 'register_ack', payload: { success: true } })); if (closeManagement) { closeManagement = false; @@ -638,7 +717,10 @@ async function startAuthRelay( } }); }); - httpTunnelServer.on('connection', (ws) => ws.on('error', () => {})); + httpTunnelServer.on('connection', (ws) => { + httpSockets.push(ws); + ws.on('error', () => {}); + }); relayServer.on('upgrade', (request, socket, head) => { const authorization = request.headers.authorization; const protocol = request.headers['sec-websocket-protocol']; @@ -669,7 +751,7 @@ async function startAuthRelay( }); const port = await listen(relayServer); cleanups.push(() => closeServer(relayServer)); - return { port, requests }; + return { port, requests, registrations, managementSockets, httpSockets }; } function listen(server: ReturnType): Promise {