From 4d034dce0e5cb260b092bb27817ef756a74bb86b Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Thu, 30 Jul 2026 14:07:38 -0700
Subject: [PATCH] fix(gateway): frame upgrade rate-limit responses (#116557)

---
 src/gateway/server-http.ts                    | 20 +++---
 ...server.plugin-node-capability-auth.test.ts | 68 ++++++++++++++++++-
 2 files changed, 76 insertions(+), 12 deletions(-)

diff --git a/src/gateway/server-http.ts b/src/gateway/server-http.ts
index d7c27362779b..9439e68f621c 100644
--- a/src/gateway/server-http.ts
+++ b/src/gateway/server-http.ts
@@ -330,22 +330,22 @@ function writeUpgradeAuthFailure(
   if (auth.rateLimited) {
     const retryAfterSeconds =
       auth.retryAfterMs && auth.retryAfterMs > 0 ? Math.ceil(auth.retryAfterMs / 1000) : undefined;
+    const body = JSON.stringify({
+      error: {
+        message: "Too many failed authentication attempts. Please try again later.",
+        type: "rate_limited",
+      },
+    });
     socket.write(
       [
         "HTTP/1.1 429 Too Many Requests",
-        retryAfterSeconds ? `Retry-After: ${retryAfterSeconds}` : undefined,
+        ...(retryAfterSeconds ? [`Retry-After: ${retryAfterSeconds}`] : []),
         "Content-Type: application/json; charset=utf-8",
+        `Content-Length: ${Buffer.byteLength(body, "utf8")}`,
         "Connection: close",
         "",
-        JSON.stringify({
-          error: {
-            message: "Too many failed authentication attempts. Please try again later.",
-            type: "rate_limited",
-          },
-        }),
-      ]
-        .filter(Boolean)
-        .join("\r\n"),
+        body,
+      ].join("\r\n"),
     );
     return;
   }
diff --git a/src/gateway/server.plugin-node-capability-auth.test.ts b/src/gateway/server.plugin-node-capability-auth.test.ts
index bc76647693fa..8a7615c342ed 100644
--- a/src/gateway/server.plugin-node-capability-auth.test.ts
+++ b/src/gateway/server.plugin-node-capability-auth.test.ts
@@ -1,6 +1,6 @@
 // Plugin node capability auth tests cover scoped canvas/A2UI HTTP and WebSocket
 // routes, preauth budgets, capability paths, and unauthorized upgrade handling.
-import type { IncomingMessage, ServerResponse } from "node:http";
+import { request, type IncomingMessage, type ServerResponse } from "node:http";
 import { connect, type Socket } from "node:net";
 import type { Duplex } from "node:stream";
 import { describe, expect, test } from "vitest";
@@ -118,6 +118,53 @@ async function expectWsRejected(
   });
 }
 
+async function requestWsUpgradeResponse(params: {
+  port: number;
+  path: string;
+  headers: Record<string, string>;
+}): Promise<{
+  statusCode: number;
+  headers: IncomingMessage["headers"];
+  body: string;
+  complete: boolean;
+}> {
+  return await new Promise((resolve, reject) => {
+    const req = request({
+      host: "127.0.0.1",
+      port: params.port,
+      path: params.path,
+      headers: {
+        ...params.headers,
+        connection: "Upgrade",
+        upgrade: "websocket",
+        "sec-websocket-key": "dGhlIHNhbXBsZSBub25jZQ==",
+        "sec-websocket-version": "13",
+      },
+    });
+    req.setTimeout(WS_REJECT_TIMEOUT_MS, () => {
+      req.destroy(new Error("timeout"));
+    });
+    req.once("response", (res) => {
+      const chunks: Buffer[] = [];
+      res.on("data", (chunk: Buffer) => chunks.push(chunk));
+      res.once("end", () => {
+        resolve({
+          statusCode: res.statusCode ?? 0,
+          headers: res.headers,
+          body: Buffer.concat(chunks).toString("utf8"),
+          complete: res.complete,
+        });
+      });
+    });
+    req.once("upgrade", (_res, socket) => {
+      socket.destroy();
+      reject(new Error("expected upgrade to reject"));
+    });
+    req.once("error", reject);
+    req.end();
+  });
+}
+
 async function expectWsConnected(url: string, headers?: Record<string, string>): Promise<void> {
   await new Promise<void>((resolve, reject) => {
     const ws = new WebSocket(url, headers ? { headers } : undefined);
@@ -660,7 +707,24 @@ describe("gateway plugin node capability auth", () => {
           const second = await expectRepeatedCanvasAuthAttemptsRateLimited(listener, headers);
           expect(second.headers.get("retry-after")).toMatch(/^\d+$/);
 
-          await expectWsRejected(`ws://127.0.0.1:${listener.port}${CANVAS_WS_PATH}`, headers, 429);
+          const upgradeResponse = await requestWsUpgradeResponse({
+            port: listener.port,
+            path: CANVAS_WS_PATH,
+            headers,
+          });
+          const expectedBody = JSON.stringify({
+            error: {
+              message: "Too many failed authentication attempts. Please try again later.",
+              type: "rate_limited",
+            },
+          });
+          expect(upgradeResponse.statusCode).toBe(429);
+          expect(upgradeResponse.headers["retry-after"]).toMatch(/^\d+$/);
+          expect(upgradeResponse.headers["content-length"]).toBe(
+            String(Buffer.byteLength(expectedBody, "utf8")),
+          );
+          expect(upgradeResponse.body).toBe(expectedBody);
+          expect(upgradeResponse.complete).toBe(true);
         },
       });
     });