pikasTech-HWLAB/internal/dev-entrypoint/cloud-web-proxy.mjs

import { request as httpRequest } from "node:http";

const STREAMING_HEADER_DEFAULTS = Object.freeze({
  "cache-control": "no-store",
  "x-accel-buffering": "no"
});

const FORWARDED_REQUEST_HEADERS = Object.freeze([
  "authorization",
  "cookie",
  "prefer",
  "traceparent",
  "tracestate",
  "x-hwlab-short-connection",
  "x-hwlab-otel-trace-id",
  "x-trace-id",
  "x-request-id",
  "x-actor-id",
  "x-source-service-id"
]);

export function isCloudWebSseRoute(pathname) {
  const value = String(pathname || "");
  return value === "/v1/agent-observer/events" || /^\/v1\/agent\/chat\/trace\/[^/]+\/stream$/u.test(value);
}

export function isEventStreamResponse(headers = {}) {
  const contentType = headerValue(headers, "content-type");
  return /^text\/event-stream(?:\s*;|$)/iu.test(contentType);
}

export function copyProxyResponseHeaders(headers = {}, { streaming = false } = {}) {
  const result = {};
  for (const [key, value] of Object.entries(headers)) {
    const lowerKey = key.toLowerCase();
    if (value === undefined) continue;
    if (lowerKey === "transfer-encoding") continue;
    if (lowerKey === "content-length") continue;
    result[key] = value;
  }
  if (streaming) {
    for (const [key, value] of Object.entries(STREAMING_HEADER_DEFAULTS)) {
      if (!hasHeader(result, key)) result[key] = value;
    }
  }
  return result;
}

export function upstreamRequestHeaders(request, body = "") {
  const headers = {
    accept: request.headers.accept || "application/json"
  };
  const hasBody = body.length > 0;
  if (request.headers["content-type"] || hasBody) {
    headers["content-type"] = request.headers["content-type"] || "application/json";
  }
  for (const name of FORWARDED_REQUEST_HEADERS) {
    if (request.headers[name] !== undefined) headers[name] = request.headers[name];
  }
  if (hasBody) headers["content-length"] = Buffer.byteLength(body);
  return headers;
}

export function proxyCloudApiRequest({
  target,
  request,
  response,
  body = "",
  timeoutMs,
  forceStream = false,
  extraResponseHeaders = {},
  streamTransform = null
}) {
  return new Promise((resolve, reject) => {
    let timedOut = false;
    let settled = false;
    let timeout = null;
    let streamTransformEnded = false;

    const endStreamTransform = () => {
      if (!streamTransformEnded && streamTransform?.end) {
        streamTransformEnded = true;
        return streamTransform.end();
      }
      return null;
    };

    const settle = (callback, value) => {
      if (settled) return;
      settled = true;
      if (timeout) clearTimeout(timeout);
      callback(value);
    };
    const armTimeout = () => {
      if (timeout) clearTimeout(timeout);
      timeout = setTimeout(() => {
        timedOut = true;
        upstream.destroy(new Error("cloud api proxy timed out after " + timeoutMs + "ms"));
      }, timeoutMs);
    };

    const upstream = httpRequest(
      target,
      {
        method: request.method,
        headers: upstreamRequestHeaders(request, body)
      },
      (upstreamResponse) => {
        const upstreamStatusCode = upstreamResponse.statusCode || 502;
        const streaming = forceStream || isEventStreamResponse(upstreamResponse.headers);
        if (streaming) {
          response.writeHead(upstreamStatusCode, mergeResponseHeaders(copyProxyResponseHeaders(upstreamResponse.headers, { streaming: true }), extraResponseHeaders));
          response.flushHeaders?.();
          upstreamResponse.on("data", (chunk) => {
            armTimeout();
            writeStreamOutput(response, streamTransform?.write ? streamTransform.write(chunk) : chunk);
          });
          upstreamResponse.on("end", () => {
            writeStreamOutput(response, endStreamTransform());
            response.end();
            settle(resolve, { statusCode: upstreamStatusCode, streaming: true, streamTransformStats: streamTransform?.stats });
          });
          upstreamResponse.on("error", (error) => {
            if (response.headersSent) {
              response.destroy(error);
              settle(resolve, { statusCode: upstreamStatusCode, streaming: true, errorCode: error?.code || "upstream_response_error" });
              return;
            }
            settle(reject, error);
          });
          upstreamResponse.on("close", () => {
            writeStreamOutput(response, endStreamTransform());
            if (!response.writableEnded) response.end();
            settle(resolve, { statusCode: upstreamStatusCode, streaming: true, streamTransformStats: streamTransform?.stats });
          });
          response.on("close", () => {
            endStreamTransform();
            upstream.destroy();
            settle(resolve, { statusCode: upstreamStatusCode, streaming: true, errorCode: "client_closed", streamTransformStats: streamTransform?.stats });
          });
          return;
        }

        const chunks = [];
        upstreamResponse.on("data", (chunk) => {
          armTimeout();
          chunks.push(chunk);
        });
        upstreamResponse.on("end", () => {
          const responseBody = Buffer.concat(chunks);
          response.writeHead(upstreamStatusCode, {
            ...mergeResponseHeaders(copyProxyResponseHeaders(upstreamResponse.headers), extraResponseHeaders),
            "content-length": responseBody.length
          });
          response.end(responseBody);
          settle(resolve, { statusCode: upstreamStatusCode, streaming: false, bodyBytes: responseBody.length });
        });
        upstreamResponse.on("error", (error) => settle(reject, error));
      }
    );

    armTimeout();
    upstream.on("error", (error) => {
      if (timedOut) error.timedOut = true;
      if (response.headersSent) {
        response.destroy(error);
        settle(resolve, { statusCode: 0, streaming: true, errorCode: error?.code || "upstream_request_error" });
        return;
      }
      settle(reject, error);
    });
    if (body) upstream.write(body);
    upstream.end();
  });
}

function writeStreamOutput(response, output) {
  if (output === undefined || output === null || output === "") return;
  if (Array.isArray(output)) {
    for (const item of output) writeStreamOutput(response, item);
    return;
  }
  response.write(output);
}

function mergeResponseHeaders(baseHeaders = {}, extraHeaders = {}) {
  const result = { ...baseHeaders };
  for (const [key, value] of Object.entries(extraHeaders)) {
    if (value === undefined || value === null || value === "") continue;
    const existingKey = Object.keys(result).find((item) => item.toLowerCase() === key.toLowerCase());
    if (existingKey && existingKey.toLowerCase() === "set-cookie") {
      const existing = Array.isArray(result[existingKey]) ? result[existingKey] : [result[existingKey]];
      const next = Array.isArray(value) ? value : [value];
      result[existingKey] = [...existing.filter(Boolean), ...next.filter(Boolean)];
      continue;
    }
    result[existingKey || key] = value;
  }
  return result;
}

function headerValue(headers, name) {
  const lowerName = name.toLowerCase();
  for (const [key, value] of Object.entries(headers)) {
    if (key.toLowerCase() !== lowerName) continue;
    if (Array.isArray(value)) return value[0] ?? "";
    return String(value ?? "");
  }
  return "";
}

function hasHeader(headers, name) {
  const lowerName = name.toLowerCase();
  return Object.keys(headers).some((key) => key.toLowerCase() === lowerName);
}