@timber-js/app 0.2.0-alpha.124 → 0.2.0-alpha.126

package/src/client/router.ts

@@ -11,7 +11,7 @@ import {
   getNavigationState,
   type NavigationState,
 } from './navigation-context.js';
-import { SegmentElementCache, cacheSegmentElements, mergeSegmentTree } from './segment-merger.js';
+
 import {
   fetchRscPayload,
   RedirectError,
@@ -19,6 +19,7 @@ import {
   VersionSkewError,
 } from './rsc-fetch.js';
 import { setHardNavigating } from './navigation-root.js';
+import { isStaleClientReference, triggerStaleReload } from './stale-reload.js';
 import type { FetchResult } from './rsc-fetch.js';
 
 // ─── Types ───────────────────────────────────────────────────────
@@ -100,8 +101,12 @@ export interface RouterDeps {
   navigateTransition?: (
     pendingUrl: string,
     perform: (
-      wrapPayload: (payload: unknown, navState: NavigationState) => unknown
-    ) => Promise<unknown>
+      wrapPayload: (
+        payload: unknown,
+        navState: NavigationState,
+        segmentUpdates?: Map<string, unknown>
+      ) => unknown
+    ) => Promise<{ element: unknown; decodePromise: Promise<void> | null }>
   ) => Promise<void>;
 
   /**
@@ -132,6 +137,13 @@ export interface RouterDeps {
    */
   completeRouterNavigation?: () => void;
 
+  /**
+   * Get the current unwrapped RSC payload element. Used for partial
+   * navigation: the router re-wraps the same element with new
+   * NavigationProvider context so mounted SegmentOutlets are preserved.
+   */
+  _getCurrentPayload?: () => unknown;
+
   /**
    * Initiate a navigation via the Navigation API (`navigation.navigate()`).
    * Fires the navigate event BEFORE committing the URL, allowing Chrome
@@ -167,12 +179,7 @@ export interface RouterInstance {
    * Called on initial hydration with segment info embedded in the HTML.
    */
   initSegmentCache(segments: SegmentInfo[]): void;
-  /**
-   * Cache segment elements from a decoded RSC element tree.
-   * Called on initial hydration to populate the element cache so the
-   * first client navigation can use partial payloads.
-   */
-  cacheElementTree(element: unknown): void;
+
   /** The segment cache (exposed for tests and <Link> prefetch) */
   segmentCache: SegmentCache;
   /** The prefetch cache (exposed for tests and <Link> prefetch) */
@@ -212,10 +219,7 @@ export function createRouter(deps: RouterDeps): RouterInstance {
   const segmentCache = new SegmentCache();
   const prefetchCache = new PrefetchCache();
   const historyStack = new HistoryStack();
-  const segmentElementCache = new SegmentElementCache();
-
   let routerPhase: RouterPhase = { phase: 'idle' };
-  let navigationId = 0;
   const pendingListeners = new Set<(pending: boolean) => void>();
 
   // AbortController for the current in-flight navigation.
@@ -286,48 +290,29 @@ export function createRouter(deps: RouterDeps): RouterInstance {
     }
   }
 
-  /**
-   * Merge a partial RSC payload with cached segment elements if segments
-   * were skipped. Returns the merged payload ready for rendering.
-   *
-   * Segment element caching is deferred to after paint by the caller
-   * so it doesn't block the render path.
-   */
-  async function mergePayload(
-    payload: unknown,
-    skippedSegments: string[] | null | undefined
-  ): Promise<unknown> {
-    if (skippedSegments && skippedSegments.length > 0) {
-      // Segment merging needs a resolved React element tree for cloneElement.
-      // Await the payload only when merging is actually needed.
-      const resolved =
-        payload != null && typeof payload === 'object' && 'then' in payload
-          ? await (payload as PromiseLike<unknown>)
-          : payload;
-      return mergeSegmentTree(resolved, skippedSegments, segmentElementCache);
-    }
-    // No merging needed — return the payload as-is (may be an unresolved
-    // Flight thenable). React handles progressive rendering via Suspense.
-    return payload;
+  function isPartialNavigation(skippedSegments: string[] | null | undefined): boolean {
+    return skippedSegments != null && skippedSegments.length > 0;
   }
 
   /**
-   * Schedule segment element caching after paint. Resolves thenable
-   * payloads first since tree walking requires resolved React elements.
-   * Guards against stale writes: if a newer navigation started before
-   * this callback fires, the cache write is skipped.
+   * Build a segment updates map for partial navigation. Identifies the
+   * first non-skipped segment and maps it to the payload content.
    */
-  function deferSegmentCaching(payload: unknown): void {
-    const capturedNavId = navigationId;
-    afterPaint(async () => {
-      if (capturedNavId !== navigationId) return;
-      const resolved =
-        payload != null && typeof payload === 'object' && 'then' in payload
-          ? await (payload as PromiseLike<unknown>)
-          : payload;
-      if (capturedNavId !== navigationId) return;
-      cacheSegmentElements(resolved, segmentElementCache);
-    });
+  function buildSegmentUpdates(result: FetchResult): Map<string, unknown> {
+    const skipped = new Set(result.skippedSegments!);
+    const segmentInfo = result.segmentInfo;
+    const updates = new Map<string, unknown>();
+    if (segmentInfo) {
+      for (const info of segmentInfo) {
+        if (!skipped.has(info.segmentId ?? info.path)) {
+          // Use segmentId (outlet key) when available — route groups
+          // use a disambiguated key like "/(marketing)" instead of "/".
+          updates.set(info.segmentId ?? info.path, result.payload);
+          break;
+        }
+      }
+    }
+    return updates;
   }
 
   /**
@@ -371,44 +356,66 @@ export function createRouter(deps: RouterDeps): RouterInstance {
   ): Promise<TransitionResult> {
     if (deps.navigateTransition) {
       let headElements: HeadElement[] | null = null;
-      let decodePromise: Promise<void> | null = null;
       await deps.navigateTransition(url, async (wrapPayload) => {
         const result = await perform();
         headElements = result.headElements;
-        decodePromise = result.decodePromise;
-        // Merge partial payload with cached segments before wrapping.
-        // When no segments were skipped, mergePayload returns the payload
-        // as-is (possibly still an unresolved Flight thenable — React
-        // handles progressive rendering via Suspense).
-        const merged = await mergePayload(result.payload, result.skippedSegments);
-        // Store the MERGED payload in history — not the partial pre-merge tree.
-        // This ensures handlePopState replays the complete tree on back/forward.
+
+        if (isPartialNavigation(result.skippedSegments)) {
+          const segmentUpdates = buildSegmentUpdates(result);
+
+          historyStack.push(url, {
+            payload: null,
+            headElements: result.headElements,
+            params: result.params,
+          });
+
+          // Re-wrap the CURRENT element with new context values.
+          // SegmentOutlets read updates from SegmentUpdateContext;
+          // NavigationProvider gets new params/pathname.
+          const element = wrapPayload(
+            deps._getCurrentPayload?.() ?? result.payload,
+            result.navState,
+            segmentUpdates
+          );
+          return { element, decodePromise: result.decodePromise };
+        }
+
+        // Full navigation — empty updates, render the new tree.
         historyStack.push(url, {
-          payload: merged,
+          payload: result.payload,
           headElements: result.headElements,
           params: result.params,
         });
-        // Defer segment element caching to after paint — it only benefits
-        // future navigations and doesn't need to block the current render.
-        deferSegmentCaching(merged);
-        // Pass navState explicitly — wrapPayload wraps element in
-        // NavigationProvider with this state, no getNavigationState() needed.
-        return wrapPayload(merged, result.navState);
+        const element = wrapPayload(result.payload, result.navState);
+        return { element, decodePromise: result.decodePromise };
       });
-      return { headElements, decodePromise };
+      // decodePromise is awaited inside the transition's startTransition
+      // callback (NavigationRoot), so we don't need it here.
+      return { headElements, decodePromise: null };
     }
     // Fallback: no transition (tests, no React tree)
     const result = await perform();
-    // Merge partial payload with cached segments before rendering
-    const merged = await mergePayload(result.payload, result.skippedSegments);
-    // Store merged payload in history
-    historyStack.push(url, {
-      payload: merged,
-      headElements: result.headElements,
-      params: result.params,
-    });
-    deferSegmentCaching(merged);
-    renderPayload(merged, result.navState);
+    if (isPartialNavigation(result.skippedSegments)) {
+      historyStack.push(url, {
+        payload: null,
+        headElements: result.headElements,
+        params: result.params,
+      });
+    } else {
+      // Resolve thenables before storing/rendering in the fallback path
+      // (tests without navigateTransition). In production, React handles
+      // thenables via Suspense; in tests, renderRoot expects resolved values.
+      const payload =
+        result.payload != null && typeof result.payload === 'object' && 'then' in result.payload
+          ? await (result.payload as PromiseLike<unknown>)
+          : result.payload;
+      historyStack.push(url, {
+        payload,
+        headElements: result.headElements,
+        params: result.params,
+      });
+      renderPayload(payload, result.navState);
+    }
     return { headElements: result.headElements, decodePromise: result.decodePromise };
   }
 
@@ -464,7 +471,7 @@ export function createRouter(deps: RouterDeps): RouterInstance {
     if (result === undefined) {
       // Fetch RSC payload with state tree for partial rendering.
       // Send current URL for intercepting route resolution (modal pattern).
-      const stateTree = segmentCache.serializeStateTree(segmentElementCache.getMergeablePaths());
+      const stateTree = segmentCache.serializeStateTree();
       const rawCurrentUrl = deps.getCurrentUrl();
       const currentUrl = rawCurrentUrl.startsWith('http')
         ? new URL(rawCurrentUrl).pathname
@@ -506,8 +513,6 @@ export function createRouter(deps: RouterDeps): RouterInstance {
     const externalSignal = options._signal as AbortSignal | undefined;
     const skipHistory = options._skipHistory === true;
 
-    navigationId++;
-
     // Create an abort controller for this navigation. Links to the external
     // signal (Navigation API's event.signal) when provided.
     const navAbort = createNavAbort(externalSignal);
@@ -541,7 +546,7 @@ export function createRouter(deps: RouterDeps): RouterInstance {
     setPending(true, url);
 
     try {
-      const { headElements, decodePromise } = await renderViaTransition(url, () =>
+      const { headElements } = await renderViaTransition(url, () =>
         performNavigationFetch(url, {
           replace,
           signal: navAbort.signal,
@@ -549,10 +554,6 @@ export function createRouter(deps: RouterDeps): RouterInstance {
         })
       );
 
-      // Await async RSC decode errors (truncated streams, Flight parse failures)
-      // so they surface through this try/catch rather than becoming unhandled.
-      if (decodePromise) await decodePromise;
-
       // Update document.title and <meta> tags with the new page's metadata
       applyHead(headElements);
 
@@ -599,6 +600,13 @@ export function createRouter(deps: RouterDeps): RouterInstance {
         window.location.href = error.url;
         return new Promise(() => {}) as never;
       }
+      // Stale client reference from streaming RSC decode — the server
+      // was redeployed with new bundles. Same recovery as VersionSkewError.
+      if (isStaleClientReference(error)) {
+        setHardNavigating(true);
+        triggerStaleReload();
+        return new Promise(() => {}) as never;
+      }
       // Abort errors are not application errors — swallow silently.
       if (isAbortError(error)) return;
       throw error;
@@ -621,13 +629,12 @@ export function createRouter(deps: RouterDeps): RouterInstance {
 
   async function refresh(): Promise<void> {
     const currentUrl = deps.getCurrentUrl();
-    navigationId++;
     const navAbort = createNavAbort();
 
     setPending(true, currentUrl);
 
     try {
-      const { headElements, decodePromise } = await renderViaTransition(currentUrl, async () => {
+      const { headElements } = await renderViaTransition(currentUrl, async () => {
         // No state tree sent — server renders the complete RSC payload
         const result = await fetchRscPayload(
           currentUrl,
@@ -642,9 +649,13 @@ export function createRouter(deps: RouterDeps): RouterInstance {
         return { ...result, navState };
       });
 
-      if (decodePromise) await decodePromise;
       applyHead(headElements);
     } catch (error) {
+      if (isStaleClientReference(error)) {
+        setHardNavigating(true);
+        triggerStaleReload();
+        return new Promise(() => {}) as never;
+      }
       // Stale transition (superseded by a newer navigation) or aborted
       // fetch — silently ignore. See TIM-629.
       if (isAbortError(error)) return;
@@ -663,7 +674,6 @@ export function createRouter(deps: RouterDeps): RouterInstance {
     scrollY: number = 0,
     externalSignal?: AbortSignal
   ): Promise<void> {
-    navigationId++;
     // Scroll position is read from history.state by the caller (browser-entry.ts)
     // and passed in. This is more reliable than tracking scroll per-URL in memory
     // because the browser maintains per-entry state even with duplicate URLs.
@@ -683,10 +693,8 @@ export function createRouter(deps: RouterDeps): RouterInstance {
       const navAbort = createNavAbort(externalSignal);
       setPending(true, url);
       try {
-        const { headElements, decodePromise } = await renderViaTransition(url, async () => {
-          const stateTree = segmentCache.serializeStateTree(
-            segmentElementCache.getMergeablePaths()
-          );
+        const { headElements } = await renderViaTransition(url, async () => {
+          const stateTree = segmentCache.serializeStateTree();
           const result = await fetchRscPayload(url, deps, stateTree, undefined, navAbort.signal);
           updateSegmentCache(result.segmentInfo);
           const navState = updateNavigationState(result.params, url);
@@ -694,10 +702,14 @@ export function createRouter(deps: RouterDeps): RouterInstance {
           return { ...result, navState };
         });
 
-        if (decodePromise) await decodePromise;
         applyHead(headElements);
         restoreScrollAfterPaint(scrollY);
       } catch (error) {
+        if (isStaleClientReference(error)) {
+          setHardNavigating(true);
+          triggerStaleReload();
+          return new Promise(() => {}) as never;
+        }
         // Stale transition (superseded by a newer navigation) or aborted
         // fetch — silently ignore. See TIM-629.
         if (isAbortError(error)) return;
@@ -721,9 +733,10 @@ export function createRouter(deps: RouterDeps): RouterInstance {
     if (historyStack.has(url)) return;
 
     // Fire-and-forget fetch
-    const stateTree = segmentCache.serializeStateTree(segmentElementCache.getMergeablePaths());
+    const stateTree = segmentCache.serializeStateTree();
     void fetchRscPayload(url, deps, stateTree).then(
       (result) => {
+        result.decodePromise?.catch(() => {});
         prefetchCache.set(url, result);
       },
       () => {
@@ -754,7 +767,7 @@ export function createRouter(deps: RouterDeps): RouterInstance {
         payload: element,
         headElements,
       });
-      deferSegmentCaching(element);
+
       // Revalidation doesn't change params/pathname — preserve current state.
       // DO NOT call updateNavigationState(null, ...) here: that normalizes
       // params to {}, clearing dynamic route params on every action response.
@@ -763,7 +776,6 @@ export function createRouter(deps: RouterDeps): RouterInstance {
       applyHead(headElements);
     },
     initSegmentCache: (segments: SegmentInfo[]) => updateSegmentCache(segments),
-    cacheElementTree: (element: unknown) => cacheSegmentElements(element, segmentElementCache),
     segmentCache,
     prefetchCache,
     historyStack,

package/src/client/rsc-fetch.ts

@@ -258,6 +258,54 @@ export class ServerErrorResponse extends Error {
   }
 }
 
+// ─── Stream Completion Tracking ───────────────────────────────────
+
+/**
+ * Wrap a response body stream to track when it's fully consumed.
+ * Returns a new body that passes all chunks through unchanged, plus
+ * a `done` promise that resolves when the last chunk is read (or
+ * rejects if the stream errors).
+ *
+ * Used to keep React transitions open for the full RSC stream
+ * duration — createFromFetch's thenable resolves on shell arrival,
+ * but we need stream completion for useOptimistic pending state.
+ */
+function trackStreamCompletion(body: ReadableStream<Uint8Array>): {
+  body: ReadableStream<Uint8Array>;
+  done: Promise<void>;
+} {
+  let resolveDone!: () => void;
+  let rejectDone!: (e: unknown) => void;
+  const done = new Promise<void>((res, rej) => {
+    resolveDone = res;
+    rejectDone = rej;
+  });
+
+  const reader = body.getReader();
+  const tracked = new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      try {
+        const result = await reader.read();
+        if (result.done) {
+          controller.close();
+          resolveDone();
+        } else {
+          controller.enqueue(result.value);
+        }
+      } catch (error) {
+        controller.error(error);
+        rejectDone(error);
+      }
+    },
+    cancel(reason) {
+      reader.cancel(reason);
+      resolveDone();
+    },
+  });
+
+  return { body: tracked, done };
+}
+
 // ─── Fetch ───────────────────────────────────────────────────────
 
 /**
@@ -289,6 +337,11 @@ export async function fetchRscPayload(
     let segmentInfo: SegmentInfo[] | null = null;
     let params: Record<string, string | string[]> | null = null;
     let skippedSegments: string[] | null = null;
+    // Track when the full RSC body stream is consumed (not just shell).
+    // Initialized to resolved for bodyless responses; overwritten when
+    // the response has a body.
+    let streamDone: Promise<void> = Promise.resolve();
+
     const wrappedPromise = fetchPromise.then((response) => {
       // Version skew detection (TIM-446): if the server signals a reload,
       // throw VersionSkewError so the caller (router navigate) can trigger
@@ -317,6 +370,20 @@ export async function fetchRscPayload(
       segmentInfo = extractSegmentInfo(response);
       params = extractParams(response);
       skippedSegments = extractSkippedSegments(response);
+
+      // Wrap the body to track full stream consumption. createFromFetch's
+      // thenable resolves when the root model (shell) arrives, but we need
+      // to know when ALL chunks are read so the React transition stays
+      // open for the full streaming duration (keeps useOptimistic alive).
+      if (response.body) {
+        const tracked = trackStreamCompletion(response.body);
+        streamDone = tracked.done;
+        streamDone.catch(() => {}); // prevent unhandled rejection
+        return new Response(tracked.body, {
+          headers: response.headers,
+          status: response.status,
+        });
+      }
       return response;
     });
     // Await headers so headElements/segmentInfo/params are populated.
@@ -326,12 +393,24 @@ export async function fetchRscPayload(
     // parts and progressively renders as chunks arrive, spreading work across
     // frames instead of blocking the main thread in one burst.
     const payload = deps.decodeRsc(wrappedPromise);
-    // Track decode completion so callers can catch async errors (truncated
-    // streams, Flight parse failures) that fire after fetchRscPayload returns.
-    const decodePromise = Promise.resolve(payload).then(() => undefined);
-    // Prevent unhandled rejection — the router monitors decodePromise separately.
-    decodePromise.catch(() => {});
-    return { payload, decodePromise, headElements, segmentInfo, params, skippedSegments };
+    // Combine stream completion with payload error propagation.
+    // streamDone keeps the transition open for the full RSC stream
+    // duration (useOptimistic pending state). payloadError propagates
+    // decode failures (stale client references, Flight parse errors)
+    // so the router's catch block can trigger recovery (stale reload).
+    const payloadError = new Promise<void>((_, reject) => {
+      Promise.resolve(payload).then(() => {}, reject);
+    });
+    payloadError.catch(() => {});
+    const decodePromise = Promise.race([streamDone, payloadError]);
+    return {
+      payload,
+      decodePromise,
+      headElements,
+      segmentInfo,
+      params,
+      skippedSegments,
+    };
   }
   // Test/fallback path: return raw text
   const response = await deps.fetch(rscUrl, { headers, redirect: 'manual', signal });

package/src/client/segment-cache.ts

@@ -111,6 +111,8 @@ function collectSyncSegments(
  */
 export interface SegmentInfo {
   path: string;
+  /** Outlet key — includes route group name when applicable (e.g., "/(marketing)"). */
+  segmentId?: string;
   isAsync: boolean;
 }
 
