fix(kimi-code): reconnect remote control tunnel after silent relay death

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<NodeJS.WriteStream, 'write'>;
   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<void> {
@@ -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<void> {
     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;

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<typeof createServer>): Promise<number> {