@rangojs/router 0.17.0 → 0.18.0

package/src/router/segment-resolution/loader-snapshot.ts

@@ -4,9 +4,11 @@
  *
  * A bake-lane loader (lane rule: see resolveLoaderData, loader-cache.ts)
  * executes during shell capture; its settled container bakes into the prelude
- * while nested promises postpone as holes. To keep a HIT's fresh
- * payload byte-identical to that frozen prelude, the capture pins the container
- * in the shell snapshot:
+ * while nested promises postpone as holes. To keep a HIT's fresh payload in
+ * agreement with that frozen prelude, the capture pins the container in the
+ * shell snapshot. Everything outside a hole is pinned; the one exception is a
+ * server-component element holding a promise, which is a hole as a whole (see
+ * elideLoaderContainer):
  *
  *   - elide:   deep-walk the settled container; a PENDING nested promise is a
  *              hole, replaced by {@link LOADER_HOLE_KEY} — and since
@@ -28,13 +30,20 @@
  *              (they cannot contradict prelude bytes that never rendered them).
  */
 
+import { isValidElement, type ReactElement } from "react";
 import { isThenable } from "../../handles/is-thenable.js";
+import {
+  cloneElementWithProps,
+  isPlainDataObject,
+  maskNestedContainerThenables,
+  type MaskReport,
+} from "./mask-nested.js";
 
 // Capture-side nested-thenable masking lives in the LEAF module mask-nested.ts
 // (shared with the capture handle-store push wrap in rsc/shell-capture.ts).
 // Re-exported here so loader-cache and the unit tests keep one import site for
 // the snapshot family.
