suaveplan/packages/api-stack/api/api-worker-rpc/src/client/client.ts

import {
    AbortError,
    attachAbortSignal,
} from "../abort-forwarding/abort-forwarding.js";
import {
    type BackpressurePolicy,
    BoundedQueue,
    CallBudget,
    DEFAULT_BACKPRESSURE,
} from "../backpressure/backpressure.js";
import {
    type CorrelationIdGenerator,
    createCorrelationIdGenerator,
    PendingRegistry,
} from "../correlation/correlation.js";
import { WorkerRpcError } from "../error-codes/error-codes.js";
import { envelopeToError } from "../error-reconstruction/error-reconstruction.js";
import {
    PROTOCOL_VERSION,
    parseFrame,
    type RequestFrame,
    type WorkerRpcFrame,
} from "../protocol/protocol.js";
import type { PortLike } from "../testing/testing.js";
import { assertCloneable } from "./cloneable-guard.js";

export interface WorkerRpcClientOptions {
    readonly idGenerator?: CorrelationIdGenerator;
    readonly callBudget?: number;
    readonly outboundQueue?: BackpressurePolicy;
    /** Per-call timeout in ms; `0` disables. Default `0`. */
    readonly defaultTimeoutMs?: number;
    /**
     * Hook invoked when a stream-chunk arrives for an open stream. Wired up
     * by `streamRequest`; here so the client can multiplex without circular
     * imports.
     */
    readonly onStreamFrame?: (frame: WorkerRpcFrame) => void;
    /** Observer invoked when a non-fatal protocol error is detected. */
    readonly onProtocolError?: (err: WorkerRpcError) => void;
}

export interface WorkerRpcClient {
    /** Invoke a method on the remote side. */
    call<R = unknown>(
        method: string,
        params: ReadonlyArray<unknown>,
        opts?: { signal?: AbortSignal; timeoutMs?: number }
    ): Promise<R>;
    /** Send a `cancel` frame for an active stream/request id. */
    cancel(id: number): void;
    /** Allocate the next correlation id (used by streaming layer). */
    nextId(): number;
    /** Send an already-built frame; used by streaming layer. */
    send(frame: WorkerRpcFrame): void;
    /** Underlying registry — exposed for streaming + tests. */
    readonly registry: PendingRegistry;
    /** Active in-flight call count. */
    readonly inFlight: () => number;
    /** Tear down: rejects pending, removes listener. */
    close(reason?: string): void;
}

/**
 * Build a worker-RPC client over a {@link PortLike}. The returned object
 * is callable via `.call(method, params)` or — through the
 * {@link createWorkerRpcClient}-derived proxy — directly as
 * `client.method(...args)`.
 */
