@rangojs/router 0.0.0-experimental.8 → 0.0.0-experimental.80daa23e

package/src/server/handle-store.ts

@@ -14,9 +14,37 @@
 export type HandleData = Record<string, Record<string, unknown[]>>;
 
 /**
- * Deep clone handle data to create a snapshot.
- * @internal
+ * Build a HandleData snapshot from a HandleStore using segment ordering.
+ * Reads data directly from the store for each segment in order.
  */
+export function buildHandleSnapshot(
+  handleStore: HandleStore,
+  segmentOrder: string[],
+): HandleData {
+  const data: HandleData = {};
+  for (const segmentId of segmentOrder) {
+    const segData = handleStore.getDataForSegment(segmentId);
+    for (const handleName in segData) {
+      if (!data[handleName]) data[handleName] = {};
+      data[handleName][segmentId] = segData[handleName];
+    }
+  }
+  return data;
+}
+
+function createLateHandlePushError(
+  handleName: string,
+  segmentId: string,
+): Error {
+  const error = new Error(
+    `Handle "${handleName}" for segment "${segmentId}" was pushed after handle collection completed. ` +
+      `This usually means an async JSX subtree suspended and later tried to push a handle during streaming. ` +
+      `Push handles from the route/layout handler or during the initial synchronous JSX render instead.`,
+  );
+  error.name = "LateHandlePushError";
+  return error;
+}
+
 function cloneHandleData(data: HandleData): HandleData {
   const clone: HandleData = {};
   for (const handleName in data) {
@@ -44,11 +72,46 @@ export interface HandleStore {
   track<T>(promise: Promise<T>): Promise<T>;
 
   /**
-   * Promise that resolves when all tracked handlers have settled.
-   * Does not reject - uses Promise.allSettled internally.
+   * Signal that no more track() calls will be made.
+   * settled will not resolve until seal() is called AND all tracked
+   * promises have settled. Calling stream() or getData() auto-seals.
+   */
+  seal(): void;
+
+  /**
+   * Promise that resolves when the store is sealed AND all tracked
+   * handlers have settled.
    */
   readonly settled: Promise<void>;
 
+  /**
+   * Optional error callback for late streaming-handle failures.
+   * Called when push() throws LateHandlePushError (handle pushed after
+   * stream completion). Allows the router to surface these errors
+   * to onError and telemetry.
+   */
+  onError?: (error: Error) => void;
+
+  /**
+   * Track a promise on the AUXILIARY (loader) settlement lane. Aux tracking
+   * keeps the store OPEN for pushes (a mid-body `ctx.handle(H)(...)` from a
+   * streaming loader is legal until this lane drains) WITHOUT joining
+   * `settled` — the handler barrier that rendered()-readers, the document
+   * handle snapshot, and prerender/cache collection wait on. Two lanes exist
+   * because loader bodies stream: joining them to `settled` would block SSR
+   * markup and hydration on the slowest loader (ssr-root.tsx and
+   * rsc-router.tsx both drain the document handle stream in blocking
+   * positions), and would self-deadlock any loader awaiting ctx.rendered().
+   */
+  trackAuxiliary<T>(promise: Promise<T>): Promise<T>;
+
+  /**
+   * Promise that resolves when the store is sealed AND BOTH lanes (handler +
+   * auxiliary) have drained. Late-push legality and default stream completion
+   * key on this, not on `settled`.
+   */
+  readonly fullySettled: Promise<void>;
+
   /**
    * Push handle data for a specific handle and segment.
    * Multiple pushes to the same handle/segment accumulate in an array.
@@ -58,18 +121,32 @@ export interface HandleStore {
 
   /**
    * Get all collected handle data after all handlers have settled.
-   * Returns a promise that waits for `settled`, then returns the data.
-   * The data may contain unresolved promises which RSC will stream.
-   * @deprecated Use stream() for progressive updates
+   * Waits for `settled` (handler lane), then returns the data at that point.
    */
   getData(): Promise<HandleData>;
 
   /**
    * Get an async iterator that yields handle data on each push.
-   * The iterator completes when all handlers have settled.
+   * Completes at the chosen settlement barrier: "fullySettled" (default —
+   * nav/action payloads, whose consumer applies each yield progressively) or
+   * "settled" (the document lane, whose consumers drain to completion in
+   * blocking positions and must not wait on loader bodies).
    * Each yield contains the full accumulated state (not just the delta).
+   * Safe for concurrent consumers (per-consumer version cursor).
+   */
+  stream(
+    until?: "settled" | "fullySettled",
+  ): AsyncGenerator<HandleData, void, unknown>;
+
+  /**
+   * Document-lane late channel: yields full-state updates for pushes that
+   * land AFTER the handler barrier (streaming loader bodies), completing at
+   * fullySettled. Returns without yielding when the auxiliary lane is empty
+   * by the time `settled` resolves. The client consumes this post-hydration
+   * (rsc-router.tsx) and merges via the same application path as nav-lane
+   * progressive handle updates.
    */
-  stream(): AsyncGenerator<HandleData, void, unknown>;
+  streamLate(): AsyncGenerator<HandleData, void, unknown>;
 
   /**
    * Get handle data for a specific segment (for caching).
@@ -81,7 +158,10 @@ export interface HandleStore {
    * Replay cached handle data back into the store (for cache hits).
    * Used to restore handle data when serving cached segments.
    */
-  replaySegmentData(segmentId: string, segmentHandles: Record<string, unknown[]>): void;
+  replaySegmentData(
+    segmentId: string,
+    segmentHandles: Record<string, unknown[]>,
+  ): void;
 }
 
 /**
@@ -105,47 +185,148 @@ export interface HandleStore {
  * ```
  */
 export function createHandleStore(): HandleStore {
-  const pending: Promise<unknown>[] = [];
   const data: HandleData = {};
 
-  // Queue for pending emissions and resolver for waiting consumer
-  let pendingEmissions: HandleData[] = [];
-  let emissionResolver: (() => void) | null = null;
+  // Settlement barriers: `settled` (handler lane) resolves when sealed AND
+  // handler inflight === 0. `fullySettled` additionally waits for the
+  // auxiliary (loader) lane. seal() signals "no more track() calls"; each
+  // track/trackAuxiliary increments its lane's count, each promise settle
+  // decrements. Barriers resolve once their conditions are met — even if
+  // tracks are added while earlier ones are still in flight.
+  let sealed = false;
+  let inflightCount = 0;
+  let auxInflightCount = 0;
+  let drainWaiters: (() => void)[] = [];
+  let fullDrainWaiters: (() => void)[] = [];
+
+  // Swap the waiter list out before resolving: a resolver may synchronously
+  // re-await and push onto the fresh list, which must not be drained by this
+  // pass. All three waiter lists (drain, full-drain, emission) flush through
+  // here so that ordering rule cannot drift between them.
+  function flushWaiters(waiters: (() => void)[]): void {
+    for (const resolve of waiters) resolve();
+  }
+
+  // Both lanes settle identically — the handler lane (`track`) gates `settled`,
+  // the auxiliary loader lane (`trackAuxiliary`) gates `fullySettled`. Shared so
+  // neither copy can lose this rule: .then() rather than .finally(), because
+  // .finally() re-throws on a NEW branch, producing an unhandled rejection that
+  // can crash the process when the tracked promise rejects.
+  function trackSettlement<T>(
+    promise: Promise<T>,
+    release: () => void,
+  ): Promise<T> {
+    const onSettle = () => {
+      release();
+      notifyDrain();
+    };
+    promise.then(onSettle, onSettle);
+    return promise;
+  }
+
+  function notifyDrain() {
+    if (sealed && inflightCount === 0 && drainWaiters.length > 0) {
+      const waiters = drainWaiters;
+      drainWaiters = [];
+      flushWaiters(waiters);
+    }
+    if (sealed && inflightCount === 0 && auxInflightCount === 0) {
+      if (streamConsumed && !completed) {
+        // Late-push guard + stream termination key on FULL drain, so a
+        // streaming loader body's ctx.handle() push stays legal until the
+        // auxiliary lane empties. Armed only once a stream consumer exists —
+        // getData()-only paths (prerender collection) keep their historical
+        // no-guard behavior.
+        completed = true;
+        signalEmission();
+      }
+      if (fullDrainWaiters.length > 0) {
+        const waiters = fullDrainWaiters;
+        fullDrainWaiters = [];
+        flushWaiters(waiters);
+      }
+    }
+  }
+
+  function sealInternal() {
+    if (sealed) return;
+    sealed = true;
+    notifyDrain();
+  }
+
+  // Emission versioning. Each push bumps `version`; every stream consumer
+  // keeps its own cursor and yields whenever the version advanced past it.
+  // Multi-consumer safe (the document "settled" stream and the late channel
+  // overlap until the handler barrier): a single dirty bit + single resolver
+  // slot would drop wakeups for all but one consumer.
+  let version = 0;
+  let emissionWaiters: (() => void)[] = [];
   let completed = false;
+  let streamConsumed = false;
 
-  // Signal that a new emission is available
+  // Wake every waiting consumer (new push or a settlement barrier fired).
   function signalEmission() {
-    if (emissionResolver) {
-      const resolver = emissionResolver;
-      emissionResolver = null;
-      resolver();
+    if (emissionWaiters.length > 0) {
+      const waiters = emissionWaiters;
+      emissionWaiters = [];
+      flushWaiters(waiters);
     }
   }
 
   // Wait for the next emission or completion
-  function waitForEmission(): Promise<void> {
-    if (pendingEmissions.length > 0 || completed) {
+  // Resolve when the consumer's cursor is behind the current version, the
+  // given done-flag reader reports termination, or a new signal arrives.
+  function waitForEmission(seen: number, isDone: () => boolean): Promise<void> {
+    if (version > seen || isDone()) {
       return Promise.resolve();
     }
     return new Promise((resolve) => {
-      emissionResolver = resolve;
+      emissionWaiters.push(resolve);
     });
   }
 
   return {
     track<T>(promise: Promise<T>): Promise<T> {
-      pending.push(promise);
-      return promise;
+      inflightCount++;
+      return trackSettlement(promise, () => {
+        inflightCount--;
+      });
+    },
+
+    trackAuxiliary<T>(promise: Promise<T>): Promise<T> {
+      auxInflightCount++;
+      return trackSettlement(promise, () => {
+        auxInflightCount--;
+      });
+    },
+
+    seal() {
+      sealInternal();
     },
 
     get settled(): Promise<void> {
-      if (pending.length === 0) {
+      if (sealed && inflightCount === 0) return Promise.resolve();
+      return new Promise<void>((resolve) => {
+        drainWaiters.push(resolve);
+      });
+    },
+
+    get fullySettled(): Promise<void> {
+      if (sealed && inflightCount === 0 && auxInflightCount === 0) {
         return Promise.resolve();
       }
-      return Promise.allSettled(pending).then(() => {});
+      return new Promise<void>((resolve) => {
+        fullDrainWaiters.push(resolve);
+      });
     },
 
     push(handleName: string, segmentId: string, value: unknown): void {
+      if (completed) {
+        const error = createLateHandlePushError(handleName, segmentId);
+        if (this.onError) this.onError(error);
+        throw error;
+      }
+
       if (!data[handleName]) {
         data[handleName] = {};
       }
@@ -154,50 +335,85 @@ export function createHandleStore(): HandleStore {
       }
       data[handleName][segmentId].push(value);
 
-      // Queue a snapshot for emission
-      pendingEmissions.push(cloneHandleData(data));
+      // Bump the version; each consumer's cursor decides when to clone+yield.
+      version++;
       signalEmission();
     },
 
     getData(): Promise<HandleData> {
-      return this.settled.then(() => data);
+      sealInternal();
+      return this.settled.then(() => cloneHandleData(data));
     },
 
-    async *stream(): AsyncGenerator<HandleData, void, unknown> {
-      // Set up completion handler
-      this.settled.then(() => {
-        completed = true;
+    async *stream(
+      until: "settled" | "fullySettled" = "fullySettled",
+    ): AsyncGenerator<HandleData, void, unknown> {
+      streamConsumed = true;
+      sealInternal();
+      // Re-evaluate: the store may have been sealed and drained BEFORE this
+      // consumer existed (seal() then stream()); sealInternal early-returns
+      // then, so the completed flag must be armed here.
+      notifyDrain();
+
+      // Per-consumer termination flag driven by the chosen barrier. The
+      // global `completed` (late-push guard) always keys on FULL drain via
+      // notifyDrain — a "settled"-scoped consumer finishing does not make
+      // later loader pushes illegal.
+      let done = false;
+      const barrier = until === "settled" ? this.settled : this.fullySettled;
+      barrier.then(() => {
+        done = true;
         signalEmission();
       });
 
-      // Initial small delay to batch rapid synchronous pushes
-      // This allows multiple handles pushing in quick succession to be batched
+      // Batch rapid synchronous pushes with initial delay
       await new Promise((resolve) => setTimeout(resolve, 0));
 
-      // If we already have data, yield the accumulated state
-      if (Object.keys(data).length > 0) {
-        // Clear pending emissions since we're yielding current state
-        pendingEmissions = [];
-        yield cloneHandleData(data);
+      let seen = 0;
+      while (true) {
+        if (version > seen && Object.keys(data).length > 0) {
+          seen = version;
+          yield cloneHandleData(data);
+        }
+        if (done) break;
+        await waitForEmission(seen, () => done);
       }
+    },
 
-      // Continue streaming on each push
-      while (!completed) {
-        await waitForEmission();
+    async *streamLate(): AsyncGenerator<HandleData, void, unknown> {
+      streamConsumed = true;
+      sealInternal();
+      // Re-evaluate: the store may have been sealed and drained BEFORE this
+      // consumer existed (seal() then stream()); sealInternal early-returns
+      // then, so the completed flag must be armed here.
+      notifyDrain();
 
-        // Yield all pending emissions (yield latest only)
-        if (pendingEmissions.length > 0) {
-          // Skip intermediate states, yield the latest
-          const latest = pendingEmissions[pendingEmissions.length - 1];
-          pendingEmissions = [];
-          yield latest;
-        }
+      // Skip everything the handler-barrier snapshot already carried; the
+      // document consumers (SSR seed + pre-hydration drain) deliver that
+      // state. This channel exists only for pushes that land AFTER `settled`
+      // — streaming loader bodies writing meta/breadcrumbs mid-body.
+      await this.settled;
+      let seen = version;
+      if (auxInflightCount === 0) {
+        // Auxiliary lane already drained: every loader push (if any) beat the
+        // handler barrier and rode the normal snapshot. Nothing late follows —
+        // handler pushes after settled are illegal by contract.
+        return;
       }
 
-      // Final yield only if there are pending emissions that weren't yielded
-      // (handles that pushed after our last yield but before completion)
-      if (pendingEmissions.length > 0) {
-        yield cloneHandleData(data);
+      let done = false;
+      this.fullySettled.then(() => {
+        done = true;
+        signalEmission();
+      });
+
+      while (true) {
+        if (version > seen) {
+          seen = version;
+          yield cloneHandleData(data);
+        }
+        if (done) break;
+        await waitForEmission(seen, () => done);
       }
     },
 
@@ -211,18 +427,20 @@ export function createHandleStore(): HandleStore {
       return result;
     },
 
-    replaySegmentData(segmentId: string, segmentHandles: Record<string, unknown[]>): void {
+    replaySegmentData(
+      segmentId: string,
+      segmentHandles: Record<string, unknown[]>,
+    ): void {
       for (const handleName in segmentHandles) {
         if (!data[handleName]) {
           data[handleName] = {};
         }
-        // Replace with replayed data (not append) to avoid handle bleeding between routes.
-        // When a cached segment is restored, its handles should replace any existing data
-        // for that segment, not accumulate on top of data from a different route.
+        // Replace (not append) to avoid handle bleeding between routes.
+        // Cached segment restoration should replace existing data for that
+        // segment, not accumulate on top of data from a different route.
         data[handleName][segmentId] = [...segmentHandles[handleName]];
       }
-      // Trigger emission for streaming
-      pendingEmissions.push(cloneHandleData(data));
+      version++;
       signalEmission();
     },
   };

package/src/server/loader-registry.ts

@@ -6,25 +6,15 @@
  */
 
 import type { LoaderFn } from "../types.js";
-import type { MiddlewareFn } from "../router/middleware.js";
-import { getFetchableLoader } from "../loader.rsc.js";
+import {
+  getFetchableLoader,
+  type LoaderRegistryEntry,
+} from "./fetchable-loader-store.js";
 
-interface RegisteredLoader {
-  fn: LoaderFn<any, any, any>;
-  middleware: MiddlewareFn[];
-}
-
-// Server-side cache - maps loader $$id to function and middleware
-// This is a CACHE populated by getLoaderLazy() when loaders are first accessed.
-// The source of truth is fetchableLoaderRegistry in loader.ts, which is populated
-// when createLoader() runs. This cache exists to:
-// 1. Avoid repeated lookups/imports for the same loader
-// 2. Support lazy loading in production (loaders imported on-demand)
-// 3. Provide a stable reference for the RSC handler
-const loaderRegistry = new Map<string, RegisteredLoader>();
+// Cache populated by getLoaderLazy() when loaders are first accessed.
+// Source of truth is fetchableLoaderRegistry in loader.ts (populated on createLoader).
+const loaderRegistry = new Map<string, LoaderRegistryEntry>();
 
-// Lazy import map - set by the loader manifest
-// Maps loader $$id to a function that imports the loader module
 type LazyLoaderImport = () => Promise<{ $$id: string }>;
 let lazyLoaderImports: Map<string, LazyLoaderImport> | null = null;
 
@@ -32,35 +22,11 @@ let lazyLoaderImports: Map<string, LazyLoaderImport> | null = null;
  * Set the lazy loader imports map (called by the loader manifest)
  */
 export function setLoaderImports(
-  imports: Record<string, LazyLoaderImport>
+  imports: Record<string, LazyLoaderImport>,
 ): void {
   lazyLoaderImports = new Map(Object.entries(imports));
 }
 
-/**
- * Register a fetchable loader by $$id
- * Called by createLoader when fetchable option is provided
- */
-export function registerLoader(
-  id: string,
-  fn: LoaderFn<any, any, any>,
-  middleware: MiddlewareFn[] = []
-): void {
-  if (loaderRegistry.has(id)) {
-    // Already registered (can happen during HMR)
-    return;
-  }
-  loaderRegistry.set(id, { fn, middleware });
-}
-
-/**
- * Get a registered loader by $$id (synchronous)
- * Returns undefined if loader is not registered
- */
-export function getLoader(id: string): RegisteredLoader | undefined {
-  return loaderRegistry.get(id);
-}
-
 /**
  * Get a loader by $$id, loading it lazily if needed
  * This is the primary method for the RSC handler to get loaders
@@ -69,81 +35,69 @@ export function getLoader(id: string): RegisteredLoader | undefined {
  * In dev: IDs are "filePath#exportName", resolved via dynamic import
  */
 export async function getLoaderLazy(
-  id: string
-): Promise<RegisteredLoader | undefined> {
-  // Check if already cached in main registry
-  const existing = loaderRegistry.get(id);
-  if (existing) {
-    return existing;
-  }
-
-  // Check the fetchable loader registry (populated by createLoader)
+  id: string,
+): Promise<LoaderRegistryEntry | undefined> {
+  // Check fetchableLoaderRegistry first — it's the source of truth.
+  // createLoader() updates it on HMR, ensuring fresh functions after file changes.
   const fetchable = getFetchableLoader(id);
   if (fetchable) {
-    // Cache in main registry for future requests
     loaderRegistry.set(id, fetchable);
     return fetchable;
   }
 
-  // Try to lazy load from the import map (production mode)
+  const existing = loaderRegistry.get(id);
+  if (existing) {
+    return existing;
+  }
+
   if (lazyLoaderImports && lazyLoaderImports.size > 0) {
     const lazyImport = lazyLoaderImports.get(id);
     if (lazyImport) {
-      try {
-        // Import the loader module - this triggers createLoader which registers fn
-        await lazyImport();
+      // A failed import is a real server breakage (broken transitive import,
+      // syntax error, throw in module top-level code), not a "loader not
+      // registered" case. Rethrow so the caller can return a 500 and route
+      // the failure through onError, instead of collapsing it to a 404.
+      await lazyImport();
 
-        // Now try to get from fetchable registry (createLoader registered it)
-        const registered = getFetchableLoader(id);
-        if (registered) {
-          loaderRegistry.set(id, registered);
-          return registered;
-        }
-      } catch (error) {
-        console.error(`[LoaderRegistry] Failed to load loader "${id}":`, error);
+      const registered = getFetchableLoader(id);
+      if (registered) {
+        loaderRegistry.set(id, registered);
+        return registered;
       }
     }
   }
 
-  // Dev mode fallback: parse the ID and use Vite's dynamic import
-  // ID format in dev: "src/path/to/file.ts#ExportName"
+  // The remaining dev fallback (parse the id as "src/path/file.ts#ExportName"
+  // and import it by path) only makes sense in dev, where ids ARE file paths
+  // and the dev loader manifest is intentionally empty. In production ids are
+  // hashed ("<hash>#ExportName") and every resolvable loader is reached above
+  // via the in-memory registry or the lazy import manifest. The hash is not a
+  // path, so a production fall-through would run import("/<hash>") and throw a
+  // misleading "No such module <hash>" 500 instead of reporting the loader as
+  // unregistered. Return undefined in production so a genuinely unknown loader
+  // is a clean 404 "not found in registry" from handleLoaderFetch.
+  if (process.env.NODE_ENV === "production") {
+    return undefined;
+  }
+
   const hashIndex = id.indexOf("#");
   if (hashIndex !== -1) {
     const filePath = id.slice(0, hashIndex);
 
-    try {
-      // In dev mode, Vite handles dynamic imports
-      // Just importing the module triggers createLoader which registers the fn
-      await import(/* @vite-ignore */ `/${filePath}`);
+    // Same as the lazy branch: a thrown import is a server error, not a
+    // not-found. Let it propagate to the caller for a 500 + onError.
+    await import(/* @vite-ignore */ `/${filePath}`);
 
-      // Now try to get from fetchable registry
-      const registered = getFetchableLoader(id);
-      if (registered) {
-        loaderRegistry.set(id, registered);
-        return registered;
-      }
-    } catch (error) {
-      console.error(`[LoaderRegistry] Failed to load loader "${id}":`, error);
+    const registered = getFetchableLoader(id);
+    if (registered) {
+      loaderRegistry.set(id, registered);
+      return registered;
     }
   }
 
   return undefined;
 }
 
-/**
- * Check if a loader is registered by $$id
- */
-export function hasLoader(id: string): boolean {
-  return loaderRegistry.has(id) || getFetchableLoader(id) !== undefined;
-}
-
-/**
- * Get all registered loader IDs (for debugging)
- */
-export function getRegisteredLoaderIds(): string[] {
-  return Array.from(loaderRegistry.keys());
-}
-
 /**
  * Register a loader by its $$id (injected by Vite plugin)
  * This is called during module loading to cache loaders
@@ -155,20 +109,17 @@ export function registerLoaderById(loader: {
   if (!loader.$$id) {
     return;
   }
-  if (loaderRegistry.has(loader.$$id)) {
-    // Already registered (can happen during HMR)
-    return;
-  }
-
-  // For fetchable loaders, fn is stored in the fetchable registry by $$id
   const fetchable = getFetchableLoader(loader.$$id);
   if (fetchable) {
     loaderRegistry.set(loader.$$id, fetchable);
     return;
   }
 
-  // Fall back to using fn from the loader object (non-fetchable loaders)
   if (loader.fn) {
-    loaderRegistry.set(loader.$$id, { fn: loader.fn, middleware: [] });
+    loaderRegistry.set(loader.$$id, {
+      fn: loader.fn,
+      middleware: [],
+      fetchable: false,
+    });
   }
 }