WRNexusJS/packages/rpc/src/stream.ts

import type { Context } from "@wrnexus/core";
import { RPC_ERROR_CODES, ServiceError } from "./errors.ts";
import { exportSubjectContext, importSubjectContext, type SubjectContext } from "./identity.ts";
import { resolveAppOrigin, RPC_INTERNAL_HEADER } from "./http.ts";
import type {
  AnyProcedures,
  InferProcedureInput,
  InferProcedureOutput,
  ServiceContract,
} from "./types.ts";

export const RPC_STREAM_PATH_PREFIX = "/__wrnexus/rpc-stream";

export function rpcStreamPath(service: string, procedure: string): string {
  return `${RPC_STREAM_PATH_PREFIX}/${service}/${procedure}`;
}

export interface StreamImplementation<Procedures extends AnyProcedures = AnyProcedures> {
  contract: ServiceContract<Procedures>;
  stream(procedure: string, payload: unknown, identity?: string): AsyncIterable<unknown>;
  streamOptions: Required<Pick<StreamImplementOptions, "maxFrameBytes" | "heartbeatMs">>;
  metrics: StreamMetrics;
}

export type StreamHandlers<Procedures extends AnyProcedures> = {
  [K in keyof Procedures]: (
    input: InferProcedureInput<Procedures[K]>,
    ctx: { subject?: import("./identity.ts").SubjectContext },
  ) => AsyncIterable<InferProcedureOutput<Procedures[K]>>;
};

export interface StreamImplementOptions {
  selfApp: string;
  checkPermission?: (permission: string, subject?: SubjectContext) => Promise<boolean> | boolean;
  /** Maximum serialized SSE data frame. Default: 64 KiB. */
  maxFrameBytes?: number;
  /** Emit an SSE comment while idle. Default: 15 seconds. */
  heartbeatMs?: number;
}

export interface StreamMetricsSnapshot {
  started: number;
  completed: number;
  failed: number;
  active: number;
}

export class StreamMetrics {
  private started = 0;
  private completed = 0;
  private failed = 0;
  private active = 0;
  begin() {
    this.started++;
    this.active++;
  }
  complete() {
    this.completed++;
    this.active = Math.max(0, this.active - 1);
  }
  fail() {
    this.failed++;
    this.active = Math.max(0, this.active - 1);
  }
  snapshot(): StreamMetricsSnapshot {
    return {
      started: this.started,
      completed: this.completed,
      failed: this.failed,
      active: this.active,
    };
  }
}

/** Define an authenticated, validated stream endpoint. */
export function implementStream<Procedures extends AnyProcedures>(
  contract: ServiceContract<Procedures>,
  handlers: StreamHandlers<Procedures>,
  options: StreamImplementOptions,
): StreamImplementation<Procedures> {
  const maxFrameBytes = options.maxFrameBytes ?? 64 * 1024;
  const heartbeatMs = options.heartbeatMs ?? 15_000;
  if (!Number.isInteger(maxFrameBytes) || maxFrameBytes < 1)
    throw new RangeError("rpc stream maxFrameBytes must be a positive integer");
  if (!Number.isInteger(heartbeatMs) || heartbeatMs < 1)
    throw new RangeError("rpc stream heartbeatMs must be a positive integer");
  return {
    contract,
    streamOptions: { maxFrameBytes, heartbeatMs },
    metrics: new StreamMetrics(),
    stream(procedure, payload, identity) {
      const handler = Object.hasOwn(handlers, procedure)
        ? handlers[procedure as keyof Procedures]
        : undefined;
      if (!handler) throw new ServiceError(RPC_ERROR_CODES.unknown, "Unknown procedure");
      return (async function* () {
        const definition = contract.procedures[procedure as keyof Procedures];
        let subject: SubjectContext | undefined;
        if (identity) {
          try {
            subject = await importSubjectContext(identity, options.selfApp);
          } catch {
            throw new ServiceError(RPC_ERROR_CODES.identity, "Invalid identity");
          }
        }
        if (definition.permission) {
          if (
            !options.checkPermission ||
            !(await options.checkPermission(definition.permission, subject))
          ) {
            throw new ServiceError(RPC_ERROR_CODES.denied, "Forbidden");
          }
        }
        let input = payload;
        if (definition.input) {
          let parsed: { ok: boolean; value?: unknown };
          try {
            parsed = definition.input.parse(payload as Record<string, unknown>);
          } catch {
            throw new ServiceError(RPC_ERROR_CODES.invalid, "Invalid input");
          }
          if (!parsed.ok) throw new ServiceError(RPC_ERROR_CODES.invalid, "Invalid input");
          input = parsed.value;
        }
        yield* (
          handler as (
            input: unknown,
            ctx: { subject?: import("./identity.ts").SubjectContext },
          ) => AsyncIterable<unknown>
        )(input, { subject });
      })();
    },
  };
}

