WRNexusJS/packages/dev-server/src/rpc-dispatch.ts

import {
  RPC_IDENTITY_HEADER,
  RPC_INTERNAL_HEADER,
  RPC_PATH_PREFIX,
  RPC_STREAM_PATH_PREFIX,
  type ServiceImplementation,
  type StreamImplementation,
} from "@wrnexus/rpc";

export { RPC_INTERNAL_HEADER };
const EDGE_HEADERS = ["x-forwarded-for", "x-forwarded-host", "x-forwarded-proto", "forwarded"];

export function isRpcPath(pathname: string): boolean {
  return (
    pathname === RPC_PATH_PREFIX ||
    pathname.startsWith(`${RPC_PATH_PREFIX}/`) ||
    pathname === RPC_STREAM_PATH_PREFIX ||
    pathname.startsWith(`${RPC_STREAM_PATH_PREFIX}/`)
  );
}

export function isInternalCaller(req: Request): boolean {
  return (
    req.headers.get(RPC_INTERNAL_HEADER) === "1" &&
    !EDGE_HEADERS.some((name) => req.headers.has(name))
  );
}

function json(body: unknown, status = 200): Response {
  return Response.json(body, { status, headers: { "cache-control": "private, no-store" } });
}

export async function handleRpcRequest(
  req: Request,
  url: URL,
  services: Map<string, ServiceImplementation | StreamImplementation>,
): Promise<Response | null> {
  if (!isRpcPath(url.pathname)) return null;
  if (!isInternalCaller(req)) return new Response("Not found", { status: 404 });
  if (req.method !== "POST") return new Response("Method not allowed", { status: 405 });
  const streaming =
    url.pathname === RPC_STREAM_PATH_PREFIX ||
    url.pathname.startsWith(`${RPC_STREAM_PATH_PREFIX}/`);
  const segments = url.pathname.split("/");
  const SAFE_SEGMENT = /^[A-Za-z0-9_-]+$/;
  const serviceName = segments[3];
  const procedure = segments[4];
  const service =
    serviceName && SAFE_SEGMENT.test(serviceName) ? services.get(serviceName) : undefined;
  if (!service || !procedure || !SAFE_SEGMENT.test(procedure) || segments.length !== 5) {
    return json({ ok: false, code: "RPC_UNKNOWN", message: "Unknown procedure", retryable: false });
  }
  let payload: unknown;
  try {
    payload = await req.json();
  } catch {
    return json({ ok: false, code: "RPC_INVALID", message: "Invalid input", retryable: false });
  }
  if (streaming) {
    if (!("stream" in service) || typeof service.stream !== "function") {
      return json({
        ok: false,
        code: "RPC_UNKNOWN",
        message: "Unknown procedure",
        retryable: false,
      });
    }
    const encoder = new TextEncoder();
    const identity = req.headers.get(RPC_IDENTITY_HEADER) ?? undefined;
    const iterator = service.stream(procedure, payload, identity)[Symbol.asyncIterator]();
    const { maxFrameBytes, heartbeatMs } = service.streamOptions;
    service.metrics.begin();
    const body = new ReadableStream<Uint8Array>({
      async start(controller) {
        try {
          let next = iterator.next();
          while (true) {
            let heartbeatTimer: ReturnType<typeof setTimeout> | undefined;
            const heartbeat = new Promise<{ kind: "heartbeat" }>((resolve) => {
              heartbeatTimer = setTimeout(() => resolve({ kind: "heartbeat" }), heartbeatMs);
            });
            const outcome = await Promise.race([
              next.then((value) => ({ kind: "data" as const, value })),
              heartbeat,
            ]);
            if (heartbeatTimer !== undefined) clearTimeout(heartbeatTimer);
            if (outcome.kind === "heartbeat") {
              controller.enqueue(encoder.encode(": keepalive\n\n"));
              continue;
            }
            if (outcome.value.done) {
              service.metrics.complete();
              controller.close();
              return;
            }
            const frame = JSON.stringify({ ok: true, value: outcome.value.value });
            if (encoder.encode(frame).byteLength > maxFrameBytes) {
              service.metrics.fail();
              controller.enqueue(
                encoder.encode(
                  'data: {"ok":false,"code":"RPC_MALFORMED","message":"Stream frame exceeds limit"}\n\n',
                ),
              );
              controller.close();
              return;
            }
            controller.enqueue(encoder.encode(`data: ${frame}\n\n`));
            next = iterator.next();
          }
        } catch {
          service.metrics.fail();
          controller.enqueue(
            encoder.encode('data: {"ok":false,"code":"RPC_HANDLER","message":"Stream failed"}\n\n'),
          );
          controller.close();
        }
      },
      async cancel() {
        service.metrics.complete();
        await iterator.return?.();
      },
    });
    return new Response(body, {
      headers: {
        "content-type": "text/event-stream; charset=utf-8",
        "cache-control": "private, no-store",
        "x-accel-buffering": "no",
      },
    });
  }
  if (!("invoke" in service) || typeof service.invoke !== "function") {
    return json({ ok: false, code: "RPC_UNKNOWN", message: "Unknown procedure", retryable: false });
  }
  return json(
    await service.invoke(procedure, payload, req.headers.get(RPC_IDENTITY_HEADER) ?? undefined),
  );
}