fix: support hwpod websocket through edge proxy

2026-06-05 20:04:28 +08:00
parent acb44b5cf0
commit 13b7621019
3 changed files with 235 additions and 56 deletions
@@ -36,6 +36,19 @@ test("edge proxy reports local health and proxies live health to upstream", asyn
  }
 });

+test("edge proxy tunnels websocket upgrade to upstream", async () => {
+  const upstream = await startWebSocketUpstream();
+  const proxy = await startEdgeProxy(upstream.port);
+  try {
+    const observed = await openWebSocket(`ws://127.0.0.1:${proxy.port}/v1/hwpod-node/ws`, () => ({ proxyStderr: proxy.stderr, upstreamEvents: upstream.events }));
+    assert.equal(observed.type, "welcome");
+    assert.equal(observed.path, "/v1/hwpod-node/ws");
+  } finally {
+    await proxy.stop();
+    await upstream.stop();
+  }
+});
+
 async function startUpstream() {
  let captured = null;
  const server = createServer(async (request, response) => {
@@ -63,8 +76,66 @@ async function startUpstream() {
  };
 }

+async function startWebSocketUpstream() {
+  const events = [];
+  const server = Bun.serve({
+    hostname: "127.0.0.1",
+    port: 0,
+    fetch(request, bunServer) {
+      const url = new URL(request.url);
+      events.push({ event: "upgrade", url: url.pathname });
+      const upgraded = bunServer.upgrade(request, { data: { path: url.pathname } });
+      return upgraded ? undefined : new Response("upgrade required", { status: 426 });
+    },
+    websocket: {
+      open(socket) {
+        setTimeout(() => {
+          socket.send(JSON.stringify({ type: "welcome", path: socket.data.path }));
+        }, 50);
+      }
+    }
+  });
+  return {
+    port: server.port,
+    events,
+    stop: () => server.stop(true)
+  };
+}
+
+async function openWebSocket(url, debug = () => ({})) {
+  return await new Promise((resolve, reject) => {
+    const socket = new WebSocket(url);
+    let finished = false;
+    const done = (error, value) => {
+      if (finished) return;
+      finished = true;
+      clearTimeout(timer);
+      try { socket.close(); } catch {}
+      if (error) reject(error);
+      else resolve(value);
+    };
+    const timer = setTimeout(() => done(new Error(`websocket did not receive message: ${url}; debug=${JSON.stringify(debug())}`)), 2500);
+    socket.addEventListener("message", (event) => {
+      dataToText(event.data)
+        .then((text) => done(null, JSON.parse(text)))
+        .catch((error) => done(error));
+    });
+    socket.addEventListener("error", () => done(new Error(`websocket error: ${url}`)));
+    socket.addEventListener("close", (event) => done(new Error(`websocket closed before message: code=${event.code} reason=${event.reason || ""}`)));
+  });
+}
+
+async function dataToText(data) {
+  if (typeof data === "string") return data;
+  if (data instanceof ArrayBuffer) return Buffer.from(data).toString("utf8");
+  if (ArrayBuffer.isView(data)) return Buffer.from(data.buffer, data.byteOffset, data.byteLength).toString("utf8");
+  if (data && typeof data.text === "function") return await data.text();
+  return String(data);
+}
+
 async function startEdgeProxy(upstreamPort) {
  const port = await freePort();
+  let stderr = "";
  const child = spawn(bunCommand, ["run", "cmd/hwlab-edge-proxy/main.ts"], {
    cwd: process.cwd(),
    env: {
@@ -75,9 +146,15 @@ async function startEdgeProxy(upstreamPort) {
    },
    stdio: ["ignore", "pipe", "pipe"]
  });
+  child.stderr.on("data", (chunk) => {
+    stderr += chunk;
+  });
  await waitForListening(child, "edge proxy");
  return {
    port,
+    get stderr() {
+      return stderr;
+    },
    stop: () => stopChild(child)
  };
 }
@@ -1,12 +1,8 @@
-#!/usr/bin/env node
+#!/usr/bin/env bun
 import {
  healthPayload,
-  listen,
-  proxyHttpRequest,
  resolveHostPort,
-  sendJson
 } from "../../internal/dev-entrypoint/http.mjs";
-import { createServer } from "node:http";

 const serviceId = "hwlab-edge-proxy";
 const upstream = process.env.HWLAB_EDGE_UPSTREAM || "http://hwlab-cloud-api.hwlab-dev.svc.cluster.local:6667";
@@ -35,59 +31,149 @@ function parseTimeout(value, fallback, { min, max }) {
  return Math.min(Math.max(parsed, min), max);
 }

-const server = createServer((request, response) => {
-  const url = new URL(request.url || "/", "http://hwlab-edge-proxy.local");
+const server = Bun.serve({
+  hostname: host,
+  port,
+  async fetch(request, bunServer) {
+    const url = new URL(request.url);

-  if (request.method === "GET" && url.pathname === "/health/live") {
-    proxyHttpRequest({ request, response, upstream, timeoutMs: proxyTimeoutMs });
-    return;
-  }
+    if (isWebSocketUpgrade(request)) {
+      if (url.pathname === "/v1/internal" || url.pathname.startsWith("/v1/internal/")) return text("not found", 404);
+      const upgraded = bunServer.upgrade(request, { data: { upstreamUrl: upstreamWebSocketUrl(request, upstream), upstreamSocket: null, upstreamOpen: false, pendingMessages: [] } });
+      return upgraded ? undefined : json({ ok: false, status: "upgrade_required" }, 426);
+    }

-  if (request.method === "GET" && (url.pathname === "/health" || url.pathname === "/edge/health")) {
-    sendJson(response, 200, healthPayload({
-      serviceId,
-      role: "public-dev-ingress",
-      details: proxyDetails()
-    }));
-    return;
-  }
+    if (request.method === "GET" && url.pathname === "/health/live") return proxyFetch(request, upstream, proxyTimeoutMs);

-  if (request.method === "GET" && (url.pathname === "/status" || url.pathname === "/routes")) {
-    sendJson(response, 200, {
-      serviceId,
-      environment: "dev",
-      status: "ok",
-      upstream,
-      routes: [
-        {
-          pathPrefix: "/",
-          upstream,
-          mode: "http-proxy"
+    if (request.method === "GET" && (url.pathname === "/health" || url.pathname === "/edge/health")) {
+      return json(healthPayload({ serviceId, role: "public-dev-ingress", details: proxyDetails() }));
+    }
+
+    if (request.method === "GET" && (url.pathname === "/status" || url.pathname === "/routes")) {
+      return json({
+        serviceId,
+        environment: "dev",
+        status: "ok",
+        upstream,
+        routes: [{ pathPrefix: "/", upstream, mode: "http-proxy" }]
+      });
+    }
+
+    if (request.method === "GET" && url.pathname === "/help") {
+      return json({ serviceId, commands: ["GET /health", "GET /edge/health", "GET /health/live", "GET /status", "GET /routes"] });
+    }
+
+    if (url.pathname === "/v1/internal" || url.pathname.startsWith("/v1/internal/")) {
+      return json({ error: { code: "not_found", message: "Internal HWLAB routes are not exposed through the public edge proxy." } }, 404);
+    }
+
+    return proxyFetch(request, upstream, proxyTimeoutMs);
+  },
+  websocket: {
+    open(socket: any) {
+      const upstreamSocket = new WebSocket(socket.data.upstreamUrl);
+      socket.data.upstreamSocket = upstreamSocket;
+      upstreamSocket.addEventListener("open", () => {
+        logWsBridge("upstream_open", socket.data.upstreamUrl);
+        socket.data.upstreamOpen = true;
+        for (const message of socket.data.pendingMessages.splice(0)) upstreamSocket.send(message);
+      });
+      upstreamSocket.addEventListener("message", (event) => {
+        logWsBridge("upstream_message", socket.data.upstreamUrl);
+        try {
+          socket.send(event.data);
+        } catch {
+          logWsBridge("client_send_failed", socket.data.upstreamUrl);
+          upstreamSocket.close();
        }
-      ]
-    });
-    return;
+      });
+      upstreamSocket.addEventListener("close", (event) => {
+        logWsBridge("upstream_close", socket.data.upstreamUrl, { code: event.code, reason: event.reason || null });
+        try {
+          socket.close(event.code || 1000, event.reason || undefined);
+        } catch {}
+      });
+      upstreamSocket.addEventListener("error", () => {
+        logWsBridge("upstream_error", socket.data.upstreamUrl);
+        try {
+          socket.close(1011, "upstream websocket error");
+        } catch {}
+      });
+    },
+    message(socket: any, message: string | Buffer) {
+      const upstreamSocket = socket.data.upstreamSocket;
+      if (socket.data.upstreamOpen && upstreamSocket?.readyState === WebSocket.OPEN) upstreamSocket.send(message);
+      else socket.data.pendingMessages.push(message);
+    },
+    close(socket: any) {
+      const upstreamSocket = socket.data.upstreamSocket;
+      if (upstreamSocket?.readyState === WebSocket.OPEN || upstreamSocket?.readyState === WebSocket.CONNECTING) upstreamSocket.close();
+    }
  }
-
-  if (request.method === "GET" && url.pathname === "/help") {
-    sendJson(response, 200, {
-      serviceId,
-      commands: ["GET /health", "GET /edge/health", "GET /health/live", "GET /status", "GET /routes"]
-    });
-    return;
-  }
-
-  if (url.pathname === "/v1/internal" || url.pathname.startsWith("/v1/internal/")) {
-    sendJson(response, 404, {
-      error: {
-        code: "not_found",
-        message: "Internal HWLAB routes are not exposed through the public edge proxy."
-      }
-    });
-    return;
-  }
-
-  proxyHttpRequest({ request, response, upstream, timeoutMs: proxyTimeoutMs });
 });

-listen(server, { serviceId, host, port });
+process.stdout.write(`${JSON.stringify({ serviceId, status: "listening", host, port: server.port })}\n`);
+
+for (const signal of ["SIGINT", "SIGTERM"]) {
+  process.on(signal, () => {
+    server.stop(true);
+    process.exit(0);
+  });
+}
+
+async function proxyFetch(request: Request, upstreamUrl: string, timeoutMs: number) {
+  const target = upstreamHttpUrl(request, upstreamUrl);
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), timeoutMs);
+  try {
+    return await fetch(target, {
+      method: request.method,
+      headers: request.headers,
+      body: request.method === "GET" || request.method === "HEAD" ? undefined : request.body,
+      redirect: "manual",
+      signal: controller.signal
+    });
+  } catch (error) {
+    const timedOut = controller.signal.aborted;
+    const code = timedOut ? "proxy_timeout" : "upstream_unavailable";
+    const userMessage = timedOut
+      ? `Code Agent 代理等待上游超过 ${timeoutMs}ms；输入已保留，可稍后重试。`
+      : "Code Agent 代理暂时无法连接上游；输入已保留，可稍后重试。";
+    return json({ status: "failed", error: { code, layer: "proxy", category: timedOut ? "timeout" : "proxy", retryable: true, userMessage, message: userMessage, blocker: { code, layer: "proxy", retryable: true, summary: userMessage } }, upstream: upstreamUrl, message: userMessage }, 502);
+  } finally {
+    clearTimeout(timeout);
+  }
+}
+
+function upstreamHttpUrl(request: Request, upstreamUrl: string) {
+  const source = new URL(request.url);
+  return new URL(source.pathname + source.search, upstreamUrl).toString();
+}
+
+function upstreamWebSocketUrl(request: Request, upstreamUrl: string) {
+  const target = new URL(upstreamHttpUrl(request, upstreamUrl));
+  target.protocol = target.protocol === "https:" ? "wss:" : "ws:";
+  return target.toString();
+}
+
+function isWebSocketUpgrade(request: Request) {
+  return request.headers.get("upgrade")?.toLowerCase() === "websocket";
+}
+
+function json(body: any, status = 200) {
+  return new Response(`${JSON.stringify(body, null, 2)}\n`, { status, headers: { "content-type": "application/json; charset=utf-8", "cache-control": "no-store" } });
+}
+
+function text(body: string, status = 200) {
+  return new Response(body, { status, headers: { "content-type": "text/plain; charset=utf-8", "cache-control": "no-store" } });
+}
+
+function logWsBridge(event: string, upstreamUrl: string, extra: Record<string, unknown> = {}) {
+  process.stderr.write(`${JSON.stringify({ event: `hwlab-edge-proxy.ws.${event}`, upstreamUrl: redactUrlToken(upstreamUrl), observedAt: new Date().toISOString(), ...extra })}\n`);
+}
+
+function redactUrlToken(urlValue: string) {
+  const parsed = new URL(urlValue);
+  if (parsed.searchParams.has("token")) parsed.searchParams.set("token", "<redacted>");
+  return parsed.toString();
+}
@@ -90,6 +90,7 @@ export function connectHwpodNodeWs(config: any, options: any = {}) {
    if (closed) return;
    socket = new WebSocket(wsUrl);
    socket.addEventListener("open", () => {
+      logWsEvent("open");
      attempt = 0;
      sendJson({
        type: "register",
@@ -103,11 +104,15 @@ export function connectHwpodNodeWs(config: any, options: any = {}) {
    });
    socket.addEventListener("message", (event: MessageEvent) => {
      void handleWsMessage(event.data).catch((error) => {
+        logWsEvent("message_error", { message: error?.message ?? String(error), code: error?.code ?? "hwpod_node_ws_message_failed" });
        sendJson({ type: "hwpod-node-error", nodeId, error: { message: error?.message ?? String(error), code: error?.code ?? "hwpod_node_ws_message_failed" }, observedAt: now() });
      });
    });
-    socket.addEventListener("close", scheduleReconnect);
-    socket.addEventListener("error", () => undefined);
+    socket.addEventListener("close", (event) => {
+      logWsEvent("close", { code: event.code, reason: event.reason || null, reconnect });
+      scheduleReconnect();
+    });
+    socket.addEventListener("error", () => logWsEvent("error", { readyState: socket?.readyState ?? null }));
  }

  function scheduleReconnect() {
@@ -122,6 +127,7 @@ export function connectHwpodNodeWs(config: any, options: any = {}) {
    const textValue = typeof raw === "string" ? raw : await new Response(raw).text();
    const message = JSON.parse(textValue);
    if (message?.type === "ack" && message.requestId === "register" && readyResolve) {
+      logWsEvent("registered", { ok: message.ok === true, message: text(message.message) || null });
      readyResolve(message);
      readyResolve = null;
      return;
@@ -160,10 +166,20 @@ export function connectHwpodNodeWs(config: any, options: any = {}) {
    socket?.close();
  }

+  function logWsEvent(event: string, extra: Record<string, unknown> = {}) {
+    process.stderr.write(`${JSON.stringify({ event: `hwpod-node.ws.${event}`, nodeId, wsUrl: redactUrlToken(wsUrl), observedAt: now(), ...extra })}\n`);
+  }
+
  connect();
  return { wsUrl, nodeId, ready, close };
 }

+function redactUrlToken(urlValue: string) {
+  const parsed = new URL(urlValue);
+  if (parsed.searchParams.has("token")) parsed.searchParams.set("token", "<redacted>");
+  return parsed.toString();
+}
+
 function hwpodNodeWsUrl(config: any) {
  const explicit = text(config.wsUrl);
  if (explicit) return appendToken(explicit, text(config.token));