@rangojs/router 0.17.0 → 0.18.0

package/src/ssr/ssr-root.tsx

@@ -189,14 +189,15 @@ export interface SsrRootOptions {
   nonce?: string;
   /**
    * Fires once when the Flight payload root settles (resolve OR reject) — the
-   * signal that every client-module load the payload references completed and
-   * fizz can start emitting the tree. The capture pass gates its abort on
-   * this: a fully REPLAYED (prerendered) route's Flight stream goes
-   * byte-quiet in ~1-3ms, but fizz cannot render even <html> until the
-   * module loads finish (real module-runner I/O in dev; 100ms+ on a cold
-   * graph), so an abort gated on Flight quiet alone fires first and freezes
-   * a zero-byte prelude. Masked-loader holes do NOT block this signal —
-   * they postpone below the root.
+   * signal that fizz can start emitting the tree. A client component used as
+   * an element type resolves lazily, so its module load can still be in
+   * flight (captureShellHTML waits for those separately, issue #949). The
+   * capture pass gates its abort on this: a fully REPLAYED (prerendered)
+   * route's Flight stream goes byte-quiet in ~1-3ms, but fizz cannot render
+   * even <html> until the module loads finish (real module-runner I/O in dev;
+   * 100ms+ on a cold graph), so an abort gated on Flight quiet alone fires
+   * first and freezes a zero-byte prelude. Masked-loader holes do NOT block
+   * this signal — they postpone below the root.
    */
   onPayloadSettled?: () => void;
   /**

package/src/testing/flight.entry.ts

@@ -18,6 +18,11 @@
  * element tree, so you can assert typed prop fidelity across the client boundary
  * (a `Date` comes back a `Date`) and detect inlined-vs-island. Serialize +
  * deserialize only — no hydration/interaction (that is the e2e tier).
+ *
+ * `serveShellRequest` serves one request through the router's production
+ * handler, so a `ppr` route's MISS runs the real shell capture and a later
+ * request is a real HIT; only the HTML step is stubbed. `resetShellTestState`
+ * clears the per-isolate PPR state between tests.
  */
 
 export {
@@ -50,3 +55,12 @@ export type {
   RenderHandlerResult,
   StateCookieSeed,
 } from "./render-handler.js";
+
+export {
+  serveShellRequest,
+  resetShellTestState,
+} from "./serve-shell-request.js";
+export type {
+  ServeShellRequestOptions,
+  ServeShellRequestResult,
+} from "./serve-shell-request.js";

package/src/testing/run-loader.ts

@@ -48,6 +48,9 @@ import { collectHandle } from "./collect-handle.js";
 import type { ThemeConfig } from "../theme/types.js";
 import type { SegmentCacheStore } from "../cache/types.js";
 import type { CacheProfile } from "../cache/profile-registry.js";
+import type { PartialCacheOptions } from "../types/cache-types.js";
+import type { HandlerContext } from "../types/handler-context.js";
+import { resolveLoaderData } from "../router/segment-resolution/loader-cache.js";
 import {
   createTestRequestContext,
   buildRunSnapshot,
@@ -111,9 +114,11 @@ export interface RunLoaderOptions<TEnv = any> {
   basename?: string;
   /**
    * Theme config in the same shape `createRouter({ theme })` takes (e.g. `true`
-   * or `{ themes: [...] }`). Seeds the request's theme config so nested handler
-   * or cache contexts created from this loader observe it. Loaders themselves do
-   * not expose `ctx.theme`/`ctx.setTheme` (those are handler/middleware-only).
+   * or `{ themes: [...] }`). Seeds the request's theme config, so the loader
+   * body reads the visitor's theme (the Cookie on `request`, or the default)
+   * with `getRequestContext().theme`. That read is guarded like `cookies()`: it
+   * throws inside a `"use cache"` function (#971). The loader ctx has no
+   * `theme`/`setTheme` (those are handler/middleware-only).
    */
   theme?: ThemeConfig | true;
   /** Environment bindings surfaced as `ctx.env`. */
@@ -151,6 +156,18 @@ export interface RunLoaderOptions<TEnv = any> {
   cacheStore?: SegmentCacheStore;
   /** Cache profiles (the `createRouter({ cacheProfiles })` shape). */
   cacheProfiles?: Record<string, CacheProfile>;
+  /**
+   * Run the loader through the `cache()` a route binds it with — the options
+   * `loader(Loader, () => [cache({...})])` takes — using the production
+   * read-through against `cache.store`, else `cacheStore`. A HIT returns the
+   * stored value without running the body and records the tags the body
+   * recorded when the entry was written, so `updateTag()` of any of them drops
+   * the entry. The write is a background task the call does not await (spy
+   * on the store's `setItem` to wait for it). A raw loader body is keyed by
+   * its function reference. `handlePushes` lists the body's pushes, so a HIT
+   * reports none: the replay into the page's handle store is not modeled.
+   */
+  cache?: PartialCacheOptions;
   /**
    * Customize the rango state cookie a loader that calls
    * `invalidateClientCache()` rotates (the name is always seeded — default
@@ -352,6 +369,47 @@ function runWithLoaderContext<R>(
   });
 }
 
+const rawLoaderIds = new WeakMap<object, string>();
+let rawLoaderSeq = 0;
+
+function cacheLoaderId(loader: RunnableLoader<unknown>): string {
+  if (typeof loader !== "function") return loader.$$id;
+  let id = rawLoaderIds.get(loader);
+  if (id === undefined) {
+    id = `runLoader#${++rawLoaderSeq}`;
+    rawLoaderIds.set(loader, id);
+  }
+  return id;
+}
+
+/**
+ * Call the loader: directly, or through its `cache()` binding (the funnel a
+ * route's DSL loaders take, loader-cache.ts resolveLoaderData), whose MISS
+ * reads the binding through the handler ctx's use().
+ */
+function invokeLoader<T>(
+  loader: RunnableLoader<T>,
+  loaderFn: (ctx: TestLoaderContext) => Promise<T> | T,
+  opts: RunLoaderOptions,
+  loaderCtx: TestLoaderContext,
+): Promise<T> {
+  const run = () => Promise.resolve(loaderFn(loaderCtx));
+  if (!opts.cache) return run();
+  const handlerCtx = {
+    params: loaderCtx.params,
+    use: run,
+  } as unknown as HandlerContext<any, any>;
+  return resolveLoaderData(
+    {
+      loader: { __brand: "loader", $$id: cacheLoaderId(loader) },
+      revalidate: [],
+      cache: { options: opts.cache },
+    },
+    handlerCtx,
+    loaderCtx.pathname,
+  );
+}
+
 export async function runLoader<T>(
   loader: RunnableLoader<T>,
   opts: RunLoaderOptions = {},
@@ -359,7 +417,7 @@ export async function runLoader<T>(
   const loaderFn = resolveLoaderFn(loader);
   const { ctx } = createTestRequestContext(buildLoaderCtxOpts(opts));
   return runWithLoaderContext(ctx as RequestContext<any>, opts, (loaderCtx) =>
-    Promise.resolve(loaderFn(loaderCtx)),
+    invokeLoader(loader, loaderFn, opts, loaderCtx),
   );
 }
 
@@ -413,7 +471,7 @@ export async function runLoaderResult<T>(
     result = await runWithLoaderContext(
       reqCtx,
       opts,
-      (loaderCtx) => Promise.resolve(loaderFn(loaderCtx)),
+      (loaderCtx) => invokeLoader(loader, loaderFn, opts, loaderCtx),
       handlePushes,
     );
   } catch (error) {

package/src/testing/serve-shell-request.ts

@@ -0,0 +1,378 @@
+/**
+ * serveShellRequest — one request through the production PPR serve path.
+ *
+ * Builds the router's request handler as `router.fetch` does
+ * (`createRSCHandler` with the router's `nonce`, `version` and `cache`
+ * config) and serves one GET through it, so middleware, the shell serve path
+ * and a MISS's background capture all run production code, on real Flight
+ * (the vendored react-server-dom serializer). The call settles every
+ * background task the request scheduled before it returns, so a MISS's
+ * `putShell` has landed.
+ *
+ * The one stub is the SSR module (the HTML step): `react-dom/server` refuses
+ * to load under the react-server condition this project runs with.
+ * - `renderHTML` passes the Flight stream through: a document body is its
+ *   Flight payload.
+ * - `captureShellHTML` returns, as the prelude, the Flight text the capture
+ *   rendered before `quiesce` (plus the post-quiesce task hops the real
+ *   capture takes). `postponed` is always null. A capture that does not
+ *   quiesce within what `ppr.captureTimeout` leaves after the handler output
+ *   settled returns no shell; the real one goes on to its abort and returns
+ *   whatever prelude rendered by then.
+ * - `resumeShellHTML` passes the tail's Flight stream through: a HIT body is
+ *   the stored prelude followed by the tail's Flight payload.
+ * So there is no HTML: the prelude's `<body` sanity gate, SSR render errors
+ * (production refuses the capture on one), fizz resume of the holes,
+ * bootstrap scripts and nonce injection stay e2e-only. A route whose real
+ * capture refuses for a root postpone (a live loader read with no boundary
+ * above it) or an SSR render error is captured here.
+ *
+ * Must run under the `react-server` condition (the rsc Vitest project), with
+ * `rangoTestAliases()` resolving `@vitejs/plugin-rsc/rsc/server` to the stub
+ * that carries the Flight runtime.
+ */
+
+import { AsyncLocalStorage } from "node:async_hooks";
+import { toInternal, type Rango } from "../router/router-interfaces.js";
+import { ensureRouterManifest } from "../route-map-builder.js";
+import type { SegmentCacheStore, ShellCacheEntry } from "../cache/types.js";
+import type { ExecutionContext } from "../types/request-scope.js";
+import type { HandlerCacheConfig, SSRModule } from "../rsc/types.js";
+import type { createRSCHandler } from "../rsc/handler.js";
+import {
+  POST_QUIESCE_TASK_HOPS,
+  SHELL_CAPTURE_MAX_WAIT_MS,
+} from "../rsc/shell-capture-constants.js";
+import { SEGMENT_FRAGMENT_CAPABILITY_HEADER } from "../segment-fragments.js";
+import { _getRequestContext } from "../server/request-context.js";
+import {
+  parsePprReplayStatus,
+  parseShellStatus,
+  shellCacheKey,
+  type PprReplayStatus,
+  type ShellStatus,
+} from "./shell-status.js";
+
+/** Options for {@link serveShellRequest}. */
+export interface ServeShellRequestOptions<TEnv = any> {
+  /**
+   * The shell and segment store for this request, in place of the router's
+   * `createRouter({ cache })` store. The rest of that config (`enabled`,
+   * `searchParams`) still applies. Omit it to serve through the router's own
+   * store.
+   */
+  cacheStore?: SegmentCacheStore;
+  /** Environment bindings, as `router.fetch(request, { env })` passes them. */
+  env?: TEnv;
+  /**
+   * Request headers. A document request defaults `accept` to `text/html`.
+   */
+  headers?: HeadersInit;
+  /**
+   * Serve the client navigation the browser sends (`?_rsc_partial`) instead
+   * of a document GET: from `from` (default `/`) with `segments` mounted
+   * (default none), advertising fragment passthrough. The replay decision is
+   * `result.replayStatus`.
+   */
+  partial?: true | { from?: string; segments?: readonly string[] };
+}
+
+/** Result of {@link serveShellRequest}. */
+export interface ServeShellRequestResult {
+  /** The handler's Response. Its body is already read: see `body`. */
+  response: Response;
+  /** The response body text. */
+  body: string;
+  /** `x-rango-shell` (`HIT` | `MISS`), or null when the header is absent. */
+  shellStatus: ShellStatus | null;
+  /** `x-rango-ppr-replay` on a partial request, or null. */
+  replayStatus: PprReplayStatus | null;
+  /**
+   * The prelude a HIT served: the capture's Flight text at quiesce (the
+   * shell as captured). Undefined unless `shellStatus` is `HIT`.
+   */
+  prelude: string | undefined;
+  /**
+   * The Flight payload this request rendered: a HIT's tail, a document's
+   * render, or a Flight response. Undefined when no Flight was rendered (a
+   * redirect or a middleware response).
+   */
+  flight: string | undefined;
+  /**
+   * The shell key production resolved for this request: for a document the
+   * serve path read and captured (a MISS or a HIT), the key it used, request
+   * partition included (a route partitioned by `cache({ key })` or the
+   * store's `keyGenerator`). Otherwise (a route without `ppr`, a request
+   * passed to axis 1 before its key was resolved, a partial request) the
+   * document URL's key without a partition, as `shellCacheKey` builds it.
+   */
+  key: string;
+  /**
+   * Read the document shell entry under `key` from the request's store
+   * (a passive `getShell`), or null. A read, so on a store with a shell memo
+   * (CFCacheStore, VercelCacheStore) it warms the memo like any other read.
+   */
+  readEntry(): Promise<ShellCacheEntry | null>;
+}
+
+type ShellHandler = ReturnType<typeof createRSCHandler>;
+
+/** What one call's cache resolution and SSR step saw. */
+interface Recorder {
+  cache?: HandlerCacheConfig;
+  /** The shell key the serve path resolved (RequestContext._shellKey). */
+  shellKey?: string;
+  /** A document render reached the HTML step. */
+  rendered?: boolean;
+  /** The HIT tail's Flight text. */
+  tail?: string;
+}
+
+const recorders = new AsyncLocalStorage<Recorder>();
+
+function macrotask(): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, 0));
+}
+
+/**
+ * Record the shell key the serve path resolved: every SSR step of a document
+ * the serve path read runs inside its request context (a capture's derived
+ * context inherits it).
+ */
+function recordShellKey(recorder: Recorder | undefined): void {
+  const shellKey = _getRequestContext()?._shellKey;
+  if (recorder && shellKey !== undefined) recorder.shellKey ??= shellKey;
+}
+
+const SSR_STUB: SSRModule = {
+  async renderHTML(rscStream) {
+    const recorder = recorders.getStore();
+    if (recorder) recorder.rendered = true;
+    recordShellKey(recorder);
+    return rscStream;
+  },
+  async captureShellHTML(rscStream, options) {
+    const reader = rscStream.getReader();
+    const decoder = new TextDecoder();
+    let frozen = "";
+    // The capture gate freezes the stream at quiesce without closing it.
+    void (async () => {
+      for (;;) {
+        const { done, value } = await reader.read();
+        if (done) return;
+        frozen += decoder.decode(value, { stream: true });
+      }
+    })().catch(() => {});
+    let timer: ReturnType<typeof setTimeout> | undefined;
+    const deadline = new Promise<false>((resolve) => {
+      timer = setTimeout(
+        () => resolve(false),
+        options.maxWaitMs ?? SHELL_CAPTURE_MAX_WAIT_MS,
+      );
+    });
+    try {
+      const quiet = await Promise.race([
+        options.quiesce.then(() => true as const),
+        deadline,
+      ]);
+      if (!quiet) return null;
+    } finally {
+      clearTimeout(timer);
+    }
+    for (let i = 0; i < POST_QUIESCE_TASK_HOPS; i++) await macrotask();
+    return { prelude: new TextEncoder().encode(frozen), postponed: null };
+  },
+  async resumeShellHTML(rscStream) {
+    const recorder = recorders.getStore();
+    recordShellKey(recorder);
+    // The served HIT's tail only, not a background re-render's.
+    if (!recorder || recorder.tail !== undefined) return rscStream;
+    recorder.tail = "";
+    const decoder = new TextDecoder();
+    return rscStream.pipeThrough(
+      new TransformStream<Uint8Array, Uint8Array>({
+        transform(chunk, controller) {
+          recorder.tail += decoder.decode(chunk, { stream: true });
+          controller.enqueue(chunk);
+        },
+        flush() {
+          recorder.tail += decoder.decode();
+        },
+      }),
+    );
+  },
+};
+
+/** Handlers per router, keyed by the `cacheStore` override or the router. */
+const handlers = new WeakMap<object, WeakMap<object, ShellHandler>>();
+
+async function getHandler(
+  router: Rango<any, any>,
+  cacheStore: SegmentCacheStore | undefined,
+): Promise<ShellHandler> {
+  let byStore = handlers.get(router);
+  if (!byStore) handlers.set(router, (byStore = new WeakMap()));
+  const cached = byStore.get(cacheStore ?? router);
+  if (cached) return cached;
+
+  const internal = toInternal(router);
+  const routerCache = internal.cache;
+  // Resolved per request, as the handler resolves the router's own config.
+  const cache = (env: unknown, ctx?: ExecutionContext): HandlerCacheConfig => {
+    const base =
+      typeof routerCache === "function" ? routerCache(env, ctx) : routerCache;
+    const config = cacheStore ? { ...base, store: cacheStore } : base!;
+    const recorder = recorders.getStore();
+    if (recorder) recorder.cache = config;
+    return config;
+  };
+  const { createRSCHandler } = await import("../rsc/handler.js");
+  const handler = createRSCHandler({
+    router: internal,
+    nonce: internal.nonce,
+    version: internal.version,
+    loadSSRModule: async () => SSR_STUB,
+    cache: routerCache || cacheStore ? cache : undefined,
+  });
+  byStore.set(cacheStore ?? router, handler);
+  return handler;
+}
+
+function buildRequest(url: URL, options: ServeShellRequestOptions): Request {
+  const headers = new Headers(options.headers);
+  const partial = options.partial;
+  if (!partial) {
+    if (!headers.has("accept")) headers.set("accept", "text/html");
+    return new Request(url, { headers });
+  }
+  const { from = "/", segments = [] } = partial === true ? {} : partial;
+  const target = new URL(url);
+  target.searchParams.set("_rsc_partial", "true");
+  target.searchParams.set("_rsc_segments", segments.join(","));
+  if (!headers.has("X-RSC-Router-Client-Path")) {
+    headers.set("X-RSC-Router-Client-Path", new URL(from, url).href);
+  }
+  if (!headers.has(SEGMENT_FRAGMENT_CAPABILITY_HEADER)) {
+    headers.set(SEGMENT_FRAGMENT_CAPABILITY_HEADER, "1");
+  }
+  return new Request(target, { headers });
+}
+
+/**
+ * Reset the per-isolate PPR state requests leave behind, so the next test
+ * starts as a fresh worker would: the capture's stampede guard, its backoff
+ * (a refused capture backs its URL off for later tests too), the capture's
+ * and the serve path's once-per-key warnings, the build-shell manifest memo,
+ * and CFCacheStore's isolate memos (shells, tag markers, tag hints), which
+ * every CFCacheStore in the process shares by namespace and URL. Call it in
+ * `beforeEach`, never while a request is in flight. VercelCacheStore's memos
+ * live on the `cache` handle it is given: a new handle starts empty.
+ */
+export async function resetShellTestState(): Promise<void> {
+  const [capture, serve, buildShells, cf] = await Promise.all([
+    import("../rsc/shell-capture.js"),
+    import("../rsc/shell-serve.js"),
+    import("../rsc/shell-build-manifest.js"),
+    import("../cache/cf/cf-cache-store.js"),
+  ]);
+  capture.resetShellCaptureStateForTests();
+  serve.resetShellServeStateForTests();
+  buildShells.resetBuildShellManifestForTests();
+  cf.resetCFShellMemoForTests();
+}
+
+/** Settle background tasks, including ones scheduled while settling. */
+async function settle(tasks: Promise<unknown>[]): Promise<void> {
+  do {
+    while (tasks.length > 0) await Promise.allSettled(tasks.splice(0));
+    await macrotask();
+  } while (tasks.length > 0);
+}
+
+/**
+ * Serve one GET for `url` through the router's production request handler
+ * and settle the background tasks it scheduled (a MISS's capture and
+ * `putShell` included). The first request for a `ppr` route is a MISS; the
+ * next one with the same store is a HIT from the stored shell. The HTML step
+ * is stubbed (see the module header), so `prelude` and `flight` are Flight
+ * text. The request starts in a later millisecond than the call: a store
+ * refuses a capture that starts in the millisecond of a tag invalidation, so
+ * an `updateTag()` made just before the call never refuses its recapture.
+ * Rejects when the handler throws or the response stream errors (a failing
+ * HIT tail), after settling the tasks scheduled so far.
+ *
+ * @example
+ * const store = new MemorySegmentCacheStore();
+ * const miss = await serveShellRequest(router, "/product/1", { cacheStore: store });
+ * expect(miss.shellStatus).toBe("MISS");
+ * const hit = await serveShellRequest(router, "/product/1", { cacheStore: store });
+ * expect(hit.shellStatus).toBe("HIT");
+ * expect(hit.prelude).toContain("Widget"); // captured shell
+ * expect(hit.flight).toContain("in stock"); // live hole, rendered now
+ */
+export async function serveShellRequest<TEnv = any>(
+  router: Rango<TEnv, any>,
+  url: string | URL,
+  options: ServeShellRequestOptions<TEnv> = {},
+): Promise<ServeShellRequestResult> {
+  const called = Date.now();
+  const target = new URL(url, "http://localhost");
+  await ensureRouterManifest(router.id);
+  const handler = await getHandler(router, options.cacheStore);
+  const request = buildRequest(target, options);
+  const tasks: Promise<unknown>[] = [];
+  const executionContext: ExecutionContext = {
+    waitUntil(promise) {
+      tasks.push(
+        Promise.resolve(promise).catch((error) =>
+          console.error("[waitUntil] Background task failed:", error),
+        ),
+      );
+    },
+    passThroughOnException() {},
+  };
+  const recorder: Recorder = {};
+  while (Date.now() === called) await macrotask();
+
+  const { response, body } = await recorders.run(recorder, async () => {
+    try {
+      const response = await handler(request, {
+        env: options.env,
+        ctx: executionContext,
+      });
+      return { response, body: await response.text() };
+    } finally {
+      await settle(tasks);
+    }
+  });
+
+  const config =
+    recorder.cache && recorder.cache.enabled !== false
+      ? recorder.cache
+      : undefined;
+  const key = recorder.shellKey ?? shellCacheKey(target, config?.searchParams);
+  const tail = recorder.tail;
+  const isFlight =
+    recorder.rendered ||
+    response.headers.get("content-type")?.includes("text/x-component");
+  return {
+    response,
+    body,
+    shellStatus: parseShellStatus(response),
+    replayStatus: parsePprReplayStatus(response),
+    prelude:
+      tail !== undefined && body.endsWith(tail)
+        ? body.slice(0, body.length - tail.length)
+        : undefined,
+    flight: tail ?? (isFlight && body ? body : undefined),
+    key,
+    async readEntry() {
+      // Passive, like the serve path's own reads: a stale entry's SWR
+      // revalidation stays unclaimed for the next request.
+      const read = await config?.store.getShell?.(key, {
+        claimRevalidation: false,
+      });
+      return read?.entry ?? null;
+    },
+  };
+}

package/src/testing/shell-status.ts

@@ -18,17 +18,25 @@
  *    Secondary unit signal: `store.getShell(shellCacheKey(url))` after a real
  *    capture flush (background `putShell`). There is no Flight flag for shell HIT.
  *
- * 3. **Stays e2e-only:** live MISS → background capture → HIT (needs real
- *    `captureShellHTML` / React `prerender` + resume fizz), browser resume of
- *    holes, Cloudflare/Vercel live stores, nonce/eligibility fall-open, and
- *    build-time producer B bake. Unit layer dogfoods the **store family + key
- *    identity + header assert**; it must never stub a shell HIT response.
+ * 3. **A real MISS → capture → HIT in unit tests:** `serveShellRequest`
+ *    (`@rangojs/router/testing/flight`, serve-shell-request.ts) serves one
+ *    request through the router's production handler and settles its
+ *    background capture; a later request with the same store is a real HIT or
+ *    partial replay. Its only stub is the HTML step (the SSR module).
+ *    **Stays e2e-only:** real HTML (the prelude's `<body` gate, fizz resume of
+ *    holes), browser resume of holes, live Cloudflare/Vercel runtimes (store
+ *    test doubles work in unit), and build-time producer B bake. The unit layer
+ *    must never stub a shell HIT response.
  *
  * Import from `@rangojs/router/testing` (Vitest) or `@rangojs/router/testing/e2e`
  * (Playwright — same pure helpers, no Vite virtuals).
  */
 
-import { sortedSearchString } from "../cache/cache-key-utils.js";
+import {
+  composeCacheKeys,
+  sortedSearchString,
+} from "../cache/cache-key-utils.js";
+import { partitionShellKey } from "../rsc/shell-capture-constants.js";
 import {
   compileSearchParamsFilter,
   type CacheSearchParams,
@@ -57,8 +65,6 @@ const PPR_REPLAY_BYPASS_REASONS = [
   "no-entry",
   "invalid-version",
   "corrupt-entry",
-  "handler-live-holes",
-  "transition-when",
   "no-segment-snapshot",
   "snapshot-miss",
   "explicit-cache-hit",
@@ -79,17 +85,27 @@ export type ShellStatusTarget = Response | { headers: Headers };
 /**
  * Shell store key for a document URL — same formula as production
  * `buildShellKey` in `rsc/shell-serve.ts` (host + pathname + sorted search +
- * `:shell`). Kept here so the testing barrel never imports the shell-serve
- * module (it pulls React). A parity test pins the two implementations.
+ * `:shell`), then the production `partitionShellKey` itself (the leaf
+ * `rsc/shell-capture-constants.ts`). The base is kept here so the testing
+ * barrel never imports the shell-serve module (it pulls React). A parity test
+ * pins the two implementations.
  *
  * Accepts a `URL` or an absolute/relative request URL string (relative strings
  * resolve against `http://localhost`). Pass the router's `cache.searchParams`
  * config as `searchParams` when the app under test sets one — the production
- * key applies it, so the expected key must too.
+ * key applies it, so the expected key must too. When the route's record is
+ * partitioned by the request — its `cache({ key })`, or the store's
+ * `keyGenerator` — its shell is too: pass the key that function returns for
+ * the request as `partition` (appended URI-encoded; a result equal to the
+ * default key partitions nothing, so pass none). Under nested keyed
+ * `cache()` scopes, pass their `key()` results as an array, outermost first:
+ * they compose as the production record key composes them
+ * (`composeCacheKeys`).
  */
 export function shellCacheKey(
   url: URL | string,
   searchParams?: CacheSearchParams,
+  partition?: string | readonly string[],
 ): string {
   const resolved =
     typeof url === "string" ? new URL(url, "http://localhost") : url;
@@ -98,7 +114,13 @@ export function shellCacheKey(
     compileSearchParamsFilter(searchParams),
   );
   const searchSuffix = sorted ? `?${sorted}` : "";
-  return `${resolved.host}${resolved.pathname}${searchSuffix}:shell`;
+  const key = `${resolved.host}${resolved.pathname}${searchSuffix}:shell`;
+  if (partition === undefined) return key;
+  if (typeof partition === "string") return partitionShellKey(key, partition);
+  // No key() results: no partition, as when `partition` is omitted.
+  return partition.length === 0
+    ? key
+    : partitionShellKey(key, composeCacheKeys(partition));
 }
 
 function getHeaders(target: ShellStatusTarget): Headers {

package/src/theme/ThemeProvider.tsx

@@ -133,12 +133,13 @@ export function ThemeProvider({
     setSystemTheme(getSystemTheme());
     // Re-sync state from an EXPLICITLY stored theme after mount. initialTheme
     // comes from the payload and can legitimately differ from the visitor's
-    // stored theme — on a PPR shell HIT it is deliberately the CAPTURE's theme
-    // (the resume tree must match the frozen prelude; see
-    // ShellCacheEntry.initialTheme). The FOUC script already applied the stored
-    // theme to the document pre-paint; this brings the provider state (toggles,
-    // useTheme readers) in line with it, and is the ONLY place the provider
-    // reads storage (the initializer must not — see the parity note above).
+    // stored theme — on a PPR shell HIT it is deliberately the CAPTURE's
+    // initialTheme, the no-cookie default (the resume tree must match the
+    // frozen prelude; see ShellCacheEntry.initialTheme). The FOUC script
+    // already applied the stored theme to the document pre-paint; this brings
+    // the provider state (toggles, useTheme readers) in line with it, and is
+    // the ONLY place the provider reads storage (the initializer must not —
+    // see the parity note above).
     // Only an explicit cookie/localStorage value re-syncs — a defaultTheme
     // fallback must not override a server-provided initialTheme when the
     // visitor never chose a theme.