feat(remote-control): cache rewritten tunnel responses with ETag validation

.changeset/rc-tunnel-cache.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Reuse unchanged Remote Control assets across page loads instead of retransferring them.

packages/remote-control/src/remote-control.ts

@@ -1,3 +1,4 @@
+import { createHash } from 'node:crypto';
 import { hostname, platform } from 'node:os';
 import { join } from 'node:path';
 import { request as httpRequest, validateHeaderName, validateHeaderValue } from 'node:http';
@@ -248,6 +249,26 @@ function isGzipCompressibleType(contentType: string): boolean {
   return mime.startsWith('text/') || GZIP_COMPRESSIBLE_TYPES.has(mime);
 }
 
+function rewrittenResponseETag(body: Buffer): string {
+  return `W/"${createHash('sha256').update(body).digest('hex')}"`;
+}
+
+function requestMatchesETag(
+  headers: readonly [string, string][],
+  etag: string,
+): boolean {
+  const candidates = [etag, etag.replace(/^W\//, '')];
+  for (const [name, value] of headers) {
+    if (name.toLowerCase() !== 'if-none-match') continue;
+    for (const token of value.split(',')) {
+      const candidate = token.trim();
+      if (candidate === '*') return true;
+      if (candidates.includes(candidate)) return true;
+    }
+  }
+  return false;
+}
+
 export async function startRemoteControl(
   options: RemoteControlOptions,
 ): Promise<RemoteControlHandle> {
@@ -856,7 +877,18 @@ function requestLocalHttp(
                 : receivedBody;
             const rewritten = body !== receivedBody;
             const headers = filterResponseHeaders(response.rawHeaders, rewritten);
-            if (rewritten) headers.push('Cache-Control', 'no-cache');
+            if (rewritten) {
+              const etag = rewrittenResponseETag(body);
+              headers.push('Cache-Control', 'no-cache', 'ETag', etag);
+              const statusCode = response.statusCode ?? 502;
+              const revalidatable =
+                (parsed.method === 'GET' || parsed.method === 'HEAD') &&
+                statusCode >= 200 &&
+                statusCode < 300;
+              if (revalidatable && requestMatchesETag(parsed.headers, etag)) {
+                return Buffer.from(`HTTP/1.1 304 Not Modified\r\n${headerLines(headers)}\r\n\r\n`);
+              }
+            }
             const negotiated =
               response.headers['content-encoding'] === undefined &&
               response.statusCode !== 206 &&
@@ -877,9 +909,6 @@ function requestLocalHttp(
             if (negotiated && acceptsGzipEncoding(parsed.headers)) {
               body = await gzipAsync(body);
               headers.push('Content-Encoding', 'gzip');
-              for (let index = headers.length - 2; index >= 0; index -= 2) {
-                if (headers[index]!.toLowerCase() === 'etag') headers.splice(index, 2);
-              }
             }
             headers.push('Content-Length', String(body.length));
             const statusCode = response.statusCode ?? 502;
@@ -898,7 +927,7 @@ function requestLocalHttp(
   });
 }
 
-function filterResponseHeaders(rawHeaders: readonly string[], blockCacheControl = false): string[] {
+function filterResponseHeaders(rawHeaders: readonly string[], blockCacheValidators = false): string[] {
   const connectionHeaders = new Set<string>();
   for (let index = 0; index < rawHeaders.length; index += 2) {
     if (rawHeaders[index]!.toLowerCase() === 'connection') {
@@ -914,7 +943,12 @@ function filterResponseHeaders(rawHeaders: readonly string[], blockCacheControl
     if (BLOCKED_RESPONSE_HEADERS.has(lower) || connectionHeaders.has(lower)) {
       continue;
     }
-    if (blockCacheControl && lower === 'cache-control') continue;
+    if (
+      blockCacheValidators &&
+      (lower === 'cache-control' || lower === 'etag' || lower === 'last-modified')
+    ) {
+      continue;
+    }
     result.push(name, rawHeaders[index + 1]!);
   }
   return result;

packages/remote-control/test/remote-control.test.ts

@@ -455,12 +455,38 @@ describe('Remote Control tunnel', () => {
     expect(gzipHead).toContain('HTTP/1.1 200 OK');
     expect(gzipHead).toContain('Content-Encoding: gzip');
     expect(gzipHead).toContain('Vary: Accept-Encoding');
-    expect(gzipHead).not.toContain('ETag');
+    expect(gzipHead).toContain('Cache-Control: no-cache');
+    const rewrittenETag = /ETag: (W\/"[0-9a-f]{64}")/.exec(gzipHead)?.[0];
+    expect(rewrittenETag).toBeDefined();
     expect(gzipHead).toContain(`Content-Length: ${gzipBody.length}`);
     expect(gunzipSync(gzipBody).toString()).toBe(
       assetJs.replaceAll('"/assets/', `"/coding-relay/devices/${handle.deviceId}/assets/`),
     );
 
+    const revalidateResponsePromise = nextJsonMessage(httpConnections[0]!);
+    httpConnections[0]!.send(
+      JSON.stringify({
+        request_id: 'request-3b',
+        type: 'request',
+        is_last: true,
+        body_base64: Buffer.from(
+          `GET /assets/index.js HTTP/1.1\r\nHost: relay.test\r\nAccept-Encoding: br, gzip\r\nIf-None-Match: ${rewrittenETag!.replace('ETag: ', '')}\r\n\r\n`,
+        ).toString('base64'),
+      }),
+    );
+    const revalidateResponse = Buffer.from(
+      (await revalidateResponsePromise)['body_base64'] as string,
+      'base64',
+    );
+    const revalidateHead = revalidateResponse
+      .subarray(0, revalidateResponse.indexOf('\r\n\r\n'))
+      .toString('latin1');
+    expect(revalidateHead).toContain('HTTP/1.1 304 Not Modified');
+    expect(revalidateHead).toContain('Cache-Control: no-cache');
+    expect(revalidateHead).toContain(rewrittenETag!);
+    expect(revalidateHead).not.toContain('Content-Encoding');
+    expect(revalidateHead).not.toContain('Content-Length');
+
     const binaryResponsePromise = nextJsonMessage(httpConnections[0]!);
     httpConnections[0]!.send(
       JSON.stringify({
@@ -501,7 +527,8 @@ describe('Remote Control tunnel', () => {
     const excludedHead = excludedResponse.subarray(0, excludedSeparator).toString('latin1');
     expect(excludedHead).not.toContain('Content-Encoding');
     expect(excludedHead).toContain('Vary: Accept-Encoding');
-    expect(excludedHead).toContain('ETag: "v1"');
+    expect(excludedHead).toContain(rewrittenETag!);
+    expect(excludedHead).not.toContain('ETag: "v1"');
     expect(excludedResponse.subarray(excludedSeparator + 4).toString()).toBe(
       assetJs.replaceAll('"/assets/', `"/coding-relay/devices/${handle.deviceId}/assets/`),
     );