feat(shared-dataflow): UpdatesStore orderBy + streaming k-way URI merge

UpdatesStore.readEntries / readUpdatedEntries gain an optional
`orderBy: "stamp" | "uri"` argument (new exported ReadOrderBy type).
Default stays "stamp" — the temporal-replay order existing callers
rely on. `"uri"` yields URI-ascending, which is the only order under
which several per-signal diff streams can be merged without buffering.

readUpstreamChanges rewritten to open one orderBy:"uri" iterator per
upstream signal of the cell and run a streaming k-way merge by URI.
Memory is now O(M) in the number of upstream signals, independent of
the number of matching URIs. URI ties between upstreams still resolve
to signal-declaration order (graph.getCellInputs). The iterators are
closed in a finally block so early `break` doesn't leak.

Tests: six new cases on InMemoryUpdatesStore covering the default vs
explicit "stamp", "uri" ordering for both read methods, and the
filter-independence property; one new case on readUpstreamChanges
asserting non-buffering behaviour via a Proxy-spied store (reads stop
when the consumer breaks). 124 → 131 tests in the package.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: mkotelnikov

packages/shared-dataflow/src/in-memory-updates-store.ts

@@ -1,5 +1,20 @@
 import type { Signal } from "./types.js";
-import type { SerializedUpdatesStore, UpdateEntry, UpdatesStore } from "./updates-store.js";
+import type {
+  ReadOrderBy,
+  SerializedUpdatesStore,
+  UpdateEntry,
+  UpdatesStore,
+} from "./updates-store.js";
+
+function compareMatches(orderBy: ReadOrderBy | undefined) {
+  if (orderBy === "uri") {
+    return (a: [string, number], b: [string, number]) => {
+      if (a[0] === b[0]) return 0;
+      return a[0] < b[0] ? -1 : 1;
+    };
+  }
+  return (a: [string, number], b: [string, number]) => a[1] - b[1];
+}
 
 /**
  * In-memory `UpdatesStore` reference implementation. State lives in this
@@ -59,6 +74,7 @@ export class InMemoryUpdatesStore implements UpdatesStore {
     signal: Signal;
     since: number;
     uriPrefix?: string;
+    orderBy?: ReadOrderBy;
   }): AsyncIterable<UpdateEntry> {
     const inner = this.state.get(opts.signal);
     if (!inner) return;
@@ -69,7 +85,7 @@ export class InMemoryUpdatesStore implements UpdatesStore {
         matches.push([uri, stamp]);
       }
     }
-    matches.sort((a, b) => a[1] - b[1]);
+    matches.sort(compareMatches(opts.orderBy));
     for (const [uri, stamp] of matches) {
       yield { signal: opts.signal, uri, stamp };
     }
@@ -79,6 +95,7 @@ export class InMemoryUpdatesStore implements UpdatesStore {
     upstreamSignal: Signal;
     currentSignal: Signal;
     uriPrefix?: string;
+    orderBy?: ReadOrderBy;
   }): AsyncIterable<UpdateEntry> {
     const upstream = this.state.get(opts.upstreamSignal);
     if (!upstream) return;
@@ -90,7 +107,7 @@ export class InMemoryUpdatesStore implements UpdatesStore {
       const currentStamp = current?.get(uri) ?? 0;
       if (upstreamStamp > currentStamp) matches.push([uri, upstreamStamp]);
     }
-    matches.sort((a, b) => a[1] - b[1]);
+    matches.sort(compareMatches(opts.orderBy));
     for (const [uri, stamp] of matches) {
       yield { signal: opts.upstreamSignal, uri, stamp };
     }

packages/shared-dataflow/src/index.ts

@@ -14,4 +14,9 @@ export type {
   UpdatesManagerOptions,
 } from "./updates-manager.js";
 export { UpdatesManager } from "./updates-manager.js";
-export type { SerializedUpdatesStore, UpdateEntry, UpdatesStore } from "./updates-store.js";
+export type {
+  ReadOrderBy,
+  SerializedUpdatesStore,
+  UpdateEntry,
+  UpdatesStore,
+} from "./updates-store.js";

packages/shared-dataflow/src/read-upstream-changes.ts

@@ -9,25 +9,23 @@ import type { UpdateEntry, UpdatesStore } from "./updates-store.js";
  * declared output (convention: cells whose primary output also serves
  * as their per-URI watermark).
  *
- * Opens one diff stream per upstream signal (via
- * `readUpdatedEntries`), merges them, and yields the combined result
- * sorted by URI (lexicographic ascending). When the same URI appears
- * in multiple upstream signals, the entries are emitted adjacently in
- * the order the upstream signals were declared in
- * `graph.getCellInputs(cellId)` — so consumers can collapse
- * per-URI in one pass.
+ * Opens one URI-ordered diff stream per upstream signal (via
+ * `readUpdatedEntries({orderBy: "uri"})`) and merges them with a
+ * streaming k-way URI merge. Memory is O(M) where M is the number of
+ * upstream signals — independent of the number of matching URIs.
+ *
+ * Output order is URI-ascending. When the same URI appears in multiple
+ * upstream signals, the entries are emitted adjacently in the order
+ * the upstream signals were declared in `graph.getCellInputs(cellId)`,
+ * so consumers can collapse per-URI in one forward pass without
+ * buffering.
  *
  * A URI updated by N upstream signals appears N times (once per
- * upstream entry). Use `aggregateByUri` if the consumer wants
- * one record per URI.
+ * upstream entry). Use `aggregateByUri` if the consumer wants one
+ * record per URI.
  *
  * Throws if the cell has inputs but no outputs (no watermark to compare
  * against). Cells with no inputs (probers) yield nothing.
- *
- * Implementation note: URI-sorted output requires buffering all
- * matching entries before yielding, since the per-signal streams are
- * stamp-ordered, not URI-ordered. Memory = O(matching URIs × upstream
- * signals).
  */
 export async function* readUpstreamChanges(
   store: UpdatesStore,
@@ -46,33 +44,64 @@ export async function* readUpstreamChanges(
   }
   const uriPrefix = opts?.uriPrefix;
 
-  // Drain every upstream's per-URI diff into one buffer so we can sort
-  // the combined stream by URI. (Within a single signal `readUpdatedEntries`
-  // yields by stamp; merging by URI across signals isn't possible without
-  // buffering.)
-  const signalOrder = new Map<string, number>();
-  for (let i = 0; i < inputs.length; i++) {
-    signalOrder.set(inputs[i] as string, i);
-  }
-  const buffered: UpdateEntry[] = [];
-  for (const upstream of inputs) {
-    for await (const entry of store.readUpdatedEntries({
+  // One URI-ordered iterator per upstream signal.
+  const iters: Array<AsyncIterator<UpdateEntry>> = inputs.map((upstream) => {
+    const iterable = store.readUpdatedEntries({
       upstreamSignal: upstream,
       currentSignal: watermark,
       uriPrefix,
-    })) {
-      buffered.push(entry);
-    }
-  }
-
-  buffered.sort((a, b) => {
-    if (a.uri !== b.uri) return a.uri < b.uri ? -1 : 1;
-    const sa = signalOrder.get(a.signal) ?? 0;
-    const sb = signalOrder.get(b.signal) ?? 0;
-    return sa - sb;
+      orderBy: "uri",
+    });
+    return iterable[Symbol.asyncIterator]();
   });
 
-  for (const entry of buffered) yield entry;
+  try {
+    // Prime the head of each stream.
+    const heads: Array<UpdateEntry | null> = await Promise.all(
+      iters.map(async (it) => {
+        const r = await it.next();
+        return r.done ? null : r.value;
+      }),
+    );
+
+    // Streaming k-way merge by URI. For URI ties, keep
+    // signal-declaration order by preferring the smaller index.
+    while (true) {
+      let minIdx = -1;
+      for (let i = 0; i < heads.length; i++) {
+        const h = heads[i];
+        if (h === null || h === undefined) continue;
+        if (minIdx === -1) {
+          minIdx = i;
+          continue;
+        }
+        // biome-ignore lint/style/noNonNullAssertion: minIdx ≥ 0 implies heads[minIdx] is not null
+        const cur = heads[minIdx]!;
+        if (h.uri < cur.uri) minIdx = i;
+      }
+      if (minIdx === -1) return;
+      // biome-ignore lint/style/noNonNullAssertion: minIdx selected from a non-null head
+      const entry = heads[minIdx]!;
+      yield entry;
+      // biome-ignore lint/style/noNonNullAssertion: iters[minIdx] paired with heads[minIdx]
+      const next = await iters[minIdx]!.next();
+      heads[minIdx] = next.done ? null : next.value;
+    }
+  } finally {
+    // Best-effort close of any non-exhausted iterators on early break.
+    await Promise.all(
+      iters.map(async (it) => {
+        if (it.return) {
+          try {
+            await it.return(undefined);
+          } catch {
+            // Swallow: a throwing iterator close must not mask the
+            // primary outcome of the merge.
+          }
+        }
+      }),
+    );
+  }
 }
 
 /**

packages/shared-dataflow/src/updates-store.ts

@@ -42,19 +42,33 @@ export type SerializedUpdatesStore = {
  * signal naming consistency with any `DataflowGraph` topology, tombstone
  * naming conventions.
  */
+/**
+ * Yield ordering for read methods on `UpdatesStore`. Both options yield
+ * the same set of entries — only the order differs.
+ *
+ * - `"stamp"` (default) — last-transaction-id ascending. Best when the
+ *   consumer wants temporal-order replay.
+ * - `"uri"` — URI ascending. Best when the consumer wants to collapse
+ *   per-URI work in one forward pass, or to merge multiple read streams
+ *   by URI without buffering (see `readUpstreamChanges`).
+ */
+export type ReadOrderBy = "stamp" | "uri";
+
 export interface UpdatesStore {
   /**
    * Read entries on a signal whose stamp > `since`, optionally filtered by
-   * URI prefix. Yields full `UpdateEntry` objects in stamp-ascending order.
+   * URI prefix. Yields full `UpdateEntry` objects.
    *
    * - `since` is exclusive: `stamp > since`. Use `since = 0` to read everything.
    * - `uriPrefix` (optional, defaults to no filter) restricts to entries whose
    *   `uri.startsWith(uriPrefix)`. An empty string is treated as no filter.
+   * - `orderBy` (optional, default `"stamp"`) — see `ReadOrderBy`.
    */
   readEntries(opts: {
     signal: Signal;
     since: number;
     uriPrefix?: string;
+    orderBy?: ReadOrderBy;
   }): AsyncIterable<UpdateEntry>;
 
   /**
@@ -72,8 +86,10 @@ export interface UpdatesStore {
    * since cleaned it up via tombstone, so the current cell has nothing
    * to do.
    *
-   * Yields in upstream-stamp-ascending order. `uriPrefix` filters the
-   * same way as `readEntries`.
+   * Default order is upstream-stamp-ascending. Pass `orderBy: "uri"`
+   * for URI-ascending order — used by `readUpstreamChanges` to merge
+   * several per-signal streams in O(1) buffering. `uriPrefix` filters
+   * the same way as `readEntries`.
    *
    * Caller contract: after handling a yielded entry, save a
    * `currentSignal` entry with stamp >= the yielded upstream stamp to
@@ -84,6 +100,7 @@ export interface UpdatesStore {
     upstreamSignal: Signal;
     currentSignal: Signal;
     uriPrefix?: string;
+    orderBy?: ReadOrderBy;
   }): AsyncIterable<UpdateEntry>;
 
   /** Upsert by `(signal, uri)`. Replaces blindly — last write wins. */

packages/shared-dataflow/tests/in-memory-updates-store.test.ts

@@ -568,6 +568,101 @@ describe("InMemoryUpdatesStore — readUpdatedEntries (per-URI cell diff)", () =
   });
 });
 