export interface StreamClientOptions {
  app?: string;
  as?: Context;
  fetch?: typeof fetch;
  signal?: AbortSignal;
  /** Reject oversized server frames before parsing. Default: 64 KiB. */
  maxFrameBytes?: number;
}
export type StreamClient<Procedures extends AnyProcedures> = {
  [K in keyof Procedures]: (
    input: InferProcedureInput<Procedures[K]>,
  ) => AsyncIterable<InferProcedureOutput<Procedures[K]>>;
};

function decodeFrames(
  body: ReadableStream<Uint8Array>,
  signal?: AbortSignal,
  maxFrameBytes = 64 * 1024,
): AsyncIterable<unknown> {
  return (async function* () {
    const reader = body.getReader();
    const decoder = new TextDecoder();
    let buffer = "";
    try {
      while (!signal?.aborted) {
        const next = await reader.read();
        if (next.done) break;
        buffer += decoder.decode(next.value, { stream: true });
        let boundary: number;
        while ((boundary = buffer.indexOf("\n\n")) >= 0) {
          const frame = buffer.slice(0, boundary);
          buffer = buffer.slice(boundary + 2);
          const data = frame
            .split("\n")
            .find((line) => line.startsWith("data: "))
            ?.slice(6);
          if (!data) continue;
          if (new TextEncoder().encode(data).byteLength > maxFrameBytes) {
            throw new ServiceError(RPC_ERROR_CODES.malformed, "Stream frame exceeds limit");
          }
          let value: unknown;
          try {
            value = JSON.parse(data);
          } catch {
            throw new ServiceError(RPC_ERROR_CODES.malformed, "Malformed stream response");
          }
          if (!value || typeof value !== "object" || !("ok" in value))
            throw new ServiceError(RPC_ERROR_CODES.malformed, "Malformed stream response");
          const result = value as { ok: boolean; value?: unknown; code?: string; message?: string };
          if (!result.ok)
            throw new ServiceError(
              result.code ?? RPC_ERROR_CODES.handler,
              result.message ?? "Stream failed",
            );
          yield result.value;
        }
      }
    } finally {
      reader.cancel().catch(() => {});
    }
  })();
}

export function streamClient<Procedures extends AnyProcedures>(
  contract: ServiceContract<Procedures>,
  options: StreamClientOptions = {},
): StreamClient<Procedures> {
  const app = options.app ?? contract.name;
  const doFetch = options.fetch ?? fetch;
  const maxFrameBytes = options.maxFrameBytes ?? 64 * 1024;
  if (!Number.isInteger(maxFrameBytes) || maxFrameBytes < 1)
    throw new RangeError("rpc stream maxFrameBytes must be a positive integer");
  return new Proxy({} as StreamClient<Procedures>, {
    get(_target, property) {
      if (typeof property !== "string" || !Object.hasOwn(contract.procedures, property))
        return undefined;
      return (input: unknown) =>
        (async function* () {
          let identity: string | undefined;
          if (options.as) identity = await exportSubjectContext(options.as, app);
          const headers: Record<string, string> = {
            "content-type": "application/json",
            [RPC_INTERNAL_HEADER]: "1",
            accept: "text/event-stream",
          };
          if (identity) headers["x-wrnexus-rpc-identity"] = identity;
          const response = await doFetch(
            `${resolveAppOrigin(app)}${rpcStreamPath(contract.name, property)}`,
            { method: "POST", headers, body: JSON.stringify(input ?? {}), signal: options.signal },
          );
          if (!response.ok || !response.body)
            throw new ServiceError(RPC_ERROR_CODES.transport, "Service unavailable");
          yield* decodeFrames(response.body, options.signal, maxFrameBytes);
        })();
    },
  });
}