-export { maskNestedContainerThenables } from "./mask-nested.js";
+export { maskNestedContainerThenables };
 
 /**
  * Marker object standing in for a pending nested promise in a recorded loader
@@ -86,12 +95,6 @@ export function isLoaderSettledMarker(
   );
 }
 
-function isPlainObject(value: unknown): value is Record<string, unknown> {
-  if (typeof value !== "object" || value === null) return false;
-  const proto = Object.getPrototypeOf(value);
-  return proto === Object.prototype || proto === null;
-}
-
 /**
  * Probe whether a promise is already settled without waiting for it: races it
  * against an immediately-resolved sentinel across two microtask hops (then
@@ -129,10 +132,14 @@ export type ElideResult =
  * promises pin their value behind a settled marker (they baked, but consumers
  * hold a promise-shaped key); pending ones become hole markers; a REJECTED
  * nested promise poisons the record (error UI must never bake into a shared
- * shell) — the caller refuses the capture. Only plain objects/arrays are
- * traversed; anything else (Date, Map, class instance) is a pinned leaf.
- * Cycles are cut as pinned references (best effort — Flight rejects true
- * cycles later regardless).
+ * shell) — the caller refuses the capture. Plain data objects, arrays and the
+ * props of host, Suspense/Fragment and client-reference elements are
+ * traversed copy-on-write (an element through cloneElementWithProps); a shared
+ * subtree or a cycle resolves to the same elided result. An element Flight
+ * renders by calling user code (see rendersOnServer) is a leaf, or ONE hole
+ * marker when a thenable sits anywhere in its props: the pin encode calls that
+ * component, so a marker inside its props would reach it. Anything else
+ * (Date, Map, class instance, lazy node) is a pinned leaf.
  *
  * The ROOT container promise-chain unwraps with NO marker: loader-cache
  * overlays against the AWAITED fresh container, so the recorded root must be
@@ -140,7 +147,7 @@ export type ElideResult =
  */
 export async function elideLoaderContainer(
   value: unknown,
-  seen: Set<object> = new Set(),
+  seen: Map<object, ElideResult> = new Map(),
 ): Promise<ElideResult> {
   if (isThenable(value)) {
     const probed = await probeSettled(value);
@@ -155,7 +162,7 @@ export async function elideLoaderContainer(
 
 async function elideNested(
   value: unknown,
-  seen: Set<object>,
+  seen: Map<object, ElideResult>,
 ): Promise<ElideResult> {
   if (isThenable(value)) {
     const probed = await probeSettled(value);
@@ -170,44 +177,146 @@ async function elideNested(
       : { [LOADER_SETTLED_KEY]: 1, value: inner.value };
     return { state: "ok", value: marker, hasHole: inner.hasHole };
   }
+  if (typeof value !== "object" || value === null) {
+    return { state: "ok", value, hasHole: false };
+  }
+  const cached = seen.get(value);
+  if (cached) return cached;
 
-  if (Array.isArray(value)) {
-    if (seen.has(value)) return { state: "ok", value, hasHole: false };
-    seen.add(value);
-    const out: unknown[] = new Array(value.length);
-    let hasHole = false;
-    for (let i = 0; i < value.length; i++) {
-      const r = await elideNested(value[i], seen);
+  if (isValidElement<Record<string, unknown>>(value)) {
+    if (rendersOnServer(value.type)) {
+      return elementHoldsThenable(value.props)
+        ? { state: "ok", value: { [LOADER_HOLE_KEY]: 1 }, hasHole: true }
+        : { state: "ok", value, hasHole: false };
+    }
+    // Cached before the walk: a cycle back here pins the original.
+    const result: ElideResult & { state: "ok" } = {
+      state: "ok",
+      value,
+      hasHole: false,
+    };
+    seen.set(value, result);
+    const props = value.props;
+    let changed: Record<string, unknown> | undefined;
+    for (const key of Object.keys(props)) {
+      const r = await elideNested(props[key], seen);
       if (r.state === "rejected") return r;
-      out[i] = r.value;
-      hasHole ||= r.hasHole;
+      if (r.value !== props[key]) (changed ??= {})[key] = r.value;
+      result.hasHole ||= r.hasHole;
     }
-    return { state: "ok", value: out, hasHole };
+    if (changed) result.value = cloneElementWithProps(value, changed);
+    return result;
   }
 
-  if (isPlainObject(value)) {
-    if (seen.has(value)) return { state: "ok", value, hasHole: false };
-    seen.add(value);
-    const out: Record<string, unknown> = {};
-    let hasHole = false;
-    for (const key of Object.keys(value)) {
-      const r = await elideNested(value[key], seen);
+  if (Array.isArray(value) || isPlainDataObject(value)) {
+    const source = value as Record<string, unknown>;
+    const out = (Array.isArray(value) ? new Array(value.length) : {}) as Record<
+      string,
+      unknown
+    >;
+    // Cached before the walk: a cycle back here pins the copy.
+    const result: ElideResult & { state: "ok" } = {
+      state: "ok",
+      value: out,
+      hasHole: false,
+    };
+    seen.set(value, result);
+    let changed = false;
+    for (const key of Object.keys(source)) {
+      const r = await elideNested(source[key], seen);
       if (r.state === "rejected") return r;
       out[key] = r.value;
-      hasHole ||= r.hasHole;
+      if (r.value !== source[key]) changed = true;
+      result.hasHole ||= r.hasHole;
     }
-    return { state: "ok", value: out, hasHole };
+    if (!changed) result.value = value;
+    return result;
   }
 
   return { state: "ok", value, hasHole: false };
 }
 
+const CLIENT_REFERENCE = Symbol.for("react.client.reference");
+const TEMPORARY_REFERENCE = Symbol.for("react.temporary.reference");
+const REACT_MEMO = Symbol.for("react.memo");
+const REACT_FORWARD_REF = Symbol.for("react.forward_ref");
+const REACT_LAZY = Symbol.for("react.lazy");
+
+/**
+ * Whether Flight calls user code to encode an element of this type (its
+ * renderElement): a function component, or a memo/forwardRef/lazy wrapper
+ * around one. Host (string), Suspense/Fragment (symbol) and client-reference
+ * types are encoded with their props as data, so a hole marker in those props
+ * reaches no component. A lazy type is resolved the way Flight resolves it:
+ * the Flight decode (a cache() hit, a "use cache" result) hands every client
+ * component back as a lazy around its client reference, and the capture render
+ * already resolved it. One that cannot resolve yet counts as a server
+ * component.
+ */
+function rendersOnServer(type: unknown): boolean {
+  if (typeof type !== "function" && (typeof type !== "object" || !type)) {
+    return false;
+  }
+  const tag = (type as { $$typeof?: unknown }).$$typeof;
+  if (tag === CLIENT_REFERENCE || tag === TEMPORARY_REFERENCE) return false;
+  if (typeof type === "function") return true;
+  if (tag === REACT_MEMO) {
+    return rendersOnServer((type as { type?: unknown }).type);
+  }
+  if (tag === REACT_LAZY) {
+    const resolved = readLazyNode(type);
+    return resolved === UNRESOLVED || rendersOnServer(resolved);
+  }
+  return tag === REACT_FORWARD_REF;
+}
+
+/**
+ * Whether the mask found a thenable in an element's props: the mask's own
+ * walk, so elide and the mask share one traversal rule.
+ */
+function elementHoldsThenable(props: unknown): boolean {
+  const report: MaskReport = { thenable: false };
+  maskNestedContainerThenables(props, undefined, report);
+  return report.thenable;
+}
+
+const UNRESOLVED: unique symbol = Symbol("unresolved");
+
+function isLazyNode(value: unknown): boolean {
+  return (
+    typeof value === "object" &&
+    value !== null &&
+    (value as { $$typeof?: unknown }).$$typeof === REACT_LAZY
+  );
+}
+
+/**
+ * The value behind a lazy node, or UNRESOLVED. Flight moves an element past
+ * ~3.2 KB of its row into a row of its own (`$L`), which the decode wraps in a
+ * lazy node, so a hole marker inside that element sits behind one. The pinned
+ * record is fully decoded when a HIT overlays it, so the read is synchronous;
+ * a node that is still pending throws and stays a leaf.
+ */
+function readLazyNode(value: unknown): unknown {
+  if (!isLazyNode(value)) return UNRESOLVED;
+  const lazy = value as { _payload: unknown; _init: (p: unknown) => unknown };
+  try {
+    return lazy._init(lazy._payload);
+  } catch {
+    return UNRESOLVED;
+  }
+}
+
 /**
  * Overlay a recorded (elided) container onto the fresh run's container for a
  * shell HIT. Recorded wins per path; hole markers take the fresh value (the
  * live nested promise); fresh-only object keys pass through. Where the shapes
  * disagree structurally, recorded wins wholesale — it is what the prelude
- * froze, and parity beats freshness inside the shell.
+ * froze, and parity beats freshness inside the shell. Copy-on-write: a
+ * recorded subtree without markers comes back by identity. A recorded element
+ * is rebuilt through cloneElementWithProps only at its marker paths, so its
+ * other props stay pinned; props the fresh element adds do not pass through
+ * (the prelude rendered the element without them).
  */
 export function overlayLoaderContainer(
   fresh: unknown,
@@ -233,27 +342,70 @@ export function overlayLoaderContainer(
 
   if (Array.isArray(recorded)) {
     const freshArr = Array.isArray(fresh) ? fresh : [];
-    return recorded.map((item, i) => overlayLoaderContainer(freshArr[i], item));
+    let out: unknown[] | undefined;
+    for (let i = 0; i < recorded.length; i++) {
+      const next = overlayLoaderContainer(freshArr[i], recorded[i]);
+      if (next !== recorded[i]) (out ??= recorded.slice())[i] = next;
+    }
+    return out ?? recorded;
   }
 
-  if (isPlainObject(recorded)) {
-    if (!isPlainObject(fresh)) {
-      // Shape drift: still honor markers (holes fall back to undefined).
-      const out: Record<string, unknown> = {};
-      for (const key of Object.keys(recorded)) {
-        out[key] = overlayLoaderContainer(undefined, recorded[key]);
-      }
-      return out;
-    }
-    const out: Record<string, unknown> = {};
+  if (isValidElement<Record<string, unknown>>(recorded)) {
+    return overlayElement(fresh, recorded);
+  }
+
+  const lazyValue = readLazyNode(recorded);
+  if (lazyValue !== UNRESOLVED) {
+    const next = overlayLoaderContainer(fresh, lazyValue);
+    return next === lazyValue ? recorded : next;
+  }
+
+  if (isPlainDataObject(recorded)) {
+    // Shape drift (fresh not a plain object): markers fall back to undefined.
+    const freshObj = isPlainDataObject(fresh) ? fresh : undefined;
+    let out: Record<string, unknown> | undefined;
     for (const key of Object.keys(recorded)) {
-      out[key] = overlayLoaderContainer(fresh[key], recorded[key]);
+      const next = overlayLoaderContainer(freshObj?.[key], recorded[key]);
+      if (next !== recorded[key]) (out ??= { ...recorded })[key] = next;
     }
-    for (const key of Object.keys(fresh)) {
-      if (!(key in out)) out[key] = fresh[key];
+    if (freshObj) {
+      for (const key of Object.keys(freshObj)) {
+        if (!(key in recorded)) (out ??= { ...recorded })[key] = freshObj[key];
+      }
     }
-    return out;
+    return out ?? recorded;
   }
 
   return recorded;
 }
+
+function overlayElement(
+  fresh: unknown,
+  recorded: ReactElement<Record<string, unknown>>,
+): unknown {
+  if (isLazyNode(fresh)) {
+    // A fresh lazy node is read only when a recorded path needs the fresh
+    // value: reading it parses a decoded row or calls a React.lazy factory
+    // the HIT render never asks for otherwise.
+    if (overlayProps(undefined, recorded) === recorded) return recorded;
+    const read = readLazyNode(fresh);
+    fresh = read === UNRESOLVED ? undefined : read;
+  }
+  return overlayProps(
+    isValidElement<Record<string, unknown>>(fresh) ? fresh.props : undefined,
+    recorded,
+  );
+}
+
+function overlayProps(
+  freshProps: Record<string, unknown> | undefined,
+  recorded: ReactElement<Record<string, unknown>>,
+): unknown {
+  const props = recorded.props;
+  let changed: Record<string, unknown> | undefined;
+  for (const key of Object.keys(props)) {
+    const next = overlayLoaderContainer(freshProps?.[key], props[key]);
+    if (next !== props[key]) (changed ??= {})[key] = next;
+  }
+  return changed ? cloneElementWithProps(recorded, changed) : recorded;
+}

package/src/router/segment-resolution/mask-nested.ts

@@ -1,14 +1,17 @@
 /**
- * Capture-side nested-thenable masking — the mechanism behind "nested-promise
- * shape is the liveness declaration" for BOTH bake-lane loader containers
- * (loader-cache.ts) and pushed handle containers (the capture handle store's
- * push wrap in rsc/shell-capture.ts).
+ * Capture-side nested thenables. Masking is the mechanism behind "nested-
+ * promise shape is the liveness declaration" for the loader lane: bake-lane
+ * loader containers (loader-cache.ts) and DSL loader handle pushes (the
+ * capture handle store's push wrap in rsc/shell-capture.ts). Handler handle
+ * pushes are handler output and bake: the capture waits for their nested
+ * promises instead (settleNestedThenables).
  *
  * Deliberately a LEAF module: loader-mask (the other natural home) imports
  * request-context, so any funnel importing the mask from there inherits that
- * graph. This module imports only is-thenable.
+ * graph. This module imports only is-thenable and react.
  */
 
+import { cloneElement, isValidElement, type ReactElement } from "react";
 import { isThenable } from "../../handles/is-thenable.js";
 
 /**
@@ -23,18 +26,102 @@ export function createMaskedLoaderPromise<T = unknown>(): Promise<T> {
   return new Promise<T>(() => {});
 }
 
-function isPlainObject(value: unknown): value is Record<string, unknown> {
+/**
+ * A data object the container walks descend into: prototype Object.prototype
+ * or null, and no symbol `$$typeof`. React elements, lazy nodes and references
+ * are plain objects too, told apart by that symbol (Flight switches on it).
+ * They are React's, not data: an element's dev-only fields (`_debugStack`,
+ * `_debugTask`, `_debugInfo`, `ref`) are non-enumerable, so a key-by-key copy
+ * drops them and dev Flight refuses it ("Attempted to render <x> without
+ * development properties", issue #942).
+ */
+export function isPlainDataObject(
+  value: unknown,
+): value is Record<string, unknown> {
   if (typeof value !== "object" || value === null) return false;
   const proto = Object.getPrototypeOf(value);
-  return proto === Object.prototype || proto === null;
+  if (proto !== Object.prototype && proto !== null) return false;
+  return typeof (value as { $$typeof?: unknown }).$$typeof !== "symbol";
 }
 
+type KeyValidated = { _store?: { validated?: unknown } };
+
 /**
- * Deep-copy a container with every NESTED thenable replaced by a masked
+ * cloneElement with only the given props replaced, keeping React's dev key
+ * validation. React 19 dev cloneElement resets `_store.validated`, so a
+ * keyless child of a static children list (validated at creation) would warn
+ * "Each child in a list should have a unique key" once cloned. Pass only the
+ * props that change: cloneElement re-parents the element (owner) when its
+ * config carries `ref`.
+ */
+export function cloneElementWithProps(
+  element: ReactElement,
+  props: Record<string, unknown>,
+): ReactElement {
+  const clone = cloneElement(element, props);
+  const from = (element as KeyValidated)._store;
+  const to = (clone as KeyValidated)._store;
+  if (from && to) to.validated = from.validated;
+  return clone;
+}
+
+/**
+ * Resolves once `value` and every thenable reachable from it through plain
+ * data objects, arrays and React element props has settled, following what
+ * each thenable resolves to. A rejection counts as settled: the consumer that
+ * encodes the value reports it. Other objects are leaves (the same walk as
+ * maskNestedContainerThenables), so a promise inside a Map is not waited for.
+ * An element is walked through its props only (`<X data={promise} />` in a
+ * handle value is handler output too), never its owner or type, and is never
+ * rebuilt: nothing is copied, only awaited.
+ *
+ * The capture waits on this for HANDLER handle pushes (shell-capture.ts):
+ * handler output is baked, so its promises must settle before the doc record
+ * encodes the handles every HIT replays.
+ */
+export function settleNestedThenables(value: unknown): Promise<void> {
+  const pending: Promise<void>[] = [];
+  const seen = new Set<object>();
+  const walk = (v: unknown): void => {
+    if (isThenable(v)) {
+      pending.push(
+        Promise.resolve(v).then(settleNestedThenables, () => undefined),
+      );
+      return;
+    }
+    if (isValidElement<Record<string, unknown>>(v)) {
+      walk(v.props);
+      return;
+    }
+    if (Array.isArray(v) || isPlainDataObject(v)) {
+      if (seen.has(v)) return;
+      seen.add(v);
+      for (const item of Array.isArray(v) ? v : Object.values(v)) walk(item);
+    }
+  };
+  walk(value);
+  return pending.length === 0
+    ? Promise.resolve()
+    : Promise.all(pending).then(() => undefined);
+}
+
+/**
+ * Optional single-walk report for maskNestedContainerThenables: `thenable`
+ * flips true when the walk masked at least one thenable. The capture push
+ * funnel reads it for a loader push: a container with a masked nested
+ * promise keeps its loader tag, so the handle encode never waits on a mask.
+ * Same single-pass shape as elideLoaderContainer's `hasHole`.
+ */
+export interface MaskReport {
+  thenable: boolean;
+}
+
+/**
+ * Copy a container with every NESTED thenable replaced by a masked
  * (never-resolving) promise. Applied during shell capture to (a) bake-lane
- * loader containers (loader-cache.ts) and (b) pushed handle containers
- * (shell-capture.ts capture store push wrap) — the two rango-owned funnels
- * where consumers declare per-request data by promise SHAPE.
+ * loader containers (loader-cache.ts) and (b) DSL loader handle pushes
+ * (shell-capture.ts capture store push wrap) — the rango-owned funnels where
+ * a loader declares per-request data by promise SHAPE.
  *
  * Why: a nested promise that happened to SETTLE before the capture's quiet
  * window closed used to bake its value into the SHARED shell (and, for
@@ -46,23 +133,16 @@ function isPlainObject(value: unknown): value is Record<string, unknown> {
  * Masking makes the consuming subtree postpone as a hole no matter when the
  * promise settles: liveness by declaration, not by racing the window.
  *
- * Only plain objects/arrays are traversed; other values are leaves. The INPUT
- * IS NEVER MUTATED — handler-side loader consumption (the consumption-lane
- * rule, semantic-matrix PPR3) shares the raw container and must keep real
- * values. Cycles are preserved as cycles in the copy.
- */
-/**
- * Optional single-walk report for maskNestedContainerThenables: `thenable`
- * flips true when the walk masked at least one thenable. The shell fast path
- * reads it at handle-push time — a pushed container with a nested thenable
- * declares per-request data, and a shell entry whose HANDLER layer made such
- * a declaration cannot be served handler-free (the hole would never fill).
- * Same single-pass shape as elideLoaderContainer's `hasHole`.
+ * Plain data objects, arrays and React element props are traversed: JSX
+ * carrying a promise (`<Reviews data={fetchReviews()} />`) declares
+ * per-request data the same way a plain container does. Other values, lazy
+ * nodes included, are leaves. Copy-on-write: a subtree without a thenable
+ * comes back by identity, so an element is rebuilt only when its props hold
+ * one, and then through cloneElementWithProps, which keeps its type, key,
+ * owner, dev fields and key validation. The INPUT IS NEVER MUTATED — handler-side loader consumption (the
+ * consumption-lane rule, semantic-matrix PPR3) shares the raw container and
+ * must keep real values. Cycles are preserved as cycles in the copy.
  */
-export interface MaskReport {
-  thenable: boolean;
-}
-
 export function maskNestedContainerThenables(
   value: unknown,
   seen: Map<object, unknown> = new Map(),
@@ -72,27 +152,41 @@ export function maskNestedContainerThenables(
     if (report) report.thenable = true;
     return createMaskedLoaderPromise();
   }
+  if (typeof value !== "object" || value === null) return value;
+  const cached = seen.get(value);
+  if (cached !== undefined) return cached;
 
-  if (Array.isArray(value)) {
-    const cached = seen.get(value);
-    if (cached !== undefined) return cached;
-    const out: unknown[] = new Array(value.length);
-    seen.set(value, out);
-    for (let i = 0; i < value.length; i++) {
-      out[i] = maskNestedContainerThenables(value[i], seen, report);
+  if (isValidElement<Record<string, unknown>>(value)) {
+    // Placeholder: a cycle back to this element through its props keeps the
+    // original.
+    seen.set(value, value);
+    const props = value.props;
+    let maskedProps: Record<string, unknown> | undefined;
+    for (const key of Object.keys(props)) {
+      const masked = maskNestedContainerThenables(props[key], seen, report);
+      if (masked !== props[key]) (maskedProps ??= {})[key] = masked;
     }
-    return out;
+    if (!maskedProps) return value;
+    const clone = cloneElementWithProps(value, maskedProps);
+    seen.set(value, clone);
+    return clone;
   }
 
-  if (isPlainObject(value)) {
-    const cached = seen.get(value);
-    if (cached !== undefined) return cached;
-    const out: Record<string, unknown> = {};
+  if (Array.isArray(value) || isPlainDataObject(value)) {
+    const source = value as Record<string, unknown>;
+    const out = (Array.isArray(value) ? new Array(value.length) : {}) as Record<
+      string,
+      unknown
+    >;
     seen.set(value, out);
-    for (const key of Object.keys(value)) {
-      out[key] = maskNestedContainerThenables(value[key], seen, report);
+    let changed = false;
+    for (const key of Object.keys(source)) {
+      out[key] = maskNestedContainerThenables(source[key], seen, report);
+      if (out[key] !== source[key]) changed = true;
     }
-    return out;
+    if (changed) return out;
+    seen.set(value, value);
+    return value;
   }
 
   return value;

package/src/router/segment-resolution/revalidation.ts

@@ -45,7 +45,8 @@ import {
 } from "./helpers.js";
 import { applyViewTransitionDefault } from "./view-transition-default.js";
 import { getRouterContext } from "../router-context.js";
-import { observeEvent, observeHandler } from "../instrument.js";
+import { observeEvent, observeSegmentHandler } from "../instrument.js";
+import { runInSegmentTagScope } from "../../cache/cache-tag.js";
 import { observeStreamedHandler } from "./streamed-handler-telemetry.js";
 import {
   track,
@@ -221,7 +222,8 @@ export async function resolveLoadersWithRevalidation<TEnv>(
   // whole-container-mask every loader on this path (a never-settling promise
   // for a loader that should have executed).
   // Outside capture the key only activates the _shellLoaderSeed overlay,
-  // which document-only serveShellHit seeds — inert for partial requests.
+  // which serveShellHit seeds for a document HIT and the PPR partial replay
+  // (matchPartialWithPprReplay) seeds once its doc record hits.
   const bakeLane = !entryLoadingMasksLoaders(entry.loading);
 
   const loadersToRun = revalidationChecks.filter((c) => c.shouldRun);
@@ -240,9 +242,10 @@ export async function resolveLoadersWithRevalidation<TEnv>(
             loaderEntry,
             ctx,
             ctx.pathname,
-            bakeLane || loaderEntry.awaitBeforeFlush === true
-              ? segmentId
-              : null,
+            // `bake`, not awaitBeforeFlush: a navigation evaluation never
+            // carries awaitBeforeFlush, and an ssr: false loader must pin
+            // here exactly as on the document HIT.
+            bakeLane || loaderEntry.bake === true ? segmentId : null,
           ),
         ),
         entry,
@@ -476,7 +479,9 @@ async function resolveParallelSlotComponent<TEnv>(args: {
       handlerRan = true;
       if (hasLoadingFallback) {
         const result =
-          typeof handler === "function" ? handler(context) : handler;
+          typeof handler === "function"
+            ? runInSegmentTagScope(parallelId, handler, context)
+            : handler;
         if (result instanceof Promise) {
           warnOnStreamedResponse(result, parallelId);
           const tracked = deps.trackHandler(result, {
@@ -497,7 +502,9 @@ async function resolveParallelSlotComponent<TEnv>(args: {
         }
       } else {
         component =
-          typeof handler === "function" ? await handler(context) : handler;
+          typeof handler === "function"
+            ? await runInSegmentTagScope(parallelId, handler, context)
+            : handler;
       }
     }
   }
@@ -842,14 +849,19 @@ export async function resolveEntryHandlerWithRevalidation<TEnv>(
           : routeEntry.handler;
       if (!routeEntry.loading) {
         const result = handleHandlerResult(
-          await observeHandler(entry.id, handler, context),
+          await observeSegmentHandler(
+            entry.shortCode,
+            entry.id,
+            handler,
+            context,
+          ),
         );
         doneHandler();
         return result;
       }
       if (!actionContext) {
         const result = handleHandlerResult(
-          observeHandler(entry.id, handler, context),
+          observeSegmentHandler(entry.shortCode, entry.id, handler, context),
         );
         if (result instanceof Promise) {
           warnOnStreamedResponse(result, routeEntry.id);
@@ -875,7 +887,12 @@ export async function resolveEntryHandlerWithRevalidation<TEnv>(
         entryId: entry.id,
       });
       const actionResult = handleHandlerResult(
-        await observeHandler(entry.id, handler, context),
+        await observeSegmentHandler(
+          entry.shortCode,
+          entry.id,
+          handler,
+          context,
+        ),
       );
       doneHandler();
       return {