statewire 0.16.0 → 0.17.1

package/src/host/socket.ts

@@ -1,3 +1,12 @@
+import { encodeDocument } from "../document";
+import {
+  DEFAULT_PROTOCOL,
+  WIRE_VERSION,
+  negotiateProtocols,
+} from "../wire/protocol";
+import type { StatewireProtocol } from "../wire/protocol";
+import { applyDocumentOps, mainDocument } from "../wire/documents";
+import type { StatewireDocuments } from "../wire/documents";
 import { useEffect, useState } from "react";
 import { createTapRoot, flushTapSync, useResource } from "@assistant-ui/tap";
 import type { ResourceElement } from "@assistant-ui/tap";
@@ -7,7 +16,7 @@ import { StatewireSendError } from "../types";
 import type { Statewire } from "../types";
 import type { StatewireHost } from "../host";
 import type { StatewireHostInternal } from "../host-internal";
-import { encodeValue, reviveValue } from "../wire/meta";
+import { reviveValue } from "../wire/meta";
 import {
   COMMAND_CAPACITY_RETRY_MS,
   MAX_COMMANDS_PER_FRAME,
@@ -17,7 +26,7 @@ import {
   MAX_INFLIGHT_COMMANDS_PER_HOST,
 } from "../wire/limits";
 import { createMacrotaskPoster } from "../util/macrotask";
-import type { StatewireMetaEntry } from "../wire/meta";
+import type { StatewireDocument } from "../document";
 import type { StatewirePath } from "../wire/ops";
 import type { StatewireCmd, StatewireEnvelope } from "../client/connection";
 import { frames as serveFrames, stream as serveStream } from "./http";
@@ -71,7 +80,7 @@ type WireOp = {
   op: string;
   path?: StatewirePath;
   value?: unknown;
-  meta?: StatewireMetaEntry[];
+  meta?: StatewireDocument.MetaEntry[];
 };
 
 type Admission = {
@@ -88,6 +97,7 @@ export type ClientRec = {
   ctxSeq: number;
   touched: number;
   context: Record<string, unknown> | null;
+  contextDocuments: StatewireDocuments.State;
   results: Map<number, { ts: number; cmd: StatewireCmd }>;
   inflight: Map<number, Admission>;
   /** The auth captured at the record's first attach; later attaches must match. */
@@ -95,6 +105,7 @@ export type ClientRec = {
 };
 
 export type Conn = {
+  selection: StatewireProtocol.Selection;
   clientId: string;
   socket: StatewireHostInternal.Socket;
   open: boolean;
@@ -125,8 +136,9 @@ const violationReason = (error: unknown): string | null =>
     ? error.message
     : null;
 
-type CtxStatement = { seq: number; value: Record<string, unknown> };
+type CtxStatement = { seq: number; ops: unknown[] };
 type CmdStatement = {
+  protocol: string;
   seq: number;
   method: string;
   params: unknown;
@@ -161,30 +173,6 @@ const parseLane = (value: unknown): Record<string, unknown>[] => {
 
 const frameEncoder = new TextEncoder();
 
-/** v1 ctx statements state the context whole: one root replace with an object value. */
-const contextValue = (ops: unknown): Record<string, unknown> => {
-  if (!Array.isArray(ops) || ops.length !== 1) {
-    throw new Error("context must be a single-op array");
-  }
-  const op = ops[0] as Record<string, unknown> | null;
-  if (
-    !isPlainObject(op) ||
-    op["op"] !== "replace" ||
-    !Array.isArray(op["path"]) ||
-    op["path"].length !== 0
-  ) {
-    throw new Error("context op must be a root replace");
-  }
-  const value =
-    op["meta"] !== undefined
-      ? reviveValue(op["value"], op["meta"], { rejectUnknown: true })
-      : op["value"];
-  if (!isPlainObject(value)) {
-    throw new Error("context value must be an object");
-  }
-  return value;
-};
-
 export const parseFrameValue = (frame: unknown): ParsedFrame => {
   if (!isPlainObject(frame)) return violation("malformed frame");
   const keys = Object.keys(frame);
@@ -199,19 +187,15 @@ export const parseFrameValue = (frame: unknown): ParsedFrame => {
       return violation("too many context statements");
     }
     for (const statement of parseLane(lane)) {
-      let value: Record<string, unknown>;
-      try {
-        value = contextValue(statement["ops"]);
-      } catch {
+      const ops = statement["ops"];
+      if (!Array.isArray(ops) || !ops.length)
         return violation("invalid-context");
-      }
       if (
-        frameEncoder.encode(JSON.stringify(value)).byteLength >
+        frameEncoder.encode(JSON.stringify(ops)).byteLength >
         MAX_CONTEXT_VALUE_BYTES
-      ) {
+      )
         return violation("oversized-context");
-      }
-      ctx.push({ seq: statement["seq"] as number, value });
+      ctx.push({ seq: statement["seq"] as number, ops });
     }
   }
   const cmd: CmdStatement[] = [];
@@ -222,10 +206,15 @@ export const parseFrameValue = (frame: unknown): ParsedFrame => {
     }
     for (const statement of parseLane(lane)) {
       const method = statement["method"];
-      if (typeof method !== "string" || !("params" in statement)) {
+      if (
+        typeof method !== "string" ||
+        typeof statement["protocol"] !== "string" ||
+        !("params" in statement)
+      ) {
         return violation("malformed statement");
       }
       cmd.push({
+        protocol: statement["protocol"] as string,
         seq: statement["seq"] as number,
         method,
         params: statement["params"],
@@ -254,9 +243,9 @@ const frameIsOversized = (raw: string): boolean =>
 
 const encodePayload = (
   payload: unknown,
-): { payload?: unknown; meta?: StatewireMetaEntry[] } => {
+): { payload?: unknown; meta?: StatewireDocument.MetaEntry[] } => {
   if (payload === undefined) return {};
-  const { value, meta } = encodeValue(payload);
+  const { value, meta } = encodeDocument(payload);
   return { payload: value, ...(meta !== undefined && { meta }) };
 };
 
@@ -334,6 +323,8 @@ const parseInitialFrame = (raw: string) => {
 // ─── host ───────────────────────────────────────────────
 
 export type SocketHostState = {
+  protocol: StatewireProtocol.Definition;
+  wireVersion: StatewireProtocol.Range;
   heartbeatMs: number;
   wsInitTimeoutMs: number;
   maxSubscribers: number;
@@ -384,7 +375,7 @@ const recordOp = (
   },
 ) => {
   if (op.op === "replace" || op.op === "add") {
-    const { value, meta } = encodeValue(op.value);
+    const { value, meta } = encodeDocument(op.value);
     h.pendingOps.push({ ...op, value, ...(meta !== undefined && { meta }) });
   } else {
     h.pendingOps.push(op);
@@ -518,7 +509,9 @@ export const flush = (h: SocketHostState, fin?: StatewireHost.Fin) => {
     ).map((seq) => ({ seq, type: "applied" }));
     if (ops.length > 0) {
       const packet: StatewireEnvelope = {
-        ops,
+        ops: conn.selection.protocols.some((p) => p.name === h.protocol.name)
+          ? ops
+          : [],
         ...(appliedEntries.length > 0 && { cmd: appliedEntries }),
         ...(edge && { idle: level }),
       };
@@ -526,7 +519,11 @@ export const flush = (h: SocketHostState, fin?: StatewireHost.Fin) => {
     } else if (appliedEntries.length > 0) {
       sendTo(h, conn, JSON.stringify({ cmd: appliedEntries }));
     }
-    const verdicts = cmds.get(conn.clientId) ?? [];
+    const verdicts = (cmds.get(conn.clientId) ?? []).map((cmd) =>
+      conn.selection.protocols.some((p) => p.name === h.protocol.name)
+        ? cmd
+        : { seq: cmd.seq, ...(cmd.dur === undefined ? {} : { dur: cmd.dur }) },
+    );
     if (verdicts.length > 0) {
       sendTo(h, conn, JSON.stringify({ cmd: verdicts }));
     }
@@ -562,6 +559,7 @@ const notifyClientRecords = (h: SocketHostState) => {
       ctxSeq: rec.ctxSeq,
       touched: rec.touched,
       context: rec.context,
+      contextDocuments: [...rec.contextDocuments.documents],
       results: [...rec.results].map(([seq, entry]) => ({
         seq,
         ts: entry.ts,
@@ -672,6 +670,7 @@ const admitCmd = (
   rec: ClientRec,
   statement: CmdStatement,
   auth: unknown,
+  selection: StatewireProtocol.Selection,
 ) => {
   const adm: Admission = {
     seq: statement.seq,
@@ -682,6 +681,18 @@ const admitCmd = (
   };
   rec.inflight.set(statement.seq, adm);
   notifyClientRecords(h);
+  if (
+    statement.protocol !== h.protocol.name ||
+    !selection.protocols.some((p) => p.name === statement.protocol)
+  ) {
+    settleCmd(h, rec, adm, {
+      seq: adm.seq,
+      type: "rejected",
+      code: "unsupported-protocol",
+      message: "command protocol was not negotiated",
+    });
+    return;
+  }
   if (!Array.isArray(statement.params)) {
     rejectStatement(h, rec, adm, "params must be an array");
     return;
@@ -749,6 +760,10 @@ export const admitFrame = (
   rec: ClientRec,
   frame: ParsedFrame,
   auth: unknown,
+  selection: StatewireProtocol.Selection = {
+    version: h.wireVersion.version,
+    protocols: [{ name: h.protocol.name, version: h.protocol.version }],
+  },
 ): FrameAdmission => {
   rec.touched = Date.now();
   const through = admittedThrough(rec);
@@ -801,11 +816,28 @@ export const admitFrame = (
       };
     }
   }
-  for (const statement of ctxNew) {
-    rec.context = statement.value;
-    rec.ctxSeq = statement.seq;
+  let contextDocuments = rec.contextDocuments;
+  try {
+    for (const statement of ctxNew)
+      contextDocuments = applyDocumentOps(
+        contextDocuments,
+        statement.ops,
+        selection.protocols,
+        true,
+      );
+    if (
+      frameEncoder.encode(JSON.stringify([...contextDocuments.documents]))
+        .byteLength > MAX_CONTEXT_VALUE_BYTES
+    )
+      return violation("oversized-context");
+  } catch {
+    return violation("invalid-context");
   }
-  for (const statement of cmdNew) admitCmd(h, rec, statement, auth);
+  rec.contextDocuments = contextDocuments;
+  rec.context = (mainDocument(contextDocuments, h.protocol.name)?.value ??
+    null) as Record<string, unknown> | null;
+  if (ctxNew.length) rec.ctxSeq = ctxNew.at(-1)!.seq;
+  for (const statement of cmdNew) admitCmd(h, rec, statement, auth, selection);
   if (cmdNew.length > 0) {
     // Memory host: admitted means retained — dur rides the admitting flush.
     queueCmd(h, rec.id, { seq: cmdNew.at(-1)!.seq, dur: true });
@@ -854,7 +886,20 @@ const message = (h: SocketHostState, conn: Conn, data: string) => {
   }
   if (frame.ctx.length === 0 && frame.cmd.length === 0) return;
   const rec = h.clients.get(conn.clientId)!;
-  const admission = admitFrame(h, rec, frame, conn.auth);
+  let admission: ReturnType<typeof admitFrame>;
+  try {
+    admission = admitFrame(h, rec, frame, conn.auth, conn.selection);
+  } catch (error) {
+    const reason = violationReason(error);
+    if (reason === null) throw error;
+    sendTo(
+      h,
+      conn,
+      JSON.stringify({ fin: { reason: "protocol-error", message: reason } }),
+    );
+    closeConn(h, conn, 1008, reason);
+    return;
+  }
   if (admission.type === "refused") {
     sendTo(h, conn, JSON.stringify(admission.syn));
   }
@@ -965,6 +1010,7 @@ const connectAuthorized = (
       socket,
       options.authorized === true,
       options.auth,
+      options.selection,
     );
     return registration?.connection ?? null;
   }
@@ -1056,7 +1102,10 @@ const connectAuthorized = (
       const frame = initial.frame;
       let headers: Record<string, string>;
       try {
-        headers = normalizedHeaders(initial.headers);
+        headers = {
+          ...normalizedHeaders(initial.headers),
+          ...normalizedHeaders(options.headers ?? {}),
+        };
       } catch {
         sendPending(
           pending,
@@ -1089,17 +1138,39 @@ const connectAuthorized = (
         .then(() => authorize?.({ headers }))
         .then((auth) => {
           if (!pending.open || h.disposed || h.finished !== null) return;
-          options.onAuthorized?.(auth);
+          const selection = negotiateProtocols(
+            new Headers(headers),
+            h.wireVersion,
+            [h.protocol],
+          );
+          if ("fin" in selection) {
+            sendPending(pending, JSON.stringify(selection));
+            closePending(h, pending, 1008, selection.fin.message);
+            return;
+          }
+          options.onAuthorized?.(auth, selection);
           pending.open = false;
           h.pendingConns.delete(pending);
           if (pending.timer !== null) clearTimeout(pending.timer);
-          const registration = register(h, clientId, socket, true, auth);
+          const registration = register(
+            h,
+            clientId,
+            socket,
+            true,
+            auth,
+            selection,
+          );
           if (registration === null) return;
           pending.connection = registration.connection;
           if (frame.ctx.length || frame.cmd.length) {
-            const admission = admitFrame(h, registration.rec, frame, auth);
-            if (admission.type === "refused")
-              sendTo(h, registration.conn, JSON.stringify(admission.syn));
+            message(
+              h,
+              registration.conn,
+              JSON.stringify({
+                ...(frame.ctx.length ? { ctx: frame.ctx } : {}),
+                ...(frame.cmd.length ? { cmd: frame.cmd } : {}),
+              }),
+            );
           }
           for (const data of pending.queued.splice(0))
             registration.connection.message(data);
@@ -1143,6 +1214,10 @@ export const register = (
   socket: StatewireHostInternal.Socket,
   sawFrame = false,
   auth: unknown = undefined,
+  selection: StatewireProtocol.Selection = {
+    version: h.wireVersion.version,
+    protocols: [{ name: h.protocol.name, version: h.protocol.version }],
+  },
 ): Registration | null => {
   const refuse = (code: number, reason: string) => {
     try {
@@ -1193,6 +1268,7 @@ export const register = (
         ctxSeq: 0,
         touched: Date.now(),
         context: null,
+        contextDocuments: { mode: "mounted", documents: new Map() },
         results: new Map(),
         inflight: new Map(),
         principal: auth,
@@ -1202,6 +1278,7 @@ export const register = (
     })();
   rec.touched = Date.now();
   const conn: Conn = {
+    selection,
     clientId,
     socket,
     open: true,
@@ -1214,22 +1291,26 @@ export const register = (
   gcClients(h);
   notifyClientRecords(h);
   ensureHeartbeat(h);
-  const { value, meta } = encodeValue(h.stateHost!.snapshot());
+  const { value, meta } = encodeDocument(h.stateHost!.snapshot());
   const hello: StatewireEnvelope = {
     ops: [
       {
-        op: "replace",
-        path: [],
+        op: "mount",
+        protocol: h.protocol.name,
+        main: true,
         value,
         ...(meta !== undefined && { meta }),
       },
     ],
     syn: {
+      ...selection,
       seq: admittedThrough(rec),
       ctxSeq: rec.ctxSeq,
       ...(h.idleWire && { idle: true }),
     },
   };
+  if (!selection.protocols.some((p) => p.name === h.protocol.name))
+    hello.ops = [];
   sendTo(h, conn, JSON.stringify(hello));
   const connection: StatewireHostInternal.SocketConnection = {
     clientId,
@@ -1287,6 +1368,10 @@ const restoreClientRecords = (
       record.id,
       {
         ...record,
+        contextDocuments: {
+          mode: "mounted",
+          documents: new Map(record.contextDocuments ?? []),
+        },
         results: new Map(
           record.results.map(({ seq, ts, cmd }) => [seq, { ts, cmd }]),
         ),
@@ -1305,6 +1390,8 @@ const useSocketHost = (
 > => {
   const [h] = useState((): SocketHostState => {
     const host: SocketHostState = {
+      protocol: { name: DEFAULT_PROTOCOL, version: WIRE_VERSION },
+      wireVersion: { version: WIRE_VERSION },
       heartbeatMs: 15_000,
       wsInitTimeoutMs: 5_000,
       maxSubscribers: MAX_SUBSCRIBERS,
@@ -1341,6 +1428,11 @@ const useSocketHost = (
     };
     return host;
   });
+  h.protocol = options?.protocol ?? {
+    name: DEFAULT_PROTOCOL,
+    version: WIRE_VERSION,
+  };
+  h.wireVersion = options?.wireVersion ?? { version: WIRE_VERSION };
   setHeartbeatMs(h, options?.heartbeatIntervalMs ?? 15_000);
   h.pushback = options?.pushback;
   h.wsInitTimeoutMs = options?.wsInitTimeoutMs ?? 5_000;

package/src/host-internal.ts

@@ -1,3 +1,5 @@
+import type { StatewireProtocol } from "./wire/protocol";
+import type { StatewireDocuments } from "./wire/documents";
 import type { StatewireHost } from "./host";
 import type { StatewireCmd } from "./client/connection";
 
@@ -17,12 +19,18 @@ export namespace StatewireHostInternal {
   };
 
   export type SocketConnectOptions = {
+    /** Upgrade headers take precedence over initialization headers. */
+    headers?: Record<string, string>;
     /** The adapter restored this socket from an already-authorized session. */
     authorized?: boolean;
+    selection?: StatewireProtocol.Selection;
     /** The restored session's captured auth value; it must satisfy the client record's principal binding. */
     auth?: unknown;
     /** Persist authorization (and the captured auth value) before the first snapshot is released. */
-    onAuthorized?: (auth: unknown) => void;
+    onAuthorized?: (
+      auth: unknown,
+      selection: StatewireProtocol.Selection,
+    ) => void;
   };
 
   export type SocketHost = {
@@ -44,6 +52,7 @@ export namespace StatewireHostInternal {
     ctxSeq: number;
     touched: number;
     context: Record<string, unknown> | null;
+    contextDocuments?: readonly [number, StatewireDocuments.Document][];
     results: readonly { seq: number; ts: number; cmd: StatewireCmd }[];
     principal: unknown;
   };
@@ -54,6 +63,8 @@ export namespace StatewireHostInternal {
   };
 
   export type SocketHostOptions<TAuth = unknown> = {
+    protocol?: StatewireProtocol.Definition;
+    wireVersion?: StatewireProtocol.Range;
     /** Gates every attach; the awaited return value is the attach's auth, exposed to command handlers as `this.caller.auth`. The first attach binds the client id's record to this value; a later attach whose value is not structurally equal is refused (HTTP 403, socket close 1008). */
     authorize?: (ctx: {
       headers: Record<string, string>;

package/src/index.ts

@@ -1,5 +1,10 @@
 export type { Statewire } from "./types";
+export { encodeDocument, decodeDocument } from "./document";
+export type { StatewireDocument } from "./document";
 export { StatewireClient, useStatewire, StatewireResource } from "./statewire";
 export { StatewireHttp } from "./transports/sse";
 export { StatewireWebsocket } from "./transports/ws";
 export { StatewireSendError } from "./types";
+
+export type { StatewireProtocol } from "./wire/protocol";
+export type { StatewireDocuments } from "./wire/documents";