@rangojs/router 0.0.0-experimental.2 → 0.0.0-experimental.20

package/src/browser/partial-update.ts

@@ -5,14 +5,20 @@ import type {
   ResolvedSegment,
 } from "./types.js";
 import type { ReactNode } from "react";
+import * as React from "react";
 import { startTransition } from "react";
+
+// addTransitionType is only available in React experimental
+const addTransitionType: ((type: string) => void) | undefined =
+  "addTransitionType" in React ? (React as any).addTransitionType : undefined;
 import type { RenderSegmentsOptions } from "../segment-system.js";
-import {
-  mergeSegmentLoaders,
-  needsLoaderMerge,
-  insertMissingDiffSegments,
-} from "./merge-segment-loaders.js";
-import type { BoundTransaction } from "./navigation-bridge.js";
+import { reconcileSegments } from "./segment-reconciler.js";
+import type { ReconcileActor } from "./segment-reconciler.js";
+import { hasActiveIntercept as hasActiveInterceptSlots } from "./intercept-utils.js";
+import type { BoundTransaction } from "./navigation-transaction.js";
+import { ServerRedirect } from "../errors.js";
+import { debugLog } from "./logging.js";
+import { validateRedirectOrigin } from "./validate-redirect-origin.js";
 
 /**
  * Configuration for creating a partial updater
@@ -41,12 +47,30 @@ export interface CommitOverrides {
   intercept?: boolean;
   /** Source URL where intercept was triggered from */
   interceptSourceUrl?: string;
+  /** Server-set location state to merge into history.pushState */
+  serverState?: Record<string, unknown>;
 }
 
 /**
- * Commit context passed to partial updater for URL updates
- * Transaction encapsulates all store mutations for atomic commit
+ * Discriminated update mode for partial updates.
  */