+describe("InMemoryUpdatesStore — orderBy 'stamp' vs 'uri'", () => {
+  it("readEntries defaults to stamp-ascending order", async () => {
+    const store = new InMemoryUpdatesStore();
+    await store.saveEntries([
+      { signal: "x", uri: "c", stamp: 9 },
+      { signal: "x", uri: "a", stamp: 3 },
+      { signal: "x", uri: "b", stamp: 5 },
+    ]);
+    const got = await collect(store.readEntries({ signal: "x", since: 0 }));
+    expect(got.map((e) => e.uri)).toEqual(["a", "b", "c"]);
+    expect(got.map((e) => e.stamp)).toEqual([3, 5, 9]);
+  });
+
+  it("readEntries with orderBy: 'uri' yields URI-ascending regardless of stamp", async () => {
+    const store = new InMemoryUpdatesStore();
+    await store.saveEntries([
+      { signal: "x", uri: "b", stamp: 9 },
+      { signal: "x", uri: "a", stamp: 100 },
+      { signal: "x", uri: "c", stamp: 1 },
+    ]);
+    const got = await collect(store.readEntries({ signal: "x", since: 0, orderBy: "uri" }));
+    expect(got.map((e) => e.uri)).toEqual(["a", "b", "c"]);
+    // The stamps reflect each URI's stored value (not re-stamped).
+    expect(got.map((e) => e.stamp)).toEqual([100, 9, 1]);
+  });
+
+  it("readEntries orderBy: 'stamp' is the explicit equivalent of the default", async () => {
+    const store = new InMemoryUpdatesStore();
+    await store.saveEntries([
+      { signal: "x", uri: "b", stamp: 9 },
+      { signal: "x", uri: "a", stamp: 100 },
+      { signal: "x", uri: "c", stamp: 1 },
+    ]);
+    const def = await collect(store.readEntries({ signal: "x", since: 0 }));
+    const explicit = await collect(store.readEntries({ signal: "x", since: 0, orderBy: "stamp" }));
+    expect(explicit).toEqual(def);
+  });
+
+  it("readUpdatedEntries defaults to stamp-ascending order", async () => {
+    const store = new InMemoryUpdatesStore();
+    await store.saveEntries([
+      { signal: "src", uri: "z", stamp: 9 },
+      { signal: "src", uri: "a", stamp: 3 },
+      { signal: "src", uri: "m", stamp: 5 },
+    ]);
+    const got = await collect(
+      store.readUpdatedEntries({ upstreamSignal: "src", currentSignal: "ext" }),
+    );
+    expect(got.map((e) => e.uri)).toEqual(["a", "m", "z"]);
+    expect(got.map((e) => e.stamp)).toEqual([3, 5, 9]);
+  });
+
+  it("readUpdatedEntries with orderBy: 'uri' yields URI-ascending", async () => {
+    const store = new InMemoryUpdatesStore();
+    await store.saveEntries([
+      { signal: "src", uri: "z", stamp: 1 },
+      { signal: "src", uri: "a", stamp: 9 },
+      { signal: "src", uri: "m", stamp: 5 },
+    ]);
+    const got = await collect(
+      store.readUpdatedEntries({
+        upstreamSignal: "src",
+        currentSignal: "ext",
+        orderBy: "uri",
+      }),
+    );
+    expect(got.map((e) => e.uri)).toEqual(["a", "m", "z"]);
+  });
+
+  it("orderBy is independent of since/uriPrefix filtering — same entries, different order", async () => {
+    const store = new InMemoryUpdatesStore();
+    await store.saveEntries([
+      { signal: "x", uri: "/keep/3", stamp: 1 },
+      { signal: "x", uri: "/keep/1", stamp: 7 },
+      { signal: "x", uri: "/keep/2", stamp: 3 },
+      { signal: "x", uri: "/drop/9", stamp: 5 },
+    ]);
+    const byStamp = await collect(
+      store.readEntries({ signal: "x", since: 0, uriPrefix: "/keep/" }),
+    );
+    const byUri = await collect(
+      store.readEntries({
+        signal: "x",
+        since: 0,
+        uriPrefix: "/keep/",
+        orderBy: "uri",
+      }),
+    );
+    // Same set, different order.
+    expect(new Set(byStamp.map((e) => e.uri))).toEqual(new Set(byUri.map((e) => e.uri)));
+    expect(byStamp.map((e) => e.uri)).toEqual(["/keep/3", "/keep/2", "/keep/1"]);
+    expect(byUri.map((e) => e.uri)).toEqual(["/keep/1", "/keep/2", "/keep/3"]);
+  });
+});
+
 describe("InMemoryUpdatesStore — stamp validation", () => {
   it("rejects a NaN stamp instead of storing an unreadable entry", async () => {
     // A NaN stamp would round-trip into the store but be invisible to