@rangojs/router 0.0.0-experimental.8 → 0.0.0-experimental.8a4d0430

package/src/server/handle-store.ts

@@ -13,6 +13,19 @@
  */
 export type HandleData = Record<string, Record<string, unknown[]>>;
 
+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;
+}
+
 /**
  * Deep clone handle data to create a snapshot.
  * @internal
@@ -44,11 +57,26 @@ 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;
+
   /**
    * Push handle data for a specific handle and segment.
    * Multiple pushes to the same handle/segment accumulate in an array.
@@ -58,9 +86,7 @@ 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`, then returns the finalized data.
    */
   getData(): Promise<HandleData>;
 
@@ -81,7 +107,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,9 +134,31 @@ export interface HandleStore {
  * ```
  */
 export function createHandleStore(): HandleStore {
-  const pending: Promise<unknown>[] = [];
   const data: HandleData = {};
 
+  // Settlement barrier: resolved only when sealed AND inflight === 0.
+  // seal() signals "no more track() calls". Each track() increments
+  // inflightCount, each promise.finally() decrements. settled resolves
+  // once both conditions are met — even if tracks are added while
+  // earlier ones are still in flight.
+  let sealed = false;
+  let inflightCount = 0;
+  let drainWaiters: (() => void)[] = [];
+
+  function notifyDrain() {
+    if (sealed && inflightCount === 0 && drainWaiters.length > 0) {
+      const waiters = drainWaiters;
+      drainWaiters = [];
+      for (const resolve of waiters) resolve();
+    }
+  }
+
+  function sealInternal() {
+    if (sealed) return;
+    sealed = true;
+    notifyDrain();
+  }
+
   // Queue for pending emissions and resolver for waiting consumer
   let pendingEmissions: HandleData[] = [];
   let emissionResolver: (() => void) | null = null;
@@ -134,18 +185,38 @@ export function createHandleStore(): HandleStore {
 
   return {
     track<T>(promise: Promise<T>): Promise<T> {
-      pending.push(promise);
+      inflightCount++;
+      // Use .then(onSettle, onSettle) instead of .finally() to avoid
+      // creating an unhandled rejection branch when the tracked promise
+      // rejects (e.g. error route handlers). .finally() re-throws the
+      // rejection on a new branch that nobody catches, which can crash
+      // the server process.
+      const onSettle = () => {
+        inflightCount--;
+        notifyDrain();
+      };
+      promise.then(onSettle, onSettle);
       return promise;
     },
 
+    seal() {
+      sealInternal();
+    },
+
     get settled(): Promise<void> {
-      if (pending.length === 0) {
-        return Promise.resolve();
-      }
-      return Promise.allSettled(pending).then(() => {});
+      if (sealed && inflightCount === 0) return Promise.resolve();
+      return new Promise<void>((resolve) => {
+        drainWaiters.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] = {};
       }
@@ -160,10 +231,14 @@ export function createHandleStore(): HandleStore {
     },
 
     getData(): Promise<HandleData> {
-      return this.settled.then(() => data);
+      sealInternal();
+      return this.settled.then(() => cloneHandleData(data));
     },
 
     async *stream(): AsyncGenerator<HandleData, void, unknown> {
+      // Auto-seal: stream() is called after all track() registrations.
+      sealInternal();
+
       // Set up completion handler
       this.settled.then(() => {
         completed = true;
@@ -178,7 +253,8 @@ export function createHandleStore(): HandleStore {
       if (Object.keys(data).length > 0) {
         // Clear pending emissions since we're yielding current state
         pendingEmissions = [];
-        yield cloneHandleData(data);
+        const snapshot = cloneHandleData(data);
+        yield snapshot;
       }
 
       // Continue streaming on each push
@@ -211,7 +287,10 @@ 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] = {};

package/src/server/loader-registry.ts

@@ -6,13 +6,10 @@
  */
 
 import type { LoaderFn } from "../types.js";
-import type { MiddlewareFn } from "../router/middleware.js";
-import { getFetchableLoader } from "../loader.rsc.js";
-
-interface RegisteredLoader {
-  fn: LoaderFn<any, any, any>;
-  middleware: MiddlewareFn[];
-}
+import {
+  getFetchableLoader,
+  type LoaderRegistryEntry,
+} from "./fetchable-loader-store.js";
 
 // Server-side cache - maps loader $$id to function and middleware
 // This is a CACHE populated by getLoaderLazy() when loaders are first accessed.
@@ -21,7 +18,7 @@ interface RegisteredLoader {
 // 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>();
+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
@@ -32,35 +29,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,8 +42,8 @@ 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> {
+  id: string,
+): Promise<LoaderRegistryEntry | undefined> {
   // Check if already cached in main registry
   const existing = loaderRegistry.get(id);
   if (existing) {
@@ -130,20 +103,6 @@ export async function getLoaderLazy(
   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,12 +114,8 @@ 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
+  // For fetchable loaders, fn is stored in the fetchable registry by $$id.
+  // Always re-check the fetchable registry so HMR picks up the new function.
   const fetchable = getFetchableLoader(loader.$$id);
   if (fetchable) {
     loaderRegistry.set(loader.$$id, fetchable);
@@ -169,6 +124,10 @@ export function registerLoaderById(loader: {
 
   // 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,
+    });
   }
 }