+export type UpdateMode =
+  | {
+      type: "navigate";
+      /** Cached segments for the target URL. When provided, these are used to build
+       * the segment map instead of the current page's segments. This ensures consistency
+       * when we send cached segment IDs to the server - if the server returns empty diff,
+       * we use the same segments we told the server we have. */
+      targetCacheSegments?: ResolvedSegment[];
+      /** Cached handle data for the target URL. When server returns empty diff and we're
+       * rendering from cache, this is passed to the UI to restore breadcrumbs etc. */
+      targetCacheHandleData?: Record<string, Record<string, unknown[]>>;
+      /** Source URL for intercept restore (popstate cache miss) */
+      interceptSourceUrl?: string;
+    }
+  | { type: "leave-intercept" }
+  | { type: "stale-revalidation"; interceptSourceUrl?: string }
+  | { type: "action"; interceptSourceUrl?: string };
 
 /**
  * Type for the fetchPartialUpdate function
@@ -56,21 +80,9 @@ export type PartialUpdater = (
   segmentIds: string[] | undefined,
   isRetry: boolean,
   signal: AbortSignal | undefined,
-  type: BoundTransaction,
-  options?: {
-    isAction?: boolean;
-    staleRevalidation?: boolean;
-    interceptSourceUrl?: string;
-    /** Cached segments for the target URL. When provided, these are used to build
-     * the segment map instead of the current page's segments. This ensures consistency
-     * when we send cached segment IDs to the server - if the server returns empty diff,
-     * we use the same segments we told the server we have. */
-    targetCacheSegments?: ResolvedSegment[];
-    /** Cached handle data for the target URL. When server returns empty diff and we're
-     * rendering from cache, this is passed to the UI to restore breadcrumbs etc. */
-    targetCacheHandleData?: Record<string, Record<string, unknown[]>>;
-  },
-) => Promise<Promise<void>>;
+  tx: BoundTransaction,
+  mode?: UpdateMode,
+) => Promise<void>;
 
 /**
  * Create a partial updater for fetching and applying RSC partial updates
@@ -80,18 +92,6 @@ export type PartialUpdater = (
  *
  * @param config - Partial update configuration
  * @returns fetchPartialUpdate function
- *
- * @example
- * ```typescript
- * const fetchPartialUpdate = createPartialUpdater({
- *   store,
- *   client,
- *   onUpdate: (update) => store.emit(update),
- *   renderSegments,
- * });
- *
- * await fetchPartialUpdate('/new-page');
- * ```
  */
 export function createPartialUpdater(
   config: PartialUpdateConfig,
@@ -99,20 +99,16 @@ export function createPartialUpdater(
   const { store, client, onUpdate, renderSegments, version } = config;
 
   /**
-   * Build a lookup map from current page's cached segments
+   * Get current page's cached segments as an array
    */
-  function getCurrentSegmentMap(): Map<string, ResolvedSegment> {
+  function getCurrentCachedSegments(): ResolvedSegment[] {
     const currentKey = store.getHistoryKey();
     const cached = store.getCachedSegments(currentKey);
-    const cachedSegments = cached?.segments || [];
-    const map = new Map<string, ResolvedSegment>();
-    cachedSegments.forEach((s) => map.set(s.id, s));
-    return map;
+    return cached?.segments || [];
   }
 
   /**
    * Fetch partial update and trigger UI update
-   * Returns a promise that resolves when the RSC stream is fully consumed
    *
    * @param tx - Transaction for committing segment state (required)
    * @param signal - AbortSignal to check if navigation is stale (not for aborting fetch)
@@ -123,102 +119,127 @@ export function createPartialUpdater(
     isRetry: boolean,
     signal: AbortSignal | undefined,
     tx: BoundTransaction,
-    options?: {
-      isAction?: boolean;
-      staleRevalidation?: boolean;
-      interceptSourceUrl?: string;
-      targetCacheSegments?: ResolvedSegment[];
-      targetCacheHandleData?: Record<string, Record<string, unknown[]>>;
-    },
-  ): Promise<Promise<void>> {
-    const {
-      isAction = false,
-      staleRevalidation = false,
-      interceptSourceUrl,
-      targetCacheSegments,
-      targetCacheHandleData,
-    } = options || {};
+    mode: UpdateMode = { type: "navigate" },
+  ): Promise<void> {
     const segmentState = store.getSegmentState();
     const url = targetUrl || window.location.href;
 
     // Capture history key at start for stale revalidation consistency check
     const historyKeyAtStart = store.getHistoryKey();
-    const segments = segmentIds ?? segmentState.currentSegmentIds;
+
+    // Derive interceptSourceUrl from modes that carry it
+    const interceptSourceUrl =
+      mode.type === "stale-revalidation" ||
+      mode.type === "action" ||
+      mode.type === "navigate"
+        ? mode.interceptSourceUrl
+        : undefined;
+
+    // When leaving intercept, filter out intercept-specific segments
+    let segments: string[];
+    if (mode.type === "leave-intercept") {
+      const currentSegments = segmentIds ?? segmentState.currentSegmentIds;
+      const currentCached = getCurrentCachedSegments();
+      const interceptIds = new Set(
+        currentCached
+          .filter((s) => s.namespace?.startsWith("intercept:"))
+          .map((s) => s.id),
+      );
+      segments = currentSegments.filter((id) => !interceptIds.has(id));
+      debugLog(
+        `[Browser] Leaving intercept - filtered segments: ${segments.join(", ")}`,
+      );
+    } else {
+      segments = segmentIds ?? segmentState.currentSegmentIds;
+    }
 
     // For intercept revalidation, use the intercept source URL as previousUrl
-    // This tells the server the route should be treated as an intercept
     const previousUrl =
       interceptSourceUrl || tx.currentUrl || segmentState.currentUrl;
 
-    console.log(`\n[Browser] >>> NAVIGATION`);
-    console.log(`[Browser] From: ${previousUrl}`);
-    console.log(`[Browser] To: ${url}`);
-    console.log(`[Browser] Segments to send: ${segments.join(", ")}`);
+    debugLog(`\n[Browser] >>> NAVIGATION`);
+    debugLog(`[Browser] From: ${previousUrl}`);
+    debugLog(`[Browser] To: ${url}`);
+    debugLog(`[Browser] Segments to send: ${segments.join(", ")}`);
     if (interceptSourceUrl) {
-      console.log(`[Browser] Intercept context from: ${interceptSourceUrl}`);
+      debugLog(`[Browser] Intercept context from: ${interceptSourceUrl}`);
     }
 
-    // Build segment map for merging with server diff.
-    // When targetCacheSegments is provided (navigating to a cached route), use those
-    // to ensure consistency - we use the same segments we told the server we have.
+    // Get cached segments for merging with server diff.
+    // When navigating with targetCacheSegments, use those for consistency.
     // Otherwise fall back to current page's segments (for same-route revalidation).
-    let currentSegmentMap: Map<string, ResolvedSegment>;
-    if (targetCacheSegments && targetCacheSegments.length > 0) {
-      currentSegmentMap = new Map();
-      targetCacheSegments.forEach((s) => currentSegmentMap.set(s.id, s));
-    } else {
-      currentSegmentMap = getCurrentSegmentMap();
-    }
-    // Mark navigation as streaming (response received, now parsing RSC)
-    // The token is ended when the stream completes
-    const streamingToken = tx.startStreaming();
+    const targetCache =
+      mode.type === "navigate" ? mode.targetCacheSegments : undefined;
+    const cachedSegs =
+      targetCache && targetCache.length > 0
+        ? targetCache
+        : getCurrentCachedSegments();
+
     // Fetch partial payload (no abort signal - RSC doesn't support it well)
-    const { payload, streamComplete: rawStreamComplete } =
-      await client.fetchPartial({
-        targetUrl: url,
-        segmentIds: segments,
-        previousUrl,
-        staleRevalidation,
-        version,
-      });
-    console.log("payload.metadata", payload.metadata);
+    let fetchResult: Awaited<ReturnType<NavigationClient["fetchPartial"]>>;
+    fetchResult = await client.fetchPartial({
+      targetUrl: url,
+      segmentIds: segments,
+      previousUrl,
+      // Mark stale when explicitly requested OR when no segments are sent
+      // (action redirect sends empty segments for a fresh render).
+      staleRevalidation:
+        mode.type === "stale-revalidation" || segments.length === 0,
+      version,
+    });
+    // Mark navigation as streaming (response received, now parsing RSC).
+    // Called after fetchPartial so pendingUrl stays set during the network wait,
+    // allowing useLinkStatus to show per-link pending indicators.
+    const streamingToken = tx.startStreaming();
+    const { payload, streamComplete: rawStreamComplete } = fetchResult;
+    debugLog("payload.metadata", payload.metadata);
 
     const streamComplete = rawStreamComplete.then(() => {
       streamingToken.end();
     });
 
+    // Handle server-side redirect with state
+    if (payload.metadata?.redirect) {
+      if (signal?.aborted) {
+        debugLog("[Browser] Ignoring stale redirect (aborted)");
+        return;
+      }
+      const redirectUrl = validateRedirectOrigin(
+        payload.metadata.redirect.url,
+        window.location.origin,
+      );
+      if (!redirectUrl) {
+        debugLog("[Browser] Ignoring blocked redirect payload");
+        return;
+      }
+      const serverState = payload.metadata.locationState;
+      throw new ServerRedirect(redirectUrl, serverState);
+    }
+
     if (payload.metadata?.isPartial) {
       const { segments: newSegments, matched, diff } = payload.metadata;
 
       // Check if this navigation is stale (a newer one started)
       if (signal?.aborted) {
-        console.log(`[Browser] Ignoring stale navigation (aborted)`);
-        return streamComplete;
+        debugLog("[Browser] Ignoring stale navigation (aborted)");
+        return;
       }
 
-      console.log(`[Browser] Partial update - matched: ${matched?.join(", ")}`);
-      console.log(`[Browser] Diff: ${diff?.join(", ")}`);
-
-      // Create lookup for new segments from server
-      const newSegmentMap = new Map<string, ResolvedSegment>();
-      (newSegments || []).forEach((s: ResolvedSegment) =>
-        newSegmentMap.set(s.id, s),
-      );
+      debugLog(`[Browser] Partial update - matched: ${matched?.join(", ")}`);
+      debugLog(`[Browser] Diff: ${diff?.join(", ")}`);
 
       // If diff is empty, nothing changed on server side.
-      // However, if we're navigating with targetCacheSegments (to a different route),
-      // we still need to render those segments since the UI is showing the old route.
       if (!diff || diff.length === 0) {
         const matchedIds = matched || [];
+        const cacheMap = new Map(cachedSegs.map((s) => [s.id, s]));
         const existingSegments = matchedIds
-          .map((id: string) => currentSegmentMap.get(id))
+          .map((id: string) => cacheMap.get(id))
           .filter(Boolean) as ResolvedSegment[];
 
         // When navigating with cached segments to a different route, render them.
-        // targetCacheSegments being provided means we're navigating to a cached route.
-        if (targetCacheSegments && targetCacheSegments.length > 0) {
-          console.log(
-            `[Browser] No diff but navigating with cached segments - rendering target route`,
+        if (mode.type === "navigate" && targetCache && targetCache.length > 0) {
+          debugLog(
+            "[Browser] No diff but navigating with cached segments - rendering target route",
           );
 
           const newTree = await renderSegments(existingSegments, {
@@ -229,98 +250,86 @@ export function createPartialUpdater(
 
           // Include cachedHandleData in metadata so NavigationProvider can restore
           // breadcrumbs and other handle data from cache.
-          // IMPORTANT: Remove `handles` from metadata to prevent NavigationProvider from
+          // Remove `handles` from metadata to prevent NavigationProvider from
           // processing an empty handles stream, which would clear the cached breadcrumbs.
-          // When rendering from cache with empty diff, we want to use cachedHandleData instead.
-          const { handles: _unusedHandles, ...metadataWithoutHandles } = payload.metadata!;
-          onUpdate({
+          const { handles: _unusedHandles, ...metadataWithoutHandles } =
+            payload.metadata!;
+          const cachedUpdate = {
             root: newTree,
             metadata: {
               ...metadataWithoutHandles,
-              cachedHandleData: targetCacheHandleData,
+              cachedHandleData: mode.targetCacheHandleData,
             },
+          };
+
+          const cachedHasTransition = existingSegments.some(
+            (s) => s.transition,
+          );
+          if (cachedHasTransition) {
+            startTransition(() => {
+              if (addTransitionType) {
+                addTransitionType("navigation");
+              }
+              onUpdate(cachedUpdate);
+            });
+          } else {
+            onUpdate(cachedUpdate);
+          }
+
+          debugLog("[Browser] Navigation complete (rendered from cache)");
+          return;
+        }
+
+        // When leaving intercept, force re-render even with empty diff
+        if (mode.type === "leave-intercept") {
+          debugLog(
+            "[Browser] Leaving intercept - forcing re-render to remove modal",
+          );
+
+          const newTree = await renderSegments(existingSegments, {
+            forceAwait: true,
+          });
+
+          tx.commit(matchedIds, existingSegments);
+
+          onUpdate({
+            root: newTree,
+            metadata: payload.metadata,
           });
 
-          console.log(`[Browser] Navigation complete (rendered from cache)\n`);
-          return streamComplete;
+          debugLog("[Browser] Navigation complete (left intercept)");
+          return;
         }
 
         // Same route revalidation with no changes - skip UI update
-        console.log(
-          `[Browser] No changes - all revalidations returned false, keeping existing UI`,
+        debugLog(
+          "[Browser] No changes - all revalidations returned false, keeping existing UI",
         );
         tx.commit(matchedIds, existingSegments);
-        console.log(`[Browser] Navigation complete (no re-render)\n`);
-        return streamComplete;
+        debugLog("[Browser] Navigation complete (no re-render)");
+        return;
       }
 
-      // Build full segment list by merging:
-      // - New/changed segments from server response (diff)
-      // - Unchanged segments from current page's cache
+      // Reconcile server segments with cached segments (single source of truth)
       const matchedIds = matched || [];
-      console.log(`[Browser] matchedIds: ${matchedIds.join(", ")}`);
-      console.log(
-        `[Browser] currentSegmentMap keys: ${[...currentSegmentMap.keys()].join(", ")}`,
-      );
-      console.log(
-        `[Browser] newSegmentMap keys: ${[...newSegmentMap.keys()].join(", ")}`,
-        newSegmentMap,
-      );
-
-      // First pass: build segments from matched IDs
-      const matchedIdSet = new Set(matchedIds);
-      const allSegments = matchedIds
-        .map((id: string) => {
-          // First check server response (new/updated segments)
-          const fromServer = newSegmentMap.get(id);
-          if (fromServer) {
-            // For partial revalidation (stale or action), merge server's new loader data
-            // with cached loader data when server returns fewer loaders than cached
-            const fromCache = currentSegmentMap.get(id);
-            if (
-              (staleRevalidation || isAction) &&
-              needsLoaderMerge(fromServer, fromCache)
-            ) {
-              return mergeSegmentLoaders(fromServer, fromCache);
-            }
-            // When server returns component: null for a layout segment, it means
-            // "this segment doesn't need re-rendering" - preserve the cached component
-            // to maintain the outlet chain and prevent React tree changes
-            if (
-              fromServer.component === null &&
-              fromServer.type === "layout" &&
-              fromCache?.component != null
-            ) {
-              console.log(
-                `[Browser] Preserving cached component for layout ${id} (server returned null)`,
-              );
-              return { ...fromServer, component: fromCache.component };
-            }
-            return fromServer;
-          }
-          // Fall back to current page's cached segments
-          const fromCache = currentSegmentMap.get(id);
-          if (!fromCache) {
-            console.warn(`[Browser] Missing segment: ${id}`);
-            return fromCache;
-          }
-          // Clear loading for cached segments to prevent suspense - server decided
-          // this segment doesn't need re-rendering, so show content as-is
-          if (fromCache.loading !== undefined) {
-            return { ...fromCache, loading: undefined };
-          }
-          return fromCache;
-        })
-        .filter(Boolean) as ResolvedSegment[];
-
-      // Insert diff segments not in matchedIds (e.g., loader segments from consolidation fetch)
-      insertMissingDiffSegments(allSegments, diff, matchedIdSet, newSegmentMap);
+      const actor: ReconcileActor =
+        mode.type === "stale-revalidation" || mode.type === "action"
+          ? "stale-revalidation"
+          : "navigation";
+
+      const reconciled = reconcileSegments({
+        actor,
+        matched: matchedIds,
+        diff: diff || [],
+        serverSegments: newSegments || [],
+        cachedSegments: cachedSegs,
+        insertMissingDiff: true,
+      });
 
       // HMR RESILIENCE: Check if we're missing any matched segments
-      // Note: allSegments may include additional diff segments, so we check matchedIds specifically
-      const allSegmentIdSet = new Set(allSegments.map((s) => s.id));
+      const reconciledIdSet = new Set(reconciled.segments.map((s) => s.id));
       const missingIds = matchedIds.filter(
-        (id: string) => !allSegmentIdSet.has(id),
+        (id: string) => !reconciledIdSet.has(id),
       );
 
       if (missingIds.length > 0) {
@@ -333,52 +342,39 @@ export function createPartialUpdater(
           );
         }
         if (signal?.aborted) {
-          console.log(
-            `[Browser] Ignoring stale navigation (aborted during HMR retry)`,
+          debugLog(
+            "[Browser] Ignoring stale navigation (aborted during HMR retry)",
           );
-          return streamComplete;
+          return;
         }
-        if (isAction) {
-          return streamComplete;
+        if (mode.type === "action") {
+          return;
         }
         console.warn(
           `[Browser] HMR detected: Missing ${missingCount} segments. Refetching all...`,
         );
 
         // Refetch with empty segments = server sends everything
-        return fetchPartialUpdate(url, [], true, signal, tx, { isAction });
+        return fetchPartialUpdate(url, [], true, signal, tx, mode);
       }
 
-      // INTERCEPT HANDLING: Separate intercept segments for explicit injection
-      // Intercept segments have namespace starting with "intercept:" or ID containing .@
-      // This makes the flow clearer and easier to debug
-      const isInterceptSegment = (s: ResolvedSegment) =>
-        s.namespace?.startsWith("intercept:") ||
-        (s.type === "parallel" && s.id.includes(".@"));
-
-      const interceptSegments = allSegments.filter(isInterceptSegment);
-      const mainSegments = allSegments.filter((s) => !isInterceptSegment(s));
-
       if (signal?.aborted) {
-        console.log(
-          `[Browser] Ignoring stale navigation (aborted before render)`,
-        );
-        return streamComplete;
+        debugLog("[Browser] Ignoring stale navigation (aborted before render)");
+        return;
       }
 
       // Rebuild tree on client (await for loader data resolution)
-      // Race against abort signal to allow cancellation during loader awaiting
-      // Pass intercept segments separately for explicit handling
-      // For stale revalidation, use forceAwait to ensure no loading fallbacks
       const renderOptions = {
-        isAction,
-        forceAwait: staleRevalidation,
+        isAction: mode.type === "action",
+        forceAwait: mode.type === "stale-revalidation",
         interceptSegments:
-          interceptSegments.length > 0 ? interceptSegments : undefined,
+          reconciled.interceptSegments.length > 0
+            ? reconciled.interceptSegments
+            : undefined,
       };
       const newTree = await (signal
         ? Promise.race([
-            renderSegments(mainSegments, renderOptions),
+            renderSegments(reconciled.mainSegments, renderOptions),
             new Promise<never>((_, reject) => {
               if (signal.aborted) {
                 reject(new DOMException("Navigation aborted", "AbortError"));
@@ -388,85 +384,79 @@ export function createPartialUpdater(
               });
             }),
           ])
-        : renderSegments(mainSegments, renderOptions));
+        : renderSegments(reconciled.mainSegments, renderOptions));
 
       // Final abort check before committing - another navigation may have started
       if (signal?.aborted) {
-        console.log(
-          `[Browser] Ignoring stale navigation (aborted before commit)`,
-        );
-        return streamComplete;
+        debugLog("[Browser] Ignoring stale navigation (aborted before commit)");
+        return;
       }
 
       // Check if this is an intercept response (any slot is active)
-      // If so, disable scroll to keep the current scroll position
-      const hasActiveIntercept = payload.metadata?.slots
-        ? Object.values(payload.metadata.slots).some((slot) => slot.active)
-        : false;
-
-      // BUG FIX: When navigating with cached target segments but receiving an intercept response,
-      // the background segments should come from the SOURCE page (where we navigated from),
-      // not the TARGET cache. This happens when:
-      // 1. User visits /product/xxx (detail page) - cached under key "/product/xxx"
-      // 2. User navigates back to /
-      // 3. User clicks product link → cache hit for "/product/xxx" (detail page)
-      // 4. But server returns intercept response (modal with index background)
-      // 5. Without this fix: background uses detail page segments (wrong!)
-      // 6. With this fix: rebuild currentSegmentMap from source page
-      if (hasActiveIntercept && targetCacheSegments) {
-        console.log(
-          `[Browser] Intercept response with target cache - rebuilding segment map from source page`,
-        );
-        currentSegmentMap = getCurrentSegmentMap();
-      }
+      const isInterceptResponse = hasActiveInterceptSlots(
+        payload.metadata?.slots,
+      );
 
-      // Track intercept context for action revalidation (only on navigation, not actions or stale revalidation)
-      if (!isAction && !staleRevalidation) {
-        if (hasActiveIntercept) {
-          // Save the source URL for action revalidation to maintain intercept context
-          store.setInterceptSourceUrl(segmentState.currentUrl);
+      // Track intercept context (only on navigation, not actions or stale revalidation)
+      // Use the authoritative source from mode/history state when restoring an
+      // intercept via popstate cache miss; fall back to the current URL for fresh
+      // intercept navigations.
+      const effectiveInterceptSource =
+        interceptSourceUrl || segmentState.currentUrl;
+      if (mode.type !== "action" && mode.type !== "stale-revalidation") {
+        if (isInterceptResponse) {
+          store.setInterceptSourceUrl(effectiveInterceptSource);
         } else {
-          // Clear intercept context when navigating to a non-intercept route
           store.setInterceptSourceUrl(null);
         }
       }
 
       // Commit navigation - transaction handles all store mutations atomically
-      // For intercept responses: disable scroll, mark as intercept, include source URL
-      // Use allSegmentIds (derived from allSegments) instead of matchedIds because
-      // we may have added diff segments (like loader segments) not in the matched array
-      const allSegmentIds = allSegments.map((s) => s.id);
-      tx.commit(
-        allSegmentIds,
-        allSegments,
-        hasActiveIntercept
-          ? {
-              scroll: false,
-              intercept: true,
-              interceptSourceUrl: segmentState.currentUrl,
-            }
-          : undefined,
-      );
+      const allSegmentIds = reconciled.segments.map((s) => s.id);
+      const serverLocationState = payload.metadata?.locationState;
+      const overrides: CommitOverrides | undefined = isInterceptResponse
+        ? {
+            scroll: false,
+            intercept: true,
+            interceptSourceUrl: effectiveInterceptSource,
+            ...(serverLocationState && { serverState: serverLocationState }),
+          }
+        : serverLocationState
+          ? { serverState: serverLocationState }
+          : undefined;
+      tx.commit(allSegmentIds, reconciled.segments, overrides);
 
       // For stale revalidation: verify history key hasn't changed before updating UI
-      // If user navigated away, skip UI update to avoid corrupting current view
-      if (staleRevalidation) {
+      if (mode.type === "stale-revalidation") {
         const historyKeyNow = store.getHistoryKey();
         if (historyKeyNow !== historyKeyAtStart) {
-          console.log(
+          debugLog(
             `[Browser] Stale revalidation: history key changed (${historyKeyAtStart} -> ${historyKeyNow}), skipping UI update`,
           );
-          return streamComplete;
+          return;
         }
       }
 
-      console.log("[partial-update] updating document");
+      debugLog("[partial-update] updating document");
 
       // Emit update to trigger React render
-      // For stale revalidation: wait for stream to complete (loaders resolved), then update
-      // For actions: wrap in startTransition to avoid UI flickering
-      if (isAction || staleRevalidation) {
+      const hasTransition = reconciled.mainSegments.some((s) => s.transition);
+
+      if (mode.type === "action" || mode.type === "stale-revalidation") {
+        startTransition(() => {
+          if (hasTransition && addTransitionType) {
+            addTransitionType("action");
+          }
+          onUpdate({
+            root: newTree,
+            metadata: payload.metadata!,
+          });
+        });
+      } else if (hasTransition) {
         startTransition(() => {
+          if (addTransitionType) {
+            addTransitionType("navigation");
+          }
           onUpdate({
             root: newTree,
             metadata: payload.metadata!,
@@ -479,77 +469,78 @@ export function createPartialUpdater(
         });
       }
 
-      console.log(`[Browser] Navigation complete\n`);
-      return streamComplete;
+      debugLog("[Browser] Navigation complete");
+      return;
     } else {
       // Full update (fallback)
       console.warn(`[Browser] Full update (fallback)`);
 
       const segments = payload.metadata?.segments || [];
 
-      // Check if this navigation is stale (a newer one started)
       if (signal?.aborted) {
-        console.log(`[Browser] Ignoring stale navigation (aborted)`);
-        return streamComplete;
-      }
-
-      // Await loader data from segments before committing URL
-      // This ensures URL only updates after loaders resolve
-      const loaderSegments = segments.filter(
-        (s: ResolvedSegment) =>
-          s.type === "loader" && s.loaderData !== undefined,
-      );
-      if (loaderSegments.length > 0) {
-        console.log(`[Browser] Awaiting ${loaderSegments.length} loader(s)...`);
-        await Promise.all(
-          loaderSegments.map((s: ResolvedSegment) =>
-            s.loaderData instanceof Promise
-              ? s.loaderData
-              : Promise.resolve(s.loaderData),
-          ),
-        );
-        console.log(`[Browser] Loaders resolved`);
+        debugLog("[Browser] Ignoring stale navigation (aborted)");
+        return;
       }
 
       const segmentIds = segments.map((s: ResolvedSegment) => s.id);
 
-      // Final abort check before committing - another navigation may have started
+      const newTree = await renderSegments(segments);
+
       if (signal?.aborted) {
-        console.log(
-          `[Browser] Ignoring stale navigation (aborted before commit)`,
-        );
-        return streamComplete;
+        debugLog("[Browser] Ignoring stale navigation (aborted before commit)");
+        return;
       }
 
-      // Commit navigation - transaction handles all store mutations atomically
-      tx.commit(segmentIds, segments);
+      const fullUpdateServerState = payload.metadata?.locationState;
+      if (fullUpdateServerState) {
+        tx.commit(segmentIds, segments, { serverState: fullUpdateServerState });
+      } else {
+        tx.commit(segmentIds, segments);
+      }
 
-      // Emit update to trigger React render
-      // For stale revalidation: wait for stream to complete, then update
-      // For actions: wrap in startTransition to avoid UI flickering
-      if (staleRevalidation) {
+      const fullHasTransition = segments.some(
+        (s: ResolvedSegment) => s.transition,
+      );
+
+      if (mode.type === "stale-revalidation") {
         await rawStreamComplete;
         startTransition(() => {
+          if (fullHasTransition && addTransitionType) {
+            addTransitionType("action");
+          }
           onUpdate({
-            root: payload.root,
+            root: newTree,
             metadata: payload.metadata!,
           });
         });
-      } else if (isAction) {
+      } else if (mode.type === "action") {
         startTransition(async () => {
+          if (fullHasTransition && addTransitionType) {
+            addTransitionType("action");
+          }
+          onUpdate({
+            root: newTree,
+            metadata: payload.metadata!,
+          });
+        });
+      } else if (fullHasTransition) {
+        startTransition(() => {
+          if (addTransitionType) {
+            addTransitionType("navigation");
+          }
           onUpdate({
-            root: payload.root,
+            root: newTree,
             metadata: payload.metadata!,
           });
         });
       } else {
         onUpdate({
-          root: payload.root,
+          root: newTree,
           metadata: payload.metadata!,
         });
       }
 
-      return streamComplete;
+      return;
     }
   }