@abraca/dabra 1.8.2 → 2.0.0

package/src/RpcClient.ts

@@ -0,0 +1,659 @@
+import EventEmitter from "./EventEmitter.ts";
+
+/**
+ * RPC v1 — request/response over `MSG_STATELESS`.
+ *
+ * Mirrors `docs/rpc-v1.md` in the abracadabra-rs server. Frames travel as
+ * `MSG_STATELESS` payloads with the literal prefix `rpc:v1:` followed by a
+ * JSON envelope. The provider's existing `sendStateless` / `stateless` event
+ * stream is the only transport surface this layer touches; everything else
+ * (state machine, deadlines, handler registry, dedupe of self-replies) lives
+ * here.
+ */
+
+export const RPC_PREFIX = "rpc:v1:";
+
+export type RpcKind =
+  | "req"
+  | "ack"
+  | "nack"
+  | "progress"
+  | "result"
+  | "cancel"
+  | "hb"
+  | "register"
+  | "unregister";
+
+export interface RpcFrame {
+  kind: RpcKind;
+  id: string;
+  ts?: number;
+  from?: string;
+  to?: string;
+  method?: string;
+  args?: unknown;
+  deadline_ms?: number;
+  data?: unknown;
+  error?: RpcErrorPayload;
+  by?: "server" | "runner";
+  runner_id?: string;
+  seq?: number;
+  methods?: string[];
+}
+
+export interface RpcErrorPayload {
+  code: string;
+  message: string;
+  details?: unknown;
+}
+
+export type RpcErrorCode =
+  | "NO_HANDLER"
+  | "HANDLER_GONE"
+  | "TIMEOUT"
+  | "CANCELLED"
+  | "RATE_LIMITED"
+  | "UNAUTHORIZED"
+  | "SCHEMA"
+  | "INTERNAL"
+  | "APP"
+  | string; // forward-compat for unknown codes from runners
+
+/** Thrown by `rpc.call(...)` on terminal nack / error / timeout. */
+export class RpcError extends Error {
+  code: RpcErrorCode;
+  details?: unknown;
+  constructor(code: RpcErrorCode, message: string, details?: unknown) {
+    super(message);
+    this.name = "RpcError";
+    this.code = code;
+    this.details = details;
+  }
+}
+
+/**
+ * Minimal provider surface RpcClient needs. The base provider already
+ * satisfies it; the indirection is here so handlers can be tested against
+ * a stub without booting a Y.Doc.
+ */
+export interface RpcTransport {
+  sendStateless(payload: string): void;
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-function-type
+  on(event: string, fn: Function): unknown;
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-function-type
+  off(event: string, fn?: Function): unknown;
+}
+
+export type RpcTarget =
+  | { kind: "role"; role: "service" }
+  | { kind: "user"; userId: string }
+  | { kind: "runner"; sessionId: string };
+
+export interface RpcCallOptions {
+  /** Defaults to `{ kind: "role", role: "service" }`. */
+  target?: RpcTarget;
+  /** Wall-clock deadline in ms. Server clamps to its own ceiling. */
+  deadline?: number;
+  /** Abort the call (sends `rpc:cancel` to the runner). */
+  signal?: AbortSignal;
+  /** Fired when a runner claims the call (after the server ack). */
+  onClaim?: (runnerId: string) => void;
+  /** Fired for every `rpc:progress` frame the runner emits. */
+  onProgress?: (data: unknown, seq: number) => void;
+}
+
+export interface RpcCallHandle<T = unknown> extends Promise<T> {
+  /** RPC id (UUID v7) — useful for log correlation. */
+  readonly id: string;
+  /** Runner session id, populated once a runner has claimed the call. */
+  readonly claimedBy: string | undefined;
+  /** Cancel the call. Equivalent to `signal.abort()` if a signal was passed. */
+  cancel(reason?: string): void;
+}
+
+interface PendingRpc {
+  id: string;
+  resolve: (value: unknown) => void;
+  reject: (err: Error) => void;
+  options: RpcCallOptions;
+  state: "pending" | "accepted" | "claimed" | "settled";
+  claimedBy: string | undefined;
+  deadlineTimer: ReturnType<typeof setTimeout> | undefined;
+}
+
+/**
+ * Handler signature for a runner-side RPC handler.
+ *
+ * Returning a value emits `rpc:result(ok)`. Throwing emits
+ * `rpc:result(err)` — `RpcError` instances pass through their `code` and
+ * `details`; any other thrown value becomes `INTERNAL`.
+ *
+ * Returning an async generator turns each `yield` into a `rpc:progress`
+ * frame (with monotonic `seq`); the generator's return value is the final
+ * `result.data`.
+ */
+export type RpcHandler = (
+  args: unknown,
+  ctx: RpcHandlerContext,
+) =>
+  | unknown
+  | Promise<unknown>
+  | AsyncGenerator<unknown, unknown, void>;
+
+export interface RpcHandlerContext {
+  id: string;
+  method: string;
+  /** Server-stamped caller user id. Trusted. */
+  from: string;
+  /** Aborted when the caller cancels or the deadline is reached. */
+  signal: AbortSignal;
+}
+
+// ── UUID v7 (small inline impl) ─────────────────────────────────────────────
+
+function uuidv7(): string {
+  // Minimal UUID v7 — time-ordered, 122 random bits.
+  const ts = Date.now();
+  const tsHex = ts.toString(16).padStart(12, "0");
+  const rand = new Uint8Array(10);
+  if (typeof crypto !== "undefined" && crypto.getRandomValues) {
+    crypto.getRandomValues(rand);
+  } else {
+    for (let i = 0; i < rand.length; i++) rand[i] = Math.floor(Math.random() * 256);
+  }
+  // version=7, variant=10
+  rand[0] = (rand[0] & 0x0f) | 0x70;
+  rand[2] = (rand[2] & 0x3f) | 0x80;
+  const hex = Array.from(rand, (b) => b.toString(16).padStart(2, "0")).join("");
+  return (
+    `${tsHex.slice(0, 8)}-${tsHex.slice(8, 12)}-` +
+    `${hex.slice(0, 4)}-${hex.slice(4, 8)}-${hex.slice(8, 20)}`
+  );
+}
+
+// ── Client ──────────────────────────────────────────────────────────────────
+
+/**
+ * Per-provider RPC layer. Construct one and attach by passing the provider
+ * (or any `RpcTransport`). Disposes on `destroy()`.
+ *
+ * Caller side: `rpc.call('ai.summarize@1', { docId }, { onProgress })`.
+ *
+ * Runner side: `rpc.handle('ai.summarize@1', async (args, ctx) => …)`.
+ */
+export class RpcClient extends EventEmitter {
+  private readonly transport: RpcTransport;
+  private readonly pending = new Map<string, PendingRpc>();
+  private readonly handlers = new Map<string, RpcHandler>();
+  /** Tracks live runner-side invocations so we can route `rpc:cancel`. */
+  private readonly runningHandlers = new Map<string, AbortController>();
+  /** Pending register frames awaiting server ack — resolved by id. */
+  private readonly pendingRegistrations = new Map<string, () => void>();
+  /** Grace timer armed on `disconnect`; cancelled on `connect`. Fires
+   * `HANDLER_GONE` on every pending RPC if the WS is still down after
+   * the grace window. v1 used to wait `deadline + 5s` which could be
+   * 35s on default deadlines — way too long for a definitively-dead
+   * connection. */
+  private disconnectGrace: ReturnType<typeof setTimeout> | undefined;
+  private static readonly DISCONNECT_GRACE_MS = 10_000;
+  private readonly onStatelessBound: (data: { payload: string }) => void;
+  /**
+   * Replays all live `register` frames whenever the underlying provider
+   * (re)syncs — covers WebSocket reconnects where the server allocated a
+   * new session id and the old capability entries were freed.
+   */
+  private readonly onSyncedBound: (data: { state: boolean }) => void;
+  private readonly onDisconnectBound: () => void;
+  private readonly onConnectBound: () => void;
+  private destroyed = false;
+
+  constructor(transport: RpcTransport) {
+    super();
+    this.transport = transport;
+    this.onStatelessBound = (data) => this.receive(data.payload);
+    this.onSyncedBound = (data) => {
+      if (data.state) this.replayRegistrations();
+    };
+    this.onDisconnectBound = () => this.armDisconnectGrace();
+    this.onConnectBound = () => this.cancelDisconnectGrace();
+    transport.on("stateless", this.onStatelessBound);
+    transport.on("synced", this.onSyncedBound);
+    transport.on("disconnect", this.onDisconnectBound);
+    transport.on("connect", this.onConnectBound);
+  }
+
+  private armDisconnectGrace(): void {
+    if (this.destroyed) return;
+    if (this.disconnectGrace) return; // already armed
+    this.disconnectGrace = setTimeout(
+      () => this.failPendingOnDisconnect(),
+      RpcClient.DISCONNECT_GRACE_MS,
+    );
+  }
+
+  private cancelDisconnectGrace(): void {
+    if (this.disconnectGrace) {
+      clearTimeout(this.disconnectGrace);
+      this.disconnectGrace = undefined;
+    }
+  }
+
+  /**
+   * Called when the WS has been down for the full grace window. Every
+   * in-flight call is dead from the server's perspective (its `close_session`
+   * already fired `HANDLER_GONE` to a now-disconnected caller; on reconnect
+   * the server-side session is fresh, so any cached cancel/result frames
+   * the server tried to deliver were lost). Fail pending immediately.
+   */
+  private failPendingOnDisconnect(): void {
+    this.disconnectGrace = undefined;
+    if (this.destroyed) return;
+    for (const [, p] of this.pending) {
+      if (p.deadlineTimer) clearTimeout(p.deadlineTimer);
+      this.settle(p, () =>
+        p.reject(new RpcError("HANDLER_GONE", "websocket disconnected")),
+      );
+    }
+  }
+
+  /**
+   * Re-emit a `register` frame for every locally-known handler. Called on
+   * every `synced` event from the provider, since a reconnected WS gets a
+   * fresh server-side session with empty capability state.
+   */
+  private replayRegistrations(): void {
+    if (this.destroyed) return;
+    // A fresh server-side session will never ack any register frame issued
+    // before this `synced` event, so prior pending entries are dead. Resolve
+    // them with the current map clear so any awaiting `ready()` doesn't
+    // hang forever — the new register below is the source of truth now.
+    if (this.pendingRegistrations.size > 0) {
+      const stale = [...this.pendingRegistrations.values()];
+      this.pendingRegistrations.clear();
+      for (const r of stale) r();
+    }
+    if (this.handlers.size === 0) return;
+    const methods = [...this.handlers.keys()];
+    const id = uuidv7();
+    // Track this re-register so `ready()` resolves once the server ack lands.
+    this.pendingRegistrations.set(id, () => {});
+    this.send({ kind: "register", id, methods });
+  }
+
+  destroy(): void {
+    if (this.destroyed) return;
+    this.destroyed = true;
+    this.transport.off("stateless", this.onStatelessBound);
+    this.transport.off("synced", this.onSyncedBound);
+    this.transport.off("disconnect", this.onDisconnectBound);
+    this.transport.off("connect", this.onConnectBound);
+    this.cancelDisconnectGrace();
+    // Abort every in-flight handler.
+    for (const [, ac] of this.runningHandlers) {
+      try { ac.abort(); } catch { /* noop */ }
+    }
+    this.runningHandlers.clear();
+    // Reject every pending caller.
+    for (const [, p] of this.pending) {
+      if (p.deadlineTimer) clearTimeout(p.deadlineTimer);
+      p.reject(new RpcError("CANCELLED", "RpcClient destroyed"));
+    }
+    this.pending.clear();
+    // Unblock any awaiters of `ready()` — the WS is gone and no register
+    // ack will ever land. Resolving (vs rejecting) keeps `ready(): Promise<void>`
+    // semantically clean; the call sites that follow `ready()` are typically
+    // diagnostic/orchestration, not request hot paths.
+    for (const [, resolve] of this.pendingRegistrations) {
+      try { resolve(); } catch { /* noop */ }
+    }
+    this.pendingRegistrations.clear();
+    this.handlers.clear();
+    this.removeAllListeners();
+  }
+
+  // ── Caller side ───────────────────────────────────────────────────────────
+
+  /**
+   * Send an RPC and return a handle that resolves with the runner's result.
+   * Throws `RpcError` on nack / handler error / timeout / cancellation.
+   */
+  call<T = unknown>(method: string, args?: unknown, options?: RpcCallOptions): RpcCallHandle<T> {
+    const opts = options ?? {};
+    const id = uuidv7();
+    const target = opts.target ?? { kind: "role", role: "service" };
+
+    let resolve!: (value: unknown) => void;
+    let reject!: (err: Error) => void;
+    const promise = new Promise<unknown>((res, rej) => {
+      resolve = res;
+      reject = rej;
+    });
+
+    // Calling on a destroyed client used to silently hang because `send()`
+    // short-circuits but the pending entry was already created. Reject now
+    // and skip the rest so consumers always get a terminal outcome.
+    if (this.destroyed) {
+      const handle = promise as RpcCallHandle<T>;
+      Object.defineProperty(handle, "id", { value: id, enumerable: true });
+      Object.defineProperty(handle, "claimedBy", { enumerable: true, get: () => undefined });
+      handle.cancel = () => {};
+      reject(new RpcError("CANCELLED", "rpc client destroyed"));
+      return handle;
+    }
+
+    // Pre-aborted signal: no point sending the `req` (the runner would do
+    // wasted work for a caller that's already given up) and no point sending
+    // a `cancel` for an id the server has never seen. Short-circuit cleanly.
+    if (opts.signal?.aborted) {
+      const reason = opts.signal.reason;
+      const handle = promise as RpcCallHandle<T>;
+      Object.defineProperty(handle, "id", { value: id, enumerable: true });
+      Object.defineProperty(handle, "claimedBy", { enumerable: true, get: () => undefined });
+      handle.cancel = () => {};
+      reject(new RpcError("CANCELLED", typeof reason === "string" ? reason : "aborted"));
+      return handle;
+    }
+
+    const pending: PendingRpc = {
+      id,
+      resolve,
+      reject,
+      options: opts,
+      state: "pending",
+      claimedBy: undefined,
+      deadlineTimer: undefined,
+    };
+    this.pending.set(id, pending);
+
+    // Local fallback timeout so we don't wait past `deadline` if the server
+    // somehow goes silent (it shouldn't — server enforces its own).
+    const deadline = opts.deadline ?? 30_000;
+    pending.deadlineTimer = setTimeout(() => {
+      if (pending.state !== "settled") {
+        this.settle(pending, () =>
+          reject(new RpcError("TIMEOUT", "rpc deadline exceeded (client-side)")),
+        );
+      }
+    }, deadline + 5_000); // grace window so server-side TIMEOUT lands first
+
+    // Pre-aborted signals were handled above (no req frame sent). For a live
+    // signal, listen for abort and cancel mid-flight when it trips.
+    if (opts.signal) {
+      opts.signal.addEventListener(
+        "abort",
+        () => this.abortInFlight(pending, opts.signal?.reason),
+        { once: true },
+      );
+    }
+
+    const frame: RpcFrame = {
+      kind: "req",
+      id,
+      ts: Date.now(),
+      method,
+      args,
+      to: serializeTarget(target),
+      deadline_ms: deadline,
+    };
+    this.send(frame);
+
+    const handle = promise as RpcCallHandle<T>;
+    Object.defineProperty(handle, "id", { value: id, enumerable: true });
+    Object.defineProperty(handle, "claimedBy", {
+      enumerable: true,
+      get: () => pending.claimedBy,
+    });
+    handle.cancel = (reason?: string) => this.abortInFlight(pending, reason);
+    return handle;
+  }
+
+  private abortInFlight(pending: PendingRpc, reason?: unknown): void {
+    if (pending.state === "settled") return;
+    // Tell the runner we're cancelling.
+    this.send({ kind: "cancel", id: pending.id, ts: Date.now(), data: reason ? { reason: String(reason) } : undefined });
+    // Reject locally; the server may still emit a CANCELLED result later
+    // (we'll see it but the pending is gone, so it's a no-op).
+    this.settle(pending, () =>
+      pending.reject(new RpcError("CANCELLED", typeof reason === "string" ? reason : "cancelled")),
+    );
+  }
+
+  // ── Runner side ───────────────────────────────────────────────────────────
+
+  /**
+   * Register a handler for `method` and advertise it to the server so
+   * incoming `to: role:service` calls resolve here. Returns an unsubscribe
+   * function. Calling `register` twice for the same method replaces the
+   * handler and re-advertises.
+   *
+   * The advertise is fire-and-forget; if you need to wait until the server
+   * has acknowledged registration (typical in tests with a separate caller
+   * connection), `await rpc.ready()` after registering all handlers.
+   */
+  handle(method: string, handler: RpcHandler): () => void {
+    this.handlers.set(method, handler);
+    const id = uuidv7();
+    this.pendingRegistrations.set(id, () => {
+      // Default no-op; overwritten when ready() is awaited.
+    });
+    this.send({ kind: "register", id, methods: [method] });
+    return () => {
+      if (this.handlers.get(method) === handler) {
+        this.handlers.delete(method);
+        this.send({ kind: "unregister", id: uuidv7(), methods: [method] });
+      }
+    };
+  }
+
+  /**
+   * Resolve once every outstanding `register` frame has been acknowledged
+   * by the server. Useful when a separate connection is about to issue
+   * `rpc.call(...)` and you need the capability registry to be live first.
+   * Resolves immediately if there's nothing pending.
+   */
+  ready(): Promise<void> {
+    if (this.pendingRegistrations.size === 0) return Promise.resolve();
+    const ids = [...this.pendingRegistrations.keys()];
+    return Promise.all(
+      ids.map(
+        (id) =>
+          new Promise<void>((resolve) => {
+            this.pendingRegistrations.set(id, resolve);
+          }),
+      ),
+    ).then(() => undefined);
+  }
+
+  // ── Wire ───────────────────────────────────────────────────────────────────
+
+  private send(frame: Partial<RpcFrame> & Pick<RpcFrame, "kind" | "id">): void {
+    if (this.destroyed) return;
+    const payload = `${RPC_PREFIX}${JSON.stringify(frame)}`;
+    this.transport.sendStateless(payload);
+  }
+
+  private receive(payload: string): void {
+    if (!payload.startsWith(RPC_PREFIX)) return;
+    const json = payload.slice(RPC_PREFIX.length);
+    let frame: RpcFrame;
+    try {
+      frame = JSON.parse(json) as RpcFrame;
+    } catch {
+      return;
+    }
+    switch (frame.kind) {
+      case "ack":     return this.onAck(frame);
+      case "progress": return this.onProgress(frame);
+      case "result":  return this.onResult(frame);
+      case "nack":    return this.onNack(frame);
+      case "req":     return void this.onReq(frame);
+      case "cancel":  return this.onCancel(frame);
+      // hb / register / unregister never travel client-bound.
+      default:        return;
+    }
+  }
+
+  private onAck(frame: RpcFrame): void {
+    // Register-ack: server confirmed a `register` frame.
+    const pendingReg = this.pendingRegistrations.get(frame.id);
+    if (pendingReg) {
+      this.pendingRegistrations.delete(frame.id);
+      pendingReg();
+      return;
+    }
+    const pending = this.pending.get(frame.id);
+    if (!pending) return;
+    if (frame.by === "server" && pending.state === "pending") {
+      pending.state = "accepted";
+    } else if (frame.by === "runner") {
+      pending.state = "claimed";
+      pending.claimedBy = frame.runner_id;
+      pending.options.onClaim?.(frame.runner_id ?? "");
+    }
+  }
+
+  private onProgress(frame: RpcFrame): void {
+    const pending = this.pending.get(frame.id);
+    if (!pending) return;
+    pending.options.onProgress?.(frame.data, frame.seq ?? 0);
+  }
+
+  private onResult(frame: RpcFrame): void {
+    const pending = this.pending.get(frame.id);
+    if (!pending) return;
+    if (frame.error) {
+      const err = new RpcError(frame.error.code, frame.error.message, frame.error.details);
+      this.settle(pending, () => pending.reject(err));
+    } else {
+      this.settle(pending, () => pending.resolve(frame.data));
+    }
+  }
+
+  private onNack(frame: RpcFrame): void {
+    const pending = this.pending.get(frame.id);
+    if (!pending) return;
+    const err = frame.error
+      ? new RpcError(frame.error.code, frame.error.message, frame.error.details)
+      : new RpcError("INTERNAL", "rpc nack");
+    this.settle(pending, () => pending.reject(err));
+  }
+
+  private settle(pending: PendingRpc, finalize: () => void): void {
+    if (pending.state === "settled") return;
+    pending.state = "settled";
+    if (pending.deadlineTimer) clearTimeout(pending.deadlineTimer);
+    this.pending.delete(pending.id);
+    finalize();
+  }
+
+  // ── Inbound request (runner side) ─────────────────────────────────────────
+
+  private async onReq(frame: RpcFrame): Promise<void> {
+    const method = frame.method ?? "";
+    const handler = this.handlers.get(method);
+    if (!handler) {
+      // Server should have nacked NO_HANDLER before we got here; ignore.
+      return;
+    }
+    if (this.runningHandlers.has(frame.id)) {
+      // Duplicate req — server retry path. Server will re-emit the cached
+      // terminal frame on its own; we just don't double-invoke.
+      return;
+    }
+    const ac = new AbortController();
+    this.runningHandlers.set(frame.id, ac);
+
+    // Send the runner-claim ack immediately.
+    this.send({
+      kind: "ack",
+      id: frame.id,
+      ts: Date.now(),
+      by: "runner",
+    });
+
+    const ctx: RpcHandlerContext = {
+      id: frame.id,
+      method,
+      from: frame.from ?? "",
+      signal: ac.signal,
+    };
+
+    // Heartbeat while the handler runs. The server doesn't enforce
+    // inactivity in v1 but we send these so tooling can observe liveness
+    // and a v2 server can fail stalled calls.
+    const hbInterval = setInterval(() => {
+      if (!ac.signal.aborted) {
+        this.send({ kind: "hb", id: frame.id, ts: Date.now() });
+      }
+    }, 5_000);
+
+    try {
+      const ret = handler(frame.args, ctx);
+      if (ret && typeof (ret as AsyncGenerator).next === "function" && typeof (ret as AsyncGenerator)[Symbol.asyncIterator] === "function") {
+        // Generator: stream progress, return final.
+        const gen = ret as AsyncGenerator<unknown, unknown, void>;
+        let seq = 0;
+        let final: unknown = undefined;
+        while (true) {
+          const { value, done } = await gen.next();
+          if (done) {
+            final = value;
+            break;
+          }
+          if (ac.signal.aborted) break;
+          this.send({ kind: "progress", id: frame.id, ts: Date.now(), seq: seq++, data: value });
+        }
+        if (ac.signal.aborted) {
+          this.send({
+            kind: "result",
+            id: frame.id,
+            ts: Date.now(),
+            error: { code: "CANCELLED", message: "handler cancelled" },
+          });
+        } else {
+          this.send({ kind: "result", id: frame.id, ts: Date.now(), data: final ?? null });
+        }
+      } else {
+        const value = await (ret as Promise<unknown>);
+        if (ac.signal.aborted) {
+          this.send({
+            kind: "result",
+            id: frame.id,
+            ts: Date.now(),
+            error: { code: "CANCELLED", message: "handler cancelled" },
+          });
+        } else {
+          this.send({ kind: "result", id: frame.id, ts: Date.now(), data: value ?? null });
+        }
+      }
+    } catch (err) {
+      const error =
+        err instanceof RpcError
+          ? { code: err.code, message: err.message, details: err.details }
+          : { code: "INTERNAL", message: err instanceof Error ? err.message : String(err) };
+      this.send({ kind: "result", id: frame.id, ts: Date.now(), error });
+    } finally {
+      clearInterval(hbInterval);
+      this.runningHandlers.delete(frame.id);
+    }
+  }
+
+  private onCancel(frame: RpcFrame): void {
+    // The runner only ever observes `cancel` for an id whose `req` it has
+    // already seen — frame ordering on a single WS preserves the
+    // (req, cancel) sequence the server emits. If `runningHandlers` has
+    // no entry for this id, the handler has already completed (cancel
+    // arrived after result), so the cancel is informational and can be
+    // dropped silently.
+    const ac = this.runningHandlers.get(frame.id);
+    if (ac) ac.abort(frame.error?.message ?? "cancelled by caller");
+  }
+}
+
+function serializeTarget(t: RpcTarget): string {
+  switch (t.kind) {
+    case "role":   return `role:${t.role}`;
+    case "user":   return `user:${t.userId}`;
+    case "runner": return `runner:${t.sessionId}`;
+  }
+}