@timber-js/app 0.1.52 → 0.1.54

package/src/client/segment-context.ts

@@ -37,6 +37,13 @@ export function useSegmentContext(): SegmentContextValue | null {
 
 interface SegmentProviderProps {
   segments: string[];
+  /**
+   * Unique identifier for this segment, used by the client-side segment
+   * merger for element caching. For route groups this includes the group
+   * name (e.g., "/(marketing)") since groups share their parent's urlPath.
+   * Falls back to the reconstructed path from `segments` if not provided.
+   */
+  segmentId?: string;
   parallelRouteKeys: string[];
   children: React.ReactNode;
 }
@@ -45,7 +52,7 @@ interface SegmentProviderProps {
  * Wraps each layout to provide segment position context.
  * Injected by rsc-entry.ts during element tree construction.
  */
-export function SegmentProvider({ segments, parallelRouteKeys, children }: SegmentProviderProps) {
+export function SegmentProvider({ segments, segmentId: _segmentId, parallelRouteKeys, children }: SegmentProviderProps) {
   const value = useMemo(
     () => ({ segments, parallelRouteKeys }),
     // segments and parallelRouteKeys are static per layout — they don't change

package/src/client/segment-merger.ts

@@ -0,0 +1,345 @@
+/**
+ * Segment Merger — client-side tree merging for partial RSC payloads.
+ *
+ * When the server skips rendering sync layouts (because the client already
+ * has them cached), the RSC payload is missing outer segment wrappers.
+ * This module reconstructs the full element tree by splicing the partial
+ * payload into cached segment subtrees.
+ *
+ * The approach:
+ * 1. After each full RSC payload render, walk the decoded element tree
+ *    and cache each segment's subtree (identified by SegmentProvider boundaries)
+ * 2. When a partial payload arrives, wrap it with cached segment elements
+ *    using React.cloneElement to preserve component identity
+ *
+ * React.cloneElement preserves the element's `type` — React sees the same
+ * component at the same tree position and reconciles (preserving state)
+ * rather than remounting. This is how layout state survives navigations.
+ *
+ * Design docs: 19-client-navigation.md §"Navigation Reconciliation"
+ * Security: access.ts runs on the server regardless of skipping — this
+ *           is a performance optimization only. See 13-security.md.
+ */
+
+import { cloneElement, isValidElement, type ReactElement, type ReactNode } from 'react';
+
+// ─── Types ───────────────────────────────────────────────────────
+
+/**
+ * A cached segment entry. Stores the full subtree rooted at a SegmentProvider
+ * and the path through the tree to the next SegmentProvider (or leaf).
+ */
+export interface CachedSegmentEntry {
+  /** The segment's URL path (e.g., "/", "/dashboard") */
+  segmentPath: string;
+  /** The SegmentProvider element for this segment */
+  element: ReactElement;
+  /**
+   * Whether this segment's cached element contains a nested SegmentProvider.
+   * Only segments with inner SegmentProviders are safe to skip — the merger
+   * can only replace inner SegmentProviders, not pages embedded in layout output.
+   * Used by the state tree serialization to exclude non-mergeable segments.
+   */
+  hasMergeableChild: boolean;
+}
+
+// ─── Segment Element Cache ───────────────────────────────────────
+
+/**
+ * Cache of React element subtrees per segment path.
+ * Updated after each navigation with the full decoded RSC element tree.
+ */
+export class SegmentElementCache {
+  private entries = new Map<string, CachedSegmentEntry>();
+
+  get(segmentPath: string): CachedSegmentEntry | undefined {
+    return this.entries.get(segmentPath);
+  }
+
+  set(segmentPath: string, entry: CachedSegmentEntry): void {
+    this.entries.set(segmentPath, entry);
+  }
+
+  has(segmentPath: string): boolean {
+    return this.entries.has(segmentPath);
+  }
+
+  clear(): void {
+    this.entries.clear();
+  }
+
+  get size(): number {
+    return this.entries.size;
+  }
+
+  /**
+   * Get the set of segment paths that are safe for the server to skip.
+   * Only segments with an inner SegmentProvider (hasMergeableChild) are
+   * included — the merger can only replace inner SegmentProviders, not
+   * pages embedded in layout output. Used to filter the state tree.
+   *
+   * Returns an empty set if the element cache is empty (no elements
+   * cached yet). This is the safe default — an empty set means no
+   * segments pass the filter, so the state tree is empty and the server
+   * does a full render. The element cache is populated lazily after the
+   * first SPA navigation (RSC-decoded elements from hydration are
+   * thenables that can't be walked until React resolves them).
+   */
+  getMergeablePaths(): Set<string> {
+    const paths = new Set<string>();
+    for (const [, entry] of this.entries) {
+      if (entry.hasMergeableChild) {
+        paths.add(entry.segmentPath);
+      }
+    }
+    return paths;
+  }
+}
+
+// ─── SegmentProvider Detection ───────────────────────────────────
+
+/**
+ * Check if a React element is a SegmentProvider by looking for the
+ * `segments` prop (an array of path segments). This is the only
+ * component that receives this prop shape.
+ */
+export function isSegmentProvider(element: unknown): element is ReactElement {
+  if (!isValidElement(element)) return false;
+  const props = element.props as Record<string, unknown>;
+  return Array.isArray(props.segments);
+}
+
+/**
+ * Extract the segment path from a SegmentProvider element.
+ *
+ * Uses the `segmentId` prop if available (set by the server for route groups
+ * to distinguish siblings that share the same urlPath). Falls back to
+ * reconstructing from the `segments` array prop.
+ */
+export function getSegmentPath(element: ReactElement): string {
+  const props = element.props as { segments: string[]; segmentId?: string };
+  // segmentId is the authoritative key — includes group name for route groups
+  if (props.segmentId) return props.segmentId;
+  const filtered = props.segments.filter(Boolean);
+  return filtered.length === 0 ? '/' : '/' + filtered.join('/');
+}
+
+// ─── Tree Walking ────────────────────────────────────────────────
+
+/**
+ * Walk a React element tree and extract all SegmentProvider boundaries.
+ * Returns an ordered list of segment entries from outermost to innermost.
+ *
+ * This only finds SegmentProviders along the main children path — it does
+ * not descend into parallel routes/slots (those are separate subtrees).
+ */
+export function extractSegments(element: unknown): CachedSegmentEntry[] {
+  const segments: CachedSegmentEntry[] = [];
+  walkForSegments(element, segments);
+  // Compute hasMergeableChild: a segment is mergeable if there's another
+  // SegmentProvider nested below it. The segments list is ordered outermost
+  // to innermost, so each segment's child is the next entry.
+  for (let i = 0; i < segments.length; i++) {
+    segments[i].hasMergeableChild = i < segments.length - 1;
+  }
+  return segments;
+}
+
+function walkForSegments(node: unknown, out: CachedSegmentEntry[]): void {
+  if (!isValidElement(node)) return;
+
+  // Use a local binding to avoid TypeScript narrowing issues with
+  // isSegmentProvider's type predicate on the same variable.
+  const el: ReactElement = node as ReactElement;
+  const props = el.props as Record<string, unknown>;
+
+  if (isSegmentProvider(node)) {
+    out.push({
+      segmentPath: getSegmentPath(el),
+      element: el,
+      hasMergeableChild: false, // computed after collection in extractSegments
+    });
+    // Continue walking into children to find nested segments
+    walkChildren(props.children as ReactNode, out);
+    return;
+  }
+
+  // Not a SegmentProvider — walk children looking for one
+  walkChildren(props.children as ReactNode, out);
+}
+
+function walkChildren(children: ReactNode, out: CachedSegmentEntry[]): void {
+  if (children == null) return;
+
+  if (Array.isArray(children)) {
+    for (const child of children) {
+      walkForSegments(child, out);
+    }
+  } else {
+    walkForSegments(children, out);
+  }
+}
+
+// ─── Cache Population ────────────────────────────────────────────
+
+/**
+ * Cache all segment subtrees from a fully-rendered RSC element tree.
+ * Call this after every full RSC payload render (navigate, refresh, hydration).
+ */
+export function cacheSegmentElements(
+  element: unknown,
+  cache: SegmentElementCache
+): void {
+  const segments = extractSegments(element);
+  for (const entry of segments) {
+    cache.set(entry.segmentPath, entry);
+  }
+}
+
+// ─── Tree Merging ────────────────────────────────────────────────
+
+/**
+ * Find a SegmentProvider nested in the children of a React element.
+ * Returns the path of elements from the given element down to the
+ * SegmentProvider, enabling reconstruction via cloneElement.
+ *
+ * The path is an array of [element, childIndex] pairs. childIndex is -1
+ * for single-child (non-array) props.children.
+ */
+type TreePath = Array<{ element: ReactElement; childIndex: number }>;
+
+function findSegmentProviderPath(
+  node: ReactElement,
+  targetPath?: string
+): TreePath | null {
+  const children = (node.props as { children?: ReactNode }).children;
+  if (children == null) return null;
+
+  if (Array.isArray(children)) {
+    for (let i = 0; i < children.length; i++) {
+      const child = children[i];
+      if (!isValidElement(child)) continue;
+
+      if (isSegmentProvider(child)) {
+        if (!targetPath || getSegmentPath(child) === targetPath) {
+          return [{ element: node, childIndex: i }];
+        }
+      }
+
+      const deeper = findSegmentProviderPath(child, targetPath);
+      if (deeper) {
+        return [{ element: node, childIndex: i }, ...deeper];
+      }
+    }
+  } else if (isValidElement(children)) {
+    if (isSegmentProvider(children)) {
+      if (!targetPath || getSegmentPath(children) === targetPath) {
+        return [{ element: node, childIndex: -1 }];
+      }
+    }
+
+    const deeper = findSegmentProviderPath(children, targetPath);
+    if (deeper) {
+      return [{ element: node, childIndex: -1 }, ...deeper];
+    }
+  }
+
+  return null;
+}
+
+/**
+ * Replace a nested SegmentProvider within a cached element tree with
+ * new content. Uses cloneElement along the path to produce a new tree
+ * with preserved component identity at every level except the replaced node.
+ *
+ * @param cachedElement The cached SegmentProvider element for this segment
+ * @param newInnerContent The new React element to splice in at the inner segment position
+ * @param innerSegmentPath The path of the inner segment to replace (optional — replaces first found)
+ * @returns New element tree with the inner segment replaced
+ */
+export function replaceInnerSegment(
+  cachedElement: ReactElement,
+  newInnerContent: ReactNode,
+  innerSegmentPath?: string
+): ReactElement {
+  const path = findSegmentProviderPath(cachedElement, innerSegmentPath);
+
+  if (!path || path.length === 0) {
+    // No inner SegmentProvider found — this segment's cached element
+    // wraps a page directly (no child layout with a SegmentProvider).
+    // We CANNOT safely replace the page because it's embedded deep in
+    // the layout's server-rendered output tree and we don't know its
+    // position. Return the cached element unchanged as a safety fallback.
+    //
+    // The server should not skip segments without a child layout below
+    // them (enforced by hasRenderedLayoutBelow in buildRouteElement).
+    // If this codepath is reached, it indicates a server/client mismatch.
+    return cachedElement;
+  }
+
+  // Reconstruct bottom-up: replace the innermost element first, then
+  // clone each ancestor with the updated child.
+  let replacement: ReactNode = newInnerContent;
+
+  for (let i = path.length - 1; i >= 0; i--) {
+    const { element, childIndex } = path[i];
+
+    if (childIndex === -1) {
+      // Single child — replace it
+      replacement = cloneElement(element, {}, replacement);
+    } else {
+      // Array children — replace the specific index
+      const children = (element.props as { children: ReactNode[] }).children;
+      const newChildren = [...children];
+      newChildren[childIndex] = replacement;
+      replacement = cloneElement(element, {}, ...newChildren);
+    }
+  }
+
+  return replacement as ReactElement;
+}
+
+/**
+ * Merge a partial RSC payload with cached segment elements.
+ *
+ * When the server skips segments, the partial payload starts from the
+ * first non-skipped segment. This function wraps it with cached elements
+ * for the skipped segments, producing a full tree that React can
+ * reconcile with the mounted tree (preserving layout state).
+ *
+ * @param partialPayload The RSC payload element (may be partial)
+ * @param skippedSegments Ordered list of segment paths that were skipped (outermost first)
+ * @param cache The segment element cache
+ * @returns The merged full element tree, or the partial payload if merging isn't possible
+ */
+export function mergeSegmentTree(
+  partialPayload: unknown,
+  skippedSegments: string[],
+  cache: SegmentElementCache
+): unknown {
+  if (!isValidElement(partialPayload)) return partialPayload;
+  if (skippedSegments.length === 0) return partialPayload;
+
+  // Build from outermost to innermost: each skipped segment's cached
+  // element wraps the next, with the partial payload at the center.
+  let result: ReactNode = partialPayload;
+
+  // Process from innermost skipped segment to outermost
+  for (let i = skippedSegments.length - 1; i >= 0; i--) {
+    const segmentPath = skippedSegments[i];
+    const cached = cache.get(segmentPath);
+
+    if (!cached) {
+      // No cached element for this segment — can't merge.
+      // This shouldn't happen (server only skips segments the client
+      // has cached), but if it does, return the partial payload as-is.
+      return partialPayload;
+    }
+
+    // Replace the inner content of the cached segment with our current result.
+    // The inner content is either the next SegmentProvider or the page.
+    result = replaceInnerSegment(cached.element, result);
+  }
+
+  return result;
+}

package/src/server/route-element-builder.ts

@@ -61,6 +61,13 @@ export interface RouteElementResult {
   segments: ManifestSegmentNode[];
   /** Max deferSuspenseFor hold window across all segments. */
   deferSuspenseFor: number;
+  /**
+   * Segment paths that were skipped because the client already has them cached.
+   * Ordered outermost to innermost. Empty when no segments were skipped.
+   * The client uses this to merge the partial payload with cached segments.
+   * See design/19-client-navigation.md §"X-Timber-State-Tree Header"
+   */
+  skippedSegments: string[];
 }
 
 /**
@@ -305,6 +312,10 @@ export async function buildRouteElement(
     layoutComponents.map(({ component, segment }) => [segment, component])
   );
 
+  // Track which segments were skipped for the X-Timber-Skipped-Segments header.
+  // The client uses this to merge the partial payload with its cached segments.
+  const skippedSegments: string[] = [];
+
   // Wrap from innermost (leaf) to outermost (root), processing every
   // segment in the chain. Each segment may contribute:
   //   1. Error boundaries (status files + error.tsx)
@@ -315,6 +326,16 @@ export async function buildRouteElement(
   // client navigation), sync layouts the client already has are skipped.
   // Access.ts already ran for ALL segments in the pre-render loop above.
   // See design/19-client-navigation.md §"X-Timber-State-Tree Header"
+  //
+  // hasRenderedLayoutBelow tracks whether a non-skipped layout has been
+  // seen below the current segment. A segment can ONLY be skipped if
+  // there is a rendered layout below it — the client merger can only
+  // replace inner SegmentProviders (client component boundaries), not
+  // page content embedded in a layout's server-rendered output.
+  // Without this guard, skipping the innermost layout causes the merger
+  // to drop the layout entirely and replace it with just the page.
+  let hasRenderedLayoutBelow = false;
+
   for (let i = segments.length - 1; i >= 0; i--) {
     const segment = segments[i];
     const isLeaf = i === segments.length - 1;
@@ -323,20 +344,38 @@ export async function buildRouteElement(
     // Check if this segment's layout can be skipped for partial rendering.
     // Skipped segments: no layout wrapping, no error boundaries, no slots,
     // no AccessGate in element tree (access already ran pre-render).
-    const skip = shouldSkipSegment(
-      segment.urlPath,
-      layoutComponent,
-      isLeaf,
-      clientStateTree ?? null
-    );
+    //
+    // Additional constraints beyond shouldSkipSegment:
+    // - Must have a rendered layout below (so the merger can find an
+    //   inner SegmentProvider to splice the new content into)
+    // - Route groups are never skipped because sibling groups share the
+    //   same urlPath (e.g., /(marketing) and /(app) both have "/"),
+    //   which would cause the wrong cached layout to be reused
+    const skip =
+      shouldSkipSegment(
+        segment.urlPath,
+        layoutComponent,
+        isLeaf,
+        clientStateTree ?? null
+      ) &&
+      hasRenderedLayoutBelow &&
+      segment.segmentType !== 'group';
 
     if (skip) {
       // Skip this segment entirely — the client uses its cached version.
       // Access.ts already ran in the pre-render loop (security guarantee).
       // Metadata was already resolved above (head elements are correct).
+      // Record for X-Timber-Skipped-Segments header (outermost first, so prepend).
+      skippedSegments.unshift(segment.urlPath);
       continue;
     }
 
+    // This segment is rendered — mark that future (outer) segments have
+    // a rendered layout below them and can safely be skipped.
+    if (layoutComponent) {
+      hasRenderedLayoutBelow = true;
+    }
+
     // Wrap with error boundaries from this segment (inside layout).
     element = await wrapSegmentWithErrorBoundaries(segment, element, h);
 
@@ -394,8 +433,18 @@ export async function buildRouteElement(
         );
       };
 
+      // segmentId uniquely identifies this segment for client-side element
+      // caching. For route groups, urlPath is shared with the parent (both "/"),
+      // so we include the group name to distinguish them. Without this, the
+      // segment merger's element cache would conflate root and group elements.
+      const segmentId =
+        segment.segmentType === 'group'
+          ? `${segment.urlPath === '/' ? '' : segment.urlPath}/${segment.segmentName}`
+          : segment.urlPath;
+
       element = h(SegmentProvider, {
         segments: segmentPath,
+        segmentId,
         parallelRouteKeys,
         children: h(TracedLayout, {
           ...slotProps,
@@ -413,5 +462,6 @@ export async function buildRouteElement(
     layoutComponents,
     segments,
     deferSuspenseFor,
+    skippedSegments,
   };
 }

package/src/server/rsc-entry/index.ts

@@ -330,7 +330,7 @@ async function renderRoute(
     throw error;
   }
 
-  const { element, headElements, layoutComponents, deferSuspenseFor } = routeResult;
+  const { element, headElements, layoutComponents, deferSuspenseFor, skippedSegments } = routeResult;
 
   // Build head HTML for injection into the SSR output.
   // Collects CSS, fonts, and modulepreload from the build manifest for matched segments.
@@ -434,7 +434,8 @@ async function renderRoute(
       layoutComponents,
       headElements,
       match,
-      responseHeaders
+      responseHeaders,
+      skippedSegments
     );
   }
 

package/src/server/rsc-entry/rsc-payload.ts

@@ -41,7 +41,8 @@ export async function buildRscPayloadResponse(
   layoutComponents: LayoutComponentEntry[],
   headElements: HeadElement[],
   match: RouteMatch,
-  responseHeaders: Headers
+  responseHeaders: Headers,
+  skippedSegments?: string[]
 ): Promise<Response> {
   // Read the first chunk from the RSC stream before committing headers.
   const reader = rscStream.getReader();
@@ -113,6 +114,12 @@ export async function buildRscPayloadResponse(
   const segmentInfo = buildSegmentInfo(segments, layoutComponents);
   responseHeaders.set('X-Timber-Segments', JSON.stringify(segmentInfo));
 
+  // Send skipped segments so the client can merge the partial RSC payload
+  // with its cached segment elements. See design/19-client-navigation.md.
+  if (skippedSegments && skippedSegments.length > 0) {
+    responseHeaders.set('X-Timber-Skipped-Segments', JSON.stringify(skippedSegments));
+  }
+
   // Send route params so the client can populate useParams() after
   // SPA navigation. Without this, useParams() returns {}.
   if (Object.keys(match.params).length > 0) {

package/src/server/state-tree-diff.ts

@@ -55,24 +55,23 @@ export function parseClientStateTree(req: Request): Set<string> | null {
  * @param clientSegments - Set of paths from X-Timber-State-Tree, or null
  */
 export function shouldSkipSegment(
-  _urlPath: string,
-  _layoutComponent: ((...args: unknown[]) => unknown) | undefined,
-  _isLeaf: boolean,
-  _clientSegments: Set<string> | null
+  urlPath: string,
+  layoutComponent: ((...args: unknown[]) => unknown) | undefined,
+  isLeaf: boolean,
+  clientSegments: Set<string> | null
 ): boolean {
-  // DISABLED: Server-side segment skipping is not safe without client-side
-  // tree merging. When the server omits a layout from the RSC payload, the
-  // client receives a tree with a different structure (no SegmentProvider,
-  // no layout wrapper). React sees a different component tree shape and
-  // re-mounts everything — destroying DOM state (input values, focus,
-  // scroll position) in layouts that should have been preserved.
-  //
-  // The correct fix requires client-side tree merging: the router must
-  // splice the partial RSC response into the existing cached tree, only
-  // replacing changed segments. Until that's implemented, always render
-  // the full tree. RSC naturally handles this efficiently — client
-  // components are sent as references, not re-serialized.
-  //
-  // See design/19-client-navigation.md §"X-Timber-State-Tree Header"
-  return false;
+  // No state tree → full render (initial load, refresh, etc.)
+  if (!clientSegments) return false;
+
+  // Leaf segments (pages) are never skipped
+  if (isLeaf) return false;
+
+  // No layout → nothing to skip
+  if (!layoutComponent) return false;
+
+  // Async layouts always re-render (they may depend on request context)
+  if (layoutComponent.constructor?.name === 'AsyncFunction') return false;
+
+  // Skip if the client already has this segment cached
+  return clientSegments.has(urlPath);
 }