@@ -137,8 +139,9 @@ export function buildSegmentTree(segments: SegmentInfo[]): SegmentNode | undefin
   let parent: SegmentNode | undefined;
 
   for (const info of layouts) {
+    const id = info.segmentId ?? info.path;
     const node: SegmentNode = {
-      segment: info.path,
+      segment: id,
       payload: null,
       isAsync: info.isAsync,
       children: new Map(),
@@ -149,7 +152,7 @@ export function buildSegmentTree(segments: SegmentInfo[]): SegmentNode | undefin
     }
 
     if (parent) {
-      parent.children.set(info.path, node);
+      parent.children.set(id, node);
     }
 
     parent = node;

package/src/client/segment-outlet.tsx

@@ -1,33 +1,31 @@
 /**
  * SegmentOutlet — client component boundary at each layout segment.
  *
- * Replaces the post-hoc tree walking in segment-merger.ts with an explicit
- * client component at each segment boundary. Each outlet:
+ * Each layout in the segment tree is wrapped with a SegmentOutlet that:
+ * 1. Knows its segment path (prop from the server)
+ * 2. Reads from SegmentUpdateContext for partial navigation updates
+ * 3. Caches its rendered content in a ref across navigations
  *
- * 1. Knows its own segment path (prop from the server)
- * 2. Caches its children in a ref across navigations
- * 3. When `keepCurrent` is true (partial navigation, this segment skipped),
- *    returns the previously cached children — layout state is preserved
- * 4. When `keepCurrent` is false (full navigation or this segment changed),
- *    stores and renders the new children
+ * On full navigation: receives new children via props, caches and renders them.
+ * On partial navigation (this segment skipped): the context map has no entry
+ * for this path, so the outlet returns cached content — preserving layout state.
+ * On partial navigation (this segment updated): the context map has content
+ * for this path, so the outlet renders the update.
  *
- * This eliminates the need for client-side element tree walking, which
- * breaks on real RSC trees due to opaque client component lazy refs,
- * Suspense thenables, and AccessGate wrappers.
+ * Uses React context instead of useSyncExternalStore to stay compatible
+ * with concurrent rendering (transitions). All outlets re-render when the
+ * context value changes, but each bails out quickly if its segment has
+ * no update — same approach as Next.js LayoutRouter.
  *
- * Architecture is similar to Next.js's `<LayoutRouter>` client component —
- * each layout boundary is an explicit client component that manages its
- * own subtree. See design/19-client-navigation.md.
- *
- * Security: This is a performance optimization only. The server always
- * runs all access.ts files regardless of segment skipping. A fabricated
- * keepCurrent prop can only cause stale layouts — never auth bypass.
+ * Security: performance optimization only. The server always runs all
+ * access.ts files regardless of segment skipping.
  * See design/13-security.md §"State tree manipulation".
  */
 
 'use client';
 
-import { useRef, type ReactNode } from 'react';
+import { useContext, useRef, type ReactNode } from 'react';
+import { SegmentUpdateContext } from './segment-update-context.js';
 
 export interface SegmentOutletProps {
   /**
@@ -38,17 +36,6 @@ export interface SegmentOutletProps {
    */
   segmentPath: string;
 
-  /**
-   * When true, the outlet returns its previously cached children instead
-   * of rendering the new children prop. Set by the server when this
-   * segment was skipped (the client already has the layout mounted).
-   *
-   * On the first render (SSR/hydration), this is always false — there's
-   * no cached content yet. On subsequent partial navigations, the server
-   * sets this to true for segments it skipped rendering.
-   */
-  keepCurrent?: boolean;
-
   /** The segment's React subtree (layout + inner content). */
   children: ReactNode;
 }
@@ -56,31 +43,35 @@ export interface SegmentOutletProps {
 /**
  * Client component boundary at each layout segment in the element tree.
  *
- * On full navigation: receives new children, stores them, renders them.
- * On partial navigation (keepCurrent=true): ignores children prop,
- * returns previously stored content — React reconciles the same elements,
- * preserving all client component state in the layout subtree.
+ * On full navigation (context map is empty): renders children, caches in ref.
+ * On partial navigation: checks the context map for an update at segmentPath.
+ *   - Update found → renders update, caches in ref.
+ *   - No update → returns cached content (layout state preserved).
  *
- * React preserves the ref across `reactRoot.render()` calls because:
- * - SegmentOutlet has a stable type (client component module reference)
- * - It appears at the same tree position on every navigation
- * - React reconciles same-type, same-position → instance preserved
+ * React preserves component instances across reactRoot.render() calls
+ * when the same type appears at the same tree position. The ref persists
+ * across navigations because SegmentOutlet is reconciled, not remounted.
  */
-export function SegmentOutlet({
-  segmentPath: _segmentPath,
-  keepCurrent = false,
-  children,
-}: SegmentOutletProps) {
-  // Store content in a ref to avoid triggering re-renders on cache updates.
-  // The ref persists across reactRoot.render() calls because React reconciles
-  // the same component type at the same tree position.
+export function SegmentOutlet({ segmentPath, children }: SegmentOutletProps) {
+  const updates = useContext(SegmentUpdateContext);
   const contentRef = useRef<ReactNode>(null);
 
-  if (!keepCurrent) {
-    // Full render or this segment was re-rendered — store and render new content
+  const update = updates.get(segmentPath);
+
+  if (update !== undefined) {
+    contentRef.current = update;
+    return update;
+  }
+
+  if (contentRef.current === null) {
+    contentRef.current = children;
+    return children;
+  }
+
+  if (children !== contentRef.current) {
     contentRef.current = children;
+    return children;
   }
-  // else: keepCurrent=true — return previously cached content
 
   return contentRef.current;
 }