@omg-dev/server 0.4.24

package/src/subscriptions.ts

@@ -0,0 +1,654 @@
+// ── Subscriptions ────────────────────────────────────────────────────────────
+//
+// Per-client, per-query reactive state. Where broker.ts broadcasts a single
+// invalidate signal to every SSE viewer regardless of what they care about,
+// this module tracks {client → set of subscriptions} and on a write only
+// re-pushes data to subscribers of the affected collection.
+//
+// Phase 4 scope: sequence numbers + reconnect resume. Every snapshot and
+// every delta carries a monotonic global `seq`; subscribers keep their
+// last-seen seq and send it as `resumeFromSeq` when reconnecting. Each
+// collection holds a small ring buffer of recent (oldRowRaw, newRowRaw)
+// write events — on resume, the server replays only the events whose seq
+// is > resumeFromSeq, re-evaluated through the resuming sub's predicate.
+// If the gap exceeds the ring (~256 entries by default), the server
+// falls back to a full snapshot. The end of a successful replay is
+// signalled by a `resumed` frame so the client knows it's caught up; a
+// snapshot frame tells the client to replace its local state.
+//
+// Transport-agnostic by design: a SubClient is just `{ send, readyState,
+// ctx }`. The vite-plugin wires Bun's native WebSocket in prod and the `ws`
+// package in dev to that shape.
+
+import type { Schema } from "@omg-dev/schema"
+import { getDbInstance } from "./db.ts"
+import { decodeRow } from "./codec.ts"
+import {
+  validatePredicate,
+  compileToSql,
+  compileToJs,
+  PredicateValidationError,
+  type Predicate,
+} from "./predicate.ts"
+
+// ── Wire protocol ────────────────────────────────────────────────────────────
+
+export type SubClientMessage =
+  | {
+      op: "sub"
+      subId: string
+      collection: string
+      /** Optional JSON predicate AST — narrows the snapshot. See predicate.ts. */
+      where?: Predicate
+      /**
+       * If provided, attempt to resume from the given monotonic seq instead
+       * of issuing a fresh snapshot. The server replays only events whose
+       * seq is > resumeFromSeq; if the gap exceeds the ring buffer it
+       * falls back to a snapshot. The client tells the two apart by the
+       * frame type it gets next (snapshot vs resumed).
+       */
+      resumeFromSeq?: number
+    }
+  | { op: "unsub"; subId: string }
+
+export type SubServerMessage =
+  | { type: "snapshot"; subId: string; collection: string; seq: number; rows: Record<string, unknown>[] }
+  | { type: "delta"; subId: string; collection: string; seq: number; op: "insert"; row: Record<string, unknown> }
+  | { type: "delta"; subId: string; collection: string; seq: number; op: "update"; row: Record<string, unknown> }
+  | { type: "delta"; subId: string; collection: string; seq: number; op: "delete"; id: string }
+  | { type: "resumed"; subId: string; collection: string; fromSeq: number; toSeq: number; replayed: number }
+  | { type: "error"; subId?: string; code: string; message: string }
+
+// ── Client adapter shape ─────────────────────────────────────────────────────
+
+export interface SubClient {
+  /** Send a single JSON-stringified frame to this client. */
+  send: (data: string) => void
+  /** WebSocket readyState convention: 1 = open, 3 = closed. */
+  readyState: number
+  /** Auth context captured at upgrade time. Reused on every snapshot. */
+  ctx: { userId: string | null; userEmail?: string; userName?: string; appId?: string }
+}
+
+interface ClientSub {
+  subId: string
+  collection: string
+  /** The predicate AST as sent on the wire (kept for echo / Inspect). */
+  predicate: Predicate | null
+  /** Pre-compiled JS row matcher — used by Phase 3 delta computation. */
+  predicateJs: ((row: Record<string, unknown>) => boolean) | null
+}
+
+// ── State ────────────────────────────────────────────────────────────────────
+
+// client → (subId → sub)
+const clients = new Map<SubClient, Map<string, ClientSub>>()
+// collection → set of clients with at least one sub on it (for fast fan-out)
+const subsByCollection = new Map<string, Set<SubClient>>()
+// Schema is injected by createVibesServer + reload paths. Subscriptions are
+// rejected when no schema is set or when the collection isn't declared —
+// without this, the SUB frame's `collection` string would be interpolated
+// straight into a SELECT (SQL injection vector).
+let currentSchema: Schema | null = null
+
+export function setSubscriptionSchema(schema: Schema | null): void {
+  currentSchema = schema
+}
+
+// ── Sequence numbers + per-collection ring buffer (Phase 4) ──────────────────
+//
+// `seqCounter` is module-global so deltas across collections have a strict
+// total order — easier reasoning for clients and matches the way the wire
+// protocol identifies "fresh writes" (any frame with seq > lastSeen).
+//
+// `ringByCollection` caches the raw write events so a reconnecting sub can
+// replay just the deltas it missed instead of re-snapshotting the whole
+// collection. Cap is RING_CAP per collection; once exceeded, the oldest
+// entry is dropped. A resume request that would dip below the ring's
+// oldest seq falls back to a fresh snapshot.
+
+interface RingEntry {
+  seq: number
+  kind: "insert" | "update" | "delete"
+  oldRowRaw: Record<string, unknown> | null
+  newRowRaw: Record<string, unknown> | null
+}
+
+let seqCounter = 0
+const ringByCollection = new Map<string, RingEntry[]>()
+let RING_CAP = 256
+
+// Hardening (Phase 6): per-connection subscription quota. Without a cap a
+// hostile or buggy client could open one WS and register thousands of
+// subs, each holding state, predicate evaluators, and ring-buffer
+// traversal cost. Default sized to "comfortably more than any real
+// dashboard needs" — a single page rendering 5-10 lists is normal.
+let MAX_SUBS_PER_CLIENT = 64
+
+export function setMaxSubsPerClient(n: number): void {
+  if (!Number.isInteger(n) || n < 1) throw new Error("max subs per client must be a positive integer")
+  MAX_SUBS_PER_CLIENT = n
+}
+
+/** Test/operator hook: override the per-collection ring cap. */
+export function setSubscriptionRingCap(cap: number): void {
+  if (!Number.isInteger(cap) || cap < 1) throw new Error("ring cap must be positive integer")
+  RING_CAP = cap
+}
+
+function nextSeq(): number {
+  seqCounter += 1
+  return seqCounter
+}
+
+function pushRing(collection: string, entry: RingEntry): void {
+  let ring = ringByCollection.get(collection)
+  if (!ring) { ring = []; ringByCollection.set(collection, ring) }
+  ring.push(entry)
+  if (ring.length > RING_CAP) ring.splice(0, ring.length - RING_CAP)
+}
+
+export function addSubClient(client: SubClient): void {
+  if (clients.has(client)) return
+  clients.set(client, new Map())
+}
+
+export function removeSubClient(client: SubClient): void {
+  const subs = clients.get(client)
+  if (!subs) return
+  for (const sub of subs.values()) {
+    const set = subsByCollection.get(sub.collection)
+    if (set) {
+      set.delete(client)
+      if (set.size === 0) subsByCollection.delete(sub.collection)
+    }
+  }
+  clients.delete(client)
+}
+
+export function subClientCount(): number {
+  return clients.size
+}
+
+export function subscriptionCount(collection?: string): number {
+  if (collection) return subsByCollection.get(collection)?.size ?? 0
+  let n = 0
+  for (const subs of clients.values()) n += subs.size
+  return n
+}
+
+// ── Inbound message dispatch ─────────────────────────────────────────────────
+
+export async function handleSubMessage(client: SubClient, raw: string): Promise<void> {
+  let msg: SubClientMessage
+  try {
+    msg = JSON.parse(raw) as SubClientMessage
+  } catch {
+    sendError(client, undefined, "bad_json", "message is not valid JSON")
+    return
+  }
+  if (!msg || typeof msg !== "object" || !("op" in msg)) {
+    sendError(client, undefined, "bad_shape", "missing op field")
+    return
+  }
+
+  if (msg.op === "sub") {
+    if (!msg.subId || typeof msg.subId !== "string") {
+      sendError(client, undefined, "bad_subId", "subId required")
+      return
+    }
+    if (!msg.collection || typeof msg.collection !== "string") {
+      sendError(client, msg.subId, "bad_collection", "collection required")
+      return
+    }
+    const subs = clients.get(client)
+    if (!subs) {
+      sendError(client, msg.subId, "client_not_registered", "client not registered — call addSubClient first")
+      return
+    }
+    if (subs.has(msg.subId)) {
+      sendError(client, msg.subId, "duplicate_subId", "subId already in use on this client")
+      return
+    }
+    if (subs.size >= MAX_SUBS_PER_CLIENT) {
+      sendError(client, msg.subId, "sub_limit_exceeded", `client has reached the subscription cap (${MAX_SUBS_PER_CLIENT})`)
+      return
+    }
+    // Collection must exist in the registered schema. This both whitelists
+    // the table name for safe SQL interpolation and surfaces typos.
+    if (!currentSchema || !currentSchema.collections[msg.collection]) {
+      sendError(client, msg.subId, "unknown_collection", `no such collection: ${msg.collection}`)
+      return
+    }
+    // Eager auth check for scoped collections — without this, a sub on a
+    // scoped collection from an anonymous client would get registered
+    // (and accumulate fan-out work) before the snapshot path errored.
+    // Phase 3 makes that visible because deltas would otherwise still be
+    // evaluated against the dangling sub.
+    const colDef = currentSchema.collections[msg.collection]
+    if (colDef.scope === "user" && !client.ctx.userId) {
+      sendError(client, msg.subId, "auth_required", `scoped collection "${msg.collection}" requires an authenticated user`)
+      return
+    }
+    // Predicate is optional. When present, validate before storing.
+    let predicate: Predicate | null = null
+    let predicateJs: ((row: Record<string, unknown>) => boolean) | null = null
+    if (msg.where !== undefined) {
+      try {
+        validatePredicate(msg.where)
+        predicate = msg.where
+        predicateJs = compileToJs(msg.where)
+      } catch (err) {
+        if (err instanceof PredicateValidationError) {
+          sendError(client, msg.subId, "bad_predicate", err.message)
+        } else {
+          sendError(client, msg.subId, "bad_predicate", (err as Error).message)
+        }
+        return
+      }
+    }
+    const sub: ClientSub = {
+      subId: msg.subId,
+      collection: msg.collection,
+      predicate,
+      predicateJs,
+    }
+    subs.set(msg.subId, sub)
+    let set = subsByCollection.get(msg.collection)
+    if (!set) { set = new Set(); subsByCollection.set(msg.collection, set) }
+    set.add(client)
+
+    // Resume path: if the client supplied resumeFromSeq AND the ring still
+    // has every event since then, replay them. Otherwise (no ring entries
+    // since that seq, or seq predates the ring's oldest entry) fall back
+    // to a fresh snapshot — same as the no-resume case.
+    if (
+      typeof msg.resumeFromSeq === "number" &&
+      msg.resumeFromSeq >= 0 &&
+      canResumeFrom(msg.collection, msg.resumeFromSeq)
+    ) {
+      await sendResumeReplay(client, sub, msg.resumeFromSeq)
+      return
+    }
+    await sendSnapshot(client, sub)
+    return
+  }
+
+  if (msg.op === "unsub") {
+    const subs = clients.get(client)
+    if (!subs) return
+    const sub = subs.get(msg.subId)
+    if (!sub) return
+    subs.delete(msg.subId)
+    // Only drop the client from the collection fan-out set when no other
+    // subscription on this client still targets the same collection.
+    let stillOnCollection = false
+    for (const remaining of subs.values()) {
+      if (remaining.collection === sub.collection) { stillOnCollection = true; break }
+    }
+    if (!stillOnCollection) {
+      const set = subsByCollection.get(sub.collection)
+      if (set) {
+        set.delete(client)
+        if (set.size === 0) subsByCollection.delete(sub.collection)
+      }
+    }
+    return
+  }
+
+  sendError(client, undefined, "unknown_op", `unknown op: ${(msg as { op: string }).op}`)
+}
+
+// ── Fan-out on write ─────────────────────────────────────────────────────────
+
+/**
+ * Push a per-row delta to every subscriber of `collection`. Auto-crud
+ * invokes this with raw (storage-shape) old/new rows; we decode through
+ * the schema once and then run each sub's predicate against the decoded
+ * row to compute the membership transition.
+ *
+ * Pass `oldRow=null` for inserts, `newRow=null` for deletes. For an update
+ * pass both; if a sub's predicate flips one direction we emit insert /
+ * delete deltas (the row entered or left the read set), if it stays on
+ * both sides we emit update.
+ *
+ * Returns a Promise (same pattern as notifyCollectionChange) so callers
+ * can fire-and-forget while tests await for determinism.
+ */
+export function notifyRowChange(
+  collection: string,
+  kind: "insert" | "update" | "delete",
+  oldRowRaw: Record<string, unknown> | null,
+  newRowRaw: Record<string, unknown> | null,
+): Promise<void> {
+  if (!currentSchema) return Promise.resolve()
+  const col = currentSchema.collections[collection]
+  if (!col) return Promise.resolve()
+
+  // Always bump seq + record in the ring BEFORE the subscriber-set check.
+  // A sub that subscribes a moment later may try to resume from a seq
+  // earlier than this write — the ring entry is how it catches up.
+  const seq = nextSeq()
+  pushRing(collection, { seq, kind, oldRowRaw, newRowRaw })
+
+  const set = subsByCollection.get(collection)
+  if (!set || set.size === 0) return Promise.resolve()
+
+  // Decode once per write — cheaper than per-subscriber, and the predicate
+  // matcher expects the same shape the client sees.
+  const oldRow = oldRowRaw ? decodeRow(oldRowRaw, col.fields) : null
+  const newRow = newRowRaw ? decodeRow(newRowRaw, col.fields) : null
+
+  // Scope enforcement: for user-scoped collections, the row only reaches a
+  // subscriber whose ctx.userId matches the row's _owner. This must happen
+  // BEFORE the predicate check — leaking a `delete <id>` delta to the
+  // wrong user would let them probe row existence by id.
+  const scoped = col.scope === "user"
+  const rowOwner =
+    scoped
+      ? (newRowRaw?._owner as string | undefined) ?? (oldRowRaw?._owner as string | undefined) ?? null
+      : null
+
+  const pending: Promise<void>[] = []
+  // Snapshot to a stable array — sends may evict clients mid-iteration.
+  const targets = Array.from(set)
+  for (const client of targets) {
+    if (client.readyState !== 1) {
+      removeSubClient(client)
+      continue
+    }
+    if (scoped) {
+      // Must be authed AND own the row. Anything else is a no-op for this
+      // sub (the sub's snapshot already excluded the row by _owner =
+      // userId, so dropping the delta keeps the local mirror correct).
+      if (!client.ctx.userId || client.ctx.userId !== rowOwner) continue
+    }
+    const subs = clients.get(client)
+    if (!subs) continue
+    for (const sub of subs.values()) {
+      if (sub.collection !== collection) continue
+      pending.push(dispatchDeltaWithSeq(client, sub, kind, oldRow, newRow, seq))
+    }
+  }
+  return Promise.allSettled(pending).then(() => undefined)
+}
+
+async function dispatchDeltaWithSeq(
+  client: SubClient,
+  sub: ClientSub,
+  kind: "insert" | "update" | "delete",
+  oldRow: Record<string, unknown> | null,
+  newRow: Record<string, unknown> | null,
+  seq: number,
+): Promise<boolean> {
+  // Returns true if a delta frame was actually emitted, false otherwise.
+  // The replay path uses this to count `replayed` correctly — entries
+  // whose predicate filtered the row out shouldn't be counted.
+  //
+  // No predicate → membership is trivially TRUE on both sides as long as
+  // the row exists at that side. The was/now transition then matches the
+  // raw kind 1-for-1.
+  const match = sub.predicateJs
+  const was = oldRow ? (match ? match(oldRow) : true) : false
+  const now = newRow ? (match ? match(newRow) : true) : false
+
+  if (kind === "insert") {
+    if (now) {
+      emitDelta(client, sub, "insert", newRow as Record<string, unknown>, seq)
+      return true
+    }
+    return false
+  }
+  if (kind === "delete") {
+    if (was) {
+      const id = (oldRow as Record<string, unknown>).id as string
+      sendMessage(client, {
+        type: "delta", subId: sub.subId, collection: sub.collection,
+        seq, op: "delete", id,
+      })
+      return true
+    }
+    return false
+  }
+  // kind === "update"
+  if (!was && now) {
+    emitDelta(client, sub, "insert", newRow as Record<string, unknown>, seq)
+    return true
+  }
+  if (was && now) {
+    emitDelta(client, sub, "update", newRow as Record<string, unknown>, seq)
+    return true
+  }
+  if (was && !now) {
+    const id = (oldRow as Record<string, unknown>).id as string
+    sendMessage(client, {
+      type: "delta", subId: sub.subId, collection: sub.collection,
+      seq, op: "delete", id,
+    })
+    return true
+  }
+  // !was && !now: row is outside this sub's read set on both sides — no-op.
+  return false
+}
+
+function emitDelta(
+  client: SubClient,
+  sub: ClientSub,
+  op: "insert" | "update",
+  row: Record<string, unknown>,
+  seq: number,
+): void {
+  // Keep _owner on the wire — REST `GET /api/<col>` returns the same
+  // shape today (auto-crud's SELECT * is unfiltered). Stripping here
+  // would diverge subscriber state from REST state.
+  sendMessage(client, {
+    type: "delta",
+    subId: sub.subId,
+    collection: sub.collection,
+    seq,
+    op,
+    row,
+  })
+}
+
+/**
+ * Coarse fan-out fallback: re-push the entire snapshot to every subscriber
+ * of `collection`. Used by user-authored raw-SQL handlers that mutate
+ * outside auto-crud — auto-crud itself prefers notifyRowChange. Kept for
+ * back-compat; calling this is strictly less efficient than the delta
+ * path but always correct.
+ */
+export function notifyCollectionChange(collection: string): Promise<void> {
+  const set = subsByCollection.get(collection)
+  if (!set || set.size === 0) return Promise.resolve()
+  const targets = Array.from(set)
+  const pending: Promise<void>[] = []
+  for (const client of targets) {
+    if (client.readyState !== 1) {
+      removeSubClient(client)
+      continue
+    }
+    const subs = clients.get(client)
+    if (!subs) continue
+    for (const sub of subs.values()) {
+      if (sub.collection !== collection) continue
+      pending.push(sendSnapshot(client, sub))
+    }
+  }
+  return Promise.allSettled(pending).then(() => undefined)
+}
+
+// ── Snapshot fetch + send ────────────────────────────────────────────────────
+
+async function sendSnapshot(client: SubClient, sub: ClientSub): Promise<void> {
+  const db = getDbInstance()
+  if (!db) {
+    sendError(client, sub.subId, "db_not_ready", "db not initialized")
+    return
+  }
+  if (!currentSchema) {
+    sendError(client, sub.subId, "no_schema", "schema not registered")
+    return
+  }
+  const col = currentSchema.collections[sub.collection]
+  if (!col) {
+    // Should not happen — sub-time validation rejects unknown collections.
+    sendError(client, sub.subId, "unknown_collection", `no such collection: ${sub.collection}`)
+    return
+  }
+
+  const scoped = col.scope === "user"
+  const conditions: string[] = []
+  const params: unknown[] = []
+
+  if (scoped) {
+    if (!client.ctx.userId) {
+      sendError(client, sub.subId, "auth_required", `scoped collection "${sub.collection}" requires an authenticated user`)
+      return
+    }
+    conditions.push("_owner = ?")
+    params.push(client.ctx.userId)
+  }
+
+  if (sub.predicate) {
+    // Reject predicate columns that don't exist in this collection's schema
+    // — same defense-in-depth as the column-identifier regex inside the
+    // predicate validator. (validatePredicate already rejected non-
+    // identifier strings; this also rejects columns whose name is well-
+    // formed but undeclared.)
+    if (!predicateColumnsExist(sub.predicate, col.fields)) {
+      sendError(client, sub.subId, "bad_predicate", "predicate references unknown column")
+      return
+    }
+    const compiled = compileToSql(sub.predicate)
+    conditions.push(compiled.sql)
+    params.push(...compiled.params)
+  }
+
+  const where = conditions.length ? `WHERE ${conditions.join(" AND ")}` : ""
+  try {
+    const rawRows = db.raw()
+      .prepare(`SELECT * FROM ${sub.collection} ${where} ORDER BY created_at DESC`)
+      .all(...params) as Record<string, unknown>[]
+    const rows = rawRows.map(r => decodeRow(r, col.fields))
+    // The snapshot's high-water mark is the current global seq — any delta
+    // with seq > hwm happened AFTER this snapshot's SELECT. Note that we
+    // don't bump the counter for snapshots; their seq is shared with the
+    // most recent write so resume-from-seq math is straightforward.
+    sendMessage(client, {
+      type: "snapshot",
+      subId: sub.subId,
+      collection: sub.collection,
+      seq: seqCounter,
+      rows,
+    })
+  } catch (err) {
+    sendError(client, sub.subId, "snapshot_failed", (err as Error).message ?? String(err))
+  }
+}
+
+// ── Resume replay ────────────────────────────────────────────────────────────
+
+function canResumeFrom(collection: string, fromSeq: number): boolean {
+  const ring = ringByCollection.get(collection)
+  // No ring (no writes since boot) is a successful resume — there's
+  // nothing to replay and the client is by definition already caught up.
+  if (!ring || ring.length === 0) return fromSeq <= seqCounter
+  const oldest = ring[0]
+  if (!oldest) return fromSeq <= seqCounter
+  return fromSeq >= oldest.seq - 1
+}
+
+async function sendResumeReplay(
+  client: SubClient,
+  sub: ClientSub,
+  fromSeq: number,
+): Promise<void> {
+  if (!currentSchema) {
+    sendError(client, sub.subId, "no_schema", "schema not registered")
+    return
+  }
+  const col = currentSchema.collections[sub.collection]
+  if (!col) {
+    sendError(client, sub.subId, "unknown_collection", `no such collection: ${sub.collection}`)
+    return
+  }
+  const ring = ringByCollection.get(sub.collection) ?? []
+  let replayed = 0
+  for (const entry of ring) {
+    if (entry.seq <= fromSeq) continue
+    // Same scope + predicate gate notifyRowChange runs. Phase 3's logic
+    // is in dispatchDelta, but it expects decoded rows. Decode here.
+    const scoped = col.scope === "user"
+    if (scoped) {
+      const owner =
+        (entry.newRowRaw?._owner as string | undefined) ??
+        (entry.oldRowRaw?._owner as string | undefined) ??
+        null
+      if (!client.ctx.userId || client.ctx.userId !== owner) continue
+    }
+    const oldRow = entry.oldRowRaw ? decodeRow(entry.oldRowRaw, col.fields) : null
+    const newRow = entry.newRowRaw ? decodeRow(entry.newRowRaw, col.fields) : null
+    const emitted = await dispatchDeltaWithSeq(client, sub, entry.kind, oldRow, newRow, entry.seq)
+    if (emitted) replayed += 1
+  }
+  sendMessage(client, {
+    type: "resumed",
+    subId: sub.subId,
+    collection: sub.collection,
+    fromSeq,
+    toSeq: seqCounter,
+    replayed,
+  })
+}
+
+const RESERVED = new Set(["id", "created_at", "updated_at", "_owner"])
+
+function predicateColumnsExist(
+  p: Predicate,
+  fields: Record<string, { type: unknown }>,
+): boolean {
+  switch (p.op) {
+    case "and":
+    case "or":
+      return p.clauses.every(c => predicateColumnsExist(c, fields))
+    case "not":
+      return predicateColumnsExist(p.clause, fields)
+    default:
+      // All leaf nodes carry a `column` property.
+      return RESERVED.has(p.column) || p.column in fields
+  }
+}
+
+function sendMessage(client: SubClient, msg: SubServerMessage): void {
+  if (client.readyState !== 1) {
+    removeSubClient(client)
+    return
+  }
+  try {
+    client.send(JSON.stringify(msg))
+  } catch {
+    removeSubClient(client)
+  }
+}
+
+function sendError(client: SubClient, subId: string | undefined, code: string, message: string): void {
+  sendMessage(client, { type: "error", subId, code, message })
+}
+
+// ── Test-only escape hatch ───────────────────────────────────────────────────
+
+/** Clear all subscription state. Used by tests; not part of the public API. */
+export function _resetSubscriptionState(): void {
+  clients.clear()
+  subsByCollection.clear()
+  ringByCollection.clear()
+  seqCounter = 0
+  RING_CAP = 256
+  MAX_SUBS_PER_CLIENT = 64
+}
+
+/** Test helper: peek at the current global seq. Not part of the public API. */
+export function _currentSeq(): number {
+  return seqCounter
+}