export function createRawWorkerRpcClient(
    port: PortLike,
    opts: WorkerRpcClientOptions = {}
): WorkerRpcClient {
    const idGen = opts.idGenerator ?? createCorrelationIdGenerator();
    const registry = new PendingRegistry();
    const budget = new CallBudget(opts.callBudget ?? 1024);
    const queue = new BoundedQueue<WorkerRpcFrame>(
        opts.outboundQueue ?? DEFAULT_BACKPRESSURE
    );
    let closed = false;

    const listener = (event: { data: unknown }): void => {
        let frame: WorkerRpcFrame;
        try {
            frame = parseFrame(event.data);
        } catch (cause) {
            opts.onProtocolError?.(
                new WorkerRpcError(
                    "WORKER_RPC.INVALID_FRAME",
                    "received malformed RPC frame",
                    { cause }
                )
            );
            return;
        }
        if (frame.v !== PROTOCOL_VERSION) {
            opts.onProtocolError?.(
                new WorkerRpcError(
                    "WORKER_RPC.VERSION_MISMATCH",
                    `expected protocol ${PROTOCOL_VERSION}, got ${frame.v}`
                )
            );
            return;
        }
        if (frame.kind === "response") {
            if (!registry.has(frame.id)) {
                opts.onProtocolError?.(
                    new WorkerRpcError(
                        "WORKER_RPC.ORPHAN_RESPONSE",
                        `no pending call for response id ${frame.id}`
                    )
                );
                return;
            }
            if (frame.ok) {
                registry.settle(frame.id, { ok: true, value: frame.result });
                budget.release();
            } else {
                registry.settle(frame.id, {
                    ok: false,
                    error: envelopeToError(frame.error),
                });
                budget.release();
            }
            return;
        }
        if (
            frame.kind === "stream-chunk" ||
            frame.kind === "stream-end" ||
            frame.kind === "progress" ||
            frame.kind === "stream-ack"
        ) {
            opts.onStreamFrame?.(frame);
            return;
        }
        // hello / goodbye / cancel arriving inbound are advisory here.
    };
    port.addEventListener("message", listener);
    port.start?.();

    const sendRaw = (frame: WorkerRpcFrame): void => {
        if (closed) {
            throw new WorkerRpcError(
                "WORKER_RPC.PORT_CLOSED",
                "cannot send on closed port"
            );
        }
        queue.enqueue(frame);
        const next = queue.dequeue();
        // dequeue is always non-undefined here because enqueue just
        // accepted (or threw) on overflow; the cast is for the type
        // narrowing only.
        port.postMessage(next as WorkerRpcFrame);
    };

    return {
        registry,
        nextId: () => idGen.next(),
        send: sendRaw,
        inFlight: () => budget.inFlight,
        call<R = unknown>(
            method: string,
            params: ReadonlyArray<unknown>,
            callOpts?: { signal?: AbortSignal; timeoutMs?: number }
        ): Promise<R> {
            assertCloneable(params, "args");
            budget.acquire();
            const id = idGen.next();
            const promise = new Promise<R>((resolve, reject) => {
                registry.register({
                    id,
                    resolve: (v) => resolve(v as R),
                    reject,
                });
                const abortLink = attachAbortSignal(callOpts?.signal, () => {
                    if (registry.has(id)) {
                        sendRaw({
                            v: PROTOCOL_VERSION,
                            kind: "cancel",
                            id,
                        });
                        registry.settle(id, {
                            ok: false,
                            error: new AbortError(`call ${method} aborted`),
                        });
                        budget.release();
                    }
                });
                const timeout =
                    callOpts?.timeoutMs ?? opts.defaultTimeoutMs ?? 0;
                if (timeout > 0) {
                    registry.attachTimeout(id, timeout);
                }
                const frame: RequestFrame = {
                    v: PROTOCOL_VERSION,
                    kind: "request",
                    id,
                    method,
                    params: [...params],
                };
                try {
                    sendRaw(frame);
                } catch (err) {
                    abortLink.dispose();
                    registry.settle(id, {
                        ok: false,
                        error:
                            err instanceof Error ? err : new Error(String(err)),
                    });
                    budget.release();
                }
            });
            return promise;
        },
        cancel(id: number): void {
            if (registry.has(id)) {
                sendRaw({ v: PROTOCOL_VERSION, kind: "cancel", id });
                registry.cancel(id);
                budget.release();
            }
        },
        close(reason = "closed"): void {
            if (closed) return;
            closed = true;
            port.removeEventListener("message", listener);
            registry.clear(
                new WorkerRpcError("WORKER_RPC.PORT_CLOSED", reason)
            );
        },
    };
}

/**
 * Typed proxy over {@link createRawWorkerRpcClient}. Every property access
 * becomes a method invocation; the returned promise resolves with the
 * remote result. Reserved fields (`close`, `__raw__`) expose lifecycle
 * controls.
 */
export type WorkerProxyClient<T extends object> = {
    [K in keyof T]: T[K] extends (...args: infer A) => infer R
        ? (...args: A) => Promise<Awaited<R>>
        : T[K];
} & {
    readonly __raw__: WorkerRpcClient;
    close(reason?: string): void;
};

export function createWorkerRpcClient<T extends object>(
    port: PortLike,
    opts: WorkerRpcClientOptions = {}
): WorkerProxyClient<T> {
    const raw = createRawWorkerRpcClient(port, opts);
    const target: Record<string, unknown> = {
        __raw__: raw,
        close: (reason?: string) => raw.close(reason),
    };
    const proxy = new Proxy(target, {
        get(t, prop) {
            if (typeof prop !== "string") return undefined;
            if (prop in t) return t[prop];
            return (...args: ReadonlyArray<unknown>) => raw.call(prop, args);
        },
    });
    return proxy as WorkerProxyClient<T>;
}