suaveplan/packages/api-stack/api/api-plugin-sse/src/middleware/sse-handler.ts

import type { Context, Handler } from "hono";
import { SSE_CODES, SseError } from "../error-codes/error-codes.js";
import { formatHeartbeatFrame, formatSseFrame } from "../formatter/format.js";
import { getSseInstruments } from "../otel/instruments.js";
import type { SseService } from "../services/service.js";
import type { OverflowPolicy } from "../types/types.js";

/**
 * Apply the configured overflow policy to the supplied backlog when
 * it has just exceeded the high-water mark. Returns `true` when the
 * incoming (newest) frame was discarded so the caller can short-circuit
 * the post-write bookkeeping. Exported so unit tests can pin every
 * branch without orchestrating a Hono request lifecycle.
 */
export function applyOverflowPolicy(
    backlog: Uint8Array[],
    queueHighWaterMark: number,
    policy: OverflowPolicy
): boolean {
    if (backlog.length <= queueHighWaterMark) return false;
    const instruments = getSseInstruments();
    if (policy === "drop-newest") {
        backlog.pop();
        instruments.queueOverflowTotal.add(1, {
            policy: "drop-newest",
            code: SSE_CODES.QUEUE_OVERFLOW,
        });
        return true;
    }
    backlog.shift();
    instruments.queueOverflowTotal.add(1, {
        policy: "drop-oldest",
        code: SSE_CODES.QUEUE_OVERFLOW,
    });
    return false;
}

/**
 * Options for `createSseRoute`. The handler reads `:channel` from the
 * request path; `resolvePrincipalId` lets consumers wire any auth
 * stack (e.g. `@suaveplan/api-plugin-auth`'s `c.get("user")?.id`)
 * without this package depending on it.
 */
export interface CreateSseRouteOptions {
    readonly heartbeatIntervalMs: number;
    readonly queueHighWaterMark: number;
    readonly overflowPolicy: OverflowPolicy;
    readonly resolvePrincipalId?: (c: Context) => string | undefined;
    /**
     * When `true`, the handler raises `SseError(SSE_UNAUTHENTICATED)`
     * if `resolvePrincipalId` returns `undefined`. The default
     * (`false`) allows anonymous connections so the route may be
     * composed with any external auth middleware.
     */
    readonly requirePrincipal?: boolean;
    readonly setIntervalImpl?: (fn: () => void, ms: number) => unknown;
    readonly clearIntervalImpl?: (handle: unknown) => void;
}

/**
 * Hono handler factory wrapping the SSE service. Returns a Response
 * with a `ReadableStream<Uint8Array>` body so any standards-compliant
 * runtime (Bun, Node 18+, Cloudflare Workers, Deno) streams without
 * an adapter shim.
 *
 * Behavior:
 *  - `:channel` path parameter is required; missing throws
 *    `SseError(SSE_CHANNEL_REQUIRED)`.
 *  - `Last-Event-ID` request header (browser-replay convention) is
 *    forwarded to `service.subscribe` for ring-buffer replay.
 *  - The request's `AbortSignal` is threaded through to the service
 *    so client disconnect unsubscribes deterministically.
 *  - Every `heartbeatIntervalMs` writes a `: heartbeat\n\n` comment
 *    frame; absent traffic, this keeps proxies + load balancers from
 *    silently closing the connection.
 *  - Backpressure: when `controller.desiredSize === 0` and the
 *    in-memory backlog exceeds `queueHighWaterMark`, applies
 *    `overflowPolicy` and increments `sse.queue_overflow_total`.
 */
export function createSseRoute(
    service: SseService,
    options: CreateSseRouteOptions
): Handler {
    const setIntervalFn =
        options.setIntervalImpl ?? ((fn, ms) => setInterval(fn, ms) as unknown);
    const clearIntervalFn =
        options.clearIntervalImpl ??
        ((h) => clearInterval(h as ReturnType<typeof setInterval>));

    return async function sseHandler(c) {
        const channel = c.req.param("channel");
        if (channel === undefined || channel === "") {
            throw new SseError(
                SSE_CODES.CHANNEL_REQUIRED,
                "SSE handler requires a :channel path parameter"
            );
        }
        const lastEventId =
            c.req.header("last-event-id") ?? c.req.header("Last-Event-ID");
        const principalId = options.resolvePrincipalId?.(c);
        if (options.requirePrincipal === true && principalId === undefined) {
            throw new SseError(
                SSE_CODES.UNAUTHENTICATED,
                "SSE handler requires an authenticated principal",
                { context: { statusCode: 401 } }
            );
        }
        const handle = service.subscribe({
            channel,
            ...(lastEventId !== undefined ? { lastEventId } : {}),
            ...(principalId !== undefined ? { principalId } : {}),
            signal: c.req.raw.signal,
        });

        const encoder = new TextEncoder();
        const instruments = getSseInstruments();
        const started = performance.now();
        let heartbeatHandle: unknown;
        let controllerRef: ReadableStreamDefaultController<Uint8Array> | null =
            null;
        const backlog: Uint8Array[] = [];

        const writeFrame = (frame: string): void => {
            if (controllerRef === null) return;
            const bytes = encoder.encode(frame);
            backlog.push(bytes);
            const evictedNewest = applyOverflowPolicy(
                backlog,
                options.queueHighWaterMark,
                options.overflowPolicy
            );
            drainInto(controllerRef);
            if (!evictedNewest) {
                instruments.messageBytesTotal.add(bytes.byteLength);
            }
        };

        // Drain queued bytes into the supplied controller while it
        // signals capacity. Called from both writeFrame (which has
        // already verified controllerRef is non-null) and the
        // ReadableStream pull callback (which fires only while the
        // controller is open).
        const drainInto = (
            ctrl: ReadableStreamDefaultController<Uint8Array>
        ): void => {
            while (
                backlog.length > 0 &&
                (ctrl.desiredSize === null || ctrl.desiredSize > 0)
            ) {
                const next = backlog.shift() as Uint8Array;
                ctrl.enqueue(next);
            }
        };

        const cleanup = async (): Promise<void> => {
            controllerRef = null;
            backlog.length = 0;
            if (heartbeatHandle !== undefined) {
                clearIntervalFn(heartbeatHandle);
                heartbeatHandle = undefined;
            }
            await handle.close();
            const duration = performance.now() - started;
            instruments.connectionDurationMs.record(duration, { channel });
        };

        const body = new ReadableStream<Uint8Array>({
            start: (controller) => {
                controllerRef = controller;
                heartbeatHandle = setIntervalFn(() => {
                    writeFrame(formatHeartbeatFrame());
                    instruments.heartbeatsTotal.add(1);
                }, options.heartbeatIntervalMs);
                void (async () => {
                    try {
                        for await (const event of handle.events) {
                            writeFrame(formatSseFrame(event));
                        }
                    } finally {
                        try {
                            controller.close();
                        } catch {
                            /* already closed */
                        }
                        await cleanup();
                    }
                })();
            },
            pull: (controller) => {
                drainInto(controller);
            },
            cancel: async () => {
                await cleanup();
            },
        });

        return c.body(body, 200, {
            "Content-Type": "text/event-stream; charset=utf-8",
            "Cache-Control": "no-cache, no-transform",
            "X-Accel-Buffering": "no",
            Connection: "keep-alive",
        });
    };
}