@rangojs/router 0.0.0-experimental.259 → 0.0.0-experimental.25dbaac3

package/src/browser/partial-update.ts

@@ -15,9 +15,18 @@ import type { RenderSegmentsOptions } from "../segment-system.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-bridge.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";
+import type { NavigationUpdate } from "./types.js";
+
+/** Build a scroll payload from the commit's scroll option */
+function toScrollPayload(
+  scroll: boolean | undefined,
+): NonNullable<NavigationUpdate["scroll"]> {
+  return { enabled: scroll !== false ? scroll : false };
+}
 
 /**
  * Configuration for creating a partial updater
@@ -30,8 +39,15 @@ export interface PartialUpdateConfig {
     segments: ResolvedSegment[],
     options?: RenderSegmentsOptions,
   ) => Promise<ReactNode> | ReactNode;
-  /** RSC version received from server (from initial payload metadata) */
-  version?: string;
+  /** RSC version getter — returns the current version (may change after HMR) */
+  getVersion?: () => string | undefined;
+  /**
+   * Replace the active app-shell when a cross-app navigation is detected.
+   * Called before the full-update tree replacement renders, so the new
+   * payload's rootLayout, basename, and version are picked up. Theme,
+   * warmup, and prefetch TTL are not part of the shell — see AppShell.
+   */
+  applyAppShell?: (next: import("./app-shell.js").AppShell) => void;
 }
 
 /**
@@ -51,9 +67,25 @@ export interface CommitOverrides {
 }
 
 /**
- * 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
@@ -63,24 +95,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[]>>;
-    /** When true, we're leaving an intercept state - don't use current segment IDs
-     * as fallback and force a fresh render from server */
-    leavingIntercept?: boolean;
-  },
-) => Promise<Promise<void>>;
+  tx: BoundTransaction,
+  mode?: UpdateMode,
+) => Promise<void>;
 
 /**
  * Create a partial updater for fetching and applying RSC partial updates
@@ -90,23 +107,18 @@ 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,
 ): PartialUpdater {
-  const { store, client, onUpdate, renderSegments, version } = config;
+  const {
+    store,
+    client,
+    onUpdate,
+    renderSegments,
+    getVersion = () => undefined,
+    applyAppShell,
+  } = config;
 
   /**
    * Get current page's cached segments as an array
@@ -119,7 +131,6 @@ export function createPartialUpdater(
 
   /**
    * 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)
@@ -130,39 +141,26 @@ 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[]>>;
-      leavingIntercept?: boolean;
-    },
-  ): Promise<Promise<void>> {
-    const {
-      isAction = false,
-      staleRevalidation = false,
-      interceptSourceUrl,
-      targetCacheSegments,
-      targetCacheHandleData,
-      leavingIntercept = false,
-    } = 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();
 
-    // When leaving intercept, don't send current segment IDs - we need fresh non-intercept segments
-    // Filter out intercept-related segments (parallel slots like @modal) from current segments
+    // 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 (leavingIntercept) {
-      // When leaving intercept, only send segments that aren't intercept-specific
-      // The server will return the non-intercept version of the route
+    if (mode.type === "leave-intercept") {
       const currentSegments = segmentIds ?? segmentState.currentSegmentIds;
-      // Use cached segment metadata (namespace) to identify intercept segments.
-      // Only intercept segments have namespace starting with "intercept:" —
-      // regular parallel segments like @sidebar are preserved.
       const currentCached = getCurrentCachedSegments();
       const interceptIds = new Set(
         currentCached
@@ -177,10 +175,16 @@ export function createPartialUpdater(
       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
+    // For intercept revalidation, use the intercept source URL as previousUrl.
+    // For leave-intercept, tx.currentUrl captures window.location.href at tx
+    // creation, which on popstate is already the destination URL and would
+    // tell the server "from == to". segmentState.currentUrl still points at
+    // the URL the cached segments render (the intercept URL), which is the
+    // correct "from" for the server's diff computation.
     const previousUrl =
-      interceptSourceUrl || tx.currentUrl || segmentState.currentUrl;
+      mode.type === "leave-intercept"
+        ? segmentState.currentUrl || tx.currentUrl
+        : interceptSourceUrl || tx.currentUrl || segmentState.currentUrl;
 
     debugLog(`\n[Browser] >>> NAVIGATION`);
     debugLog(`[Browser] From: ${previousUrl}`);
@@ -191,36 +195,37 @@ export function createPartialUpdater(
     }
 
     // Get cached segments 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.
+    // When navigating with targetCacheSegments, use those for consistency.
     // Otherwise fall back to current page's segments (for same-route revalidation).
+    const targetCache =
+      mode.type === "navigate" ? mode.targetCacheSegments : undefined;
     const cachedSegs =
-      targetCacheSegments && targetCacheSegments.length > 0
-        ? targetCacheSegments
+      targetCache && targetCache.length > 0
+        ? targetCache
         : getCurrentCachedSegments();
-    // Mark navigation as streaming (response received, now parsing RSC)
-    // The token is ended when the stream completes
-    const streamingToken = tx.startStreaming();
+    const cachedSegsSource =
+      targetCache && targetCache.length > 0 ? "history-cache" : "current-page";
+    debugLog(
+      `[Browser] cachedSegs source: ${cachedSegsSource} (${cachedSegs.length} segments: ${cachedSegs.map((s) => s.id).join(", ")})`,
+    );
+
     // Fetch partial payload (no abort signal - RSC doesn't support it well)
-    // Wrapped in try/catch to ensure streamingToken.end() is called if fetch throws,
-    // preventing isStreaming from being permanently stuck as true.
     let fetchResult: Awaited<ReturnType<NavigationClient["fetchPartial"]>>;
-    try {
-      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).
-        // Only the fetch URL param is affected here — behavioral side effects
-        // (forceAwait, history key check) are controlled by the staleRevalidation variable.
-        staleRevalidation: staleRevalidation || segments.length === 0,
-        version,
-      });
-    } catch (err) {
-      streamingToken.end();
-      throw err;
-    }
+    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: getVersion(),
+      routerId: store.getRouterId?.(),
+    });
+    // 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);
 
@@ -228,17 +233,46 @@ export function createPartialUpdater(
       streamingToken.end();
     });
 
-    // Handle server-side redirect with state: the server returned a 200 with
-    // a redirect payload instead of a 3xx so that location state is preserved.
-    // Throw ServerRedirect to let navigate() catch it and re-navigate with state.
-    // Check signal.aborted first — a newer navigation may have started, and we
-    // must not redirect from a stale response.
+    // Detect app switch: if routerId changed, the navigation crossed into
+    // a different router (e.g., via host router path mount). Downgrade
+    // partial to full so the entire tree is replaced without reconciliation
+    // against stale segments from the previous app, and replace the app
+    // shell (rootLayout, basename, version) so the target app's document
+    // and router config take effect instead of remaining captured from the
+    // initial load. Theme, warmup, and prefetch TTL are intentionally
+    // document-lifetime (see AppShell doc); a new document navigation
+    // applies them.
+    if (payload.metadata?.routerId) {
+      const prevRouterId = store.getRouterId?.();
+      if (prevRouterId && prevRouterId !== payload.metadata.routerId) {
+        debugLog(
+          `[Browser] App switch detected (${prevRouterId} → ${payload.metadata.routerId}), forcing full update`,
+        );
+        payload.metadata.isPartial = false;
+        applyAppShell?.({
+          routerId: payload.metadata.routerId,
+          rootLayout: payload.metadata.rootLayout,
+          basename: payload.metadata.basename,
+          version: payload.metadata.version,
+        });
+      }
+      store.setRouterId?.(payload.metadata.routerId);
+    }
+
+    // Handle server-side redirect with state
     if (payload.metadata?.redirect) {
       if (signal?.aborted) {
-        console.log(`[Browser] Ignoring stale redirect (aborted)`);
-        return streamComplete;
+        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 { url: redirectUrl } = payload.metadata.redirect;
       const serverState = payload.metadata.locationState;
       throw new ServerRedirect(redirectUrl, serverState);
     }
@@ -249,15 +283,13 @@ export function createPartialUpdater(
       // Check if this navigation is stale (a newer one started)
       if (signal?.aborted) {
         debugLog("[Browser] Ignoring stale navigation (aborted)");
-        return streamComplete;
+        return;
       }
 
       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]));
@@ -266,8 +298,7 @@ export function createPartialUpdater(
           .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) {
+        if (mode.type === "navigate" && targetCache && targetCache.length > 0) {
           debugLog(
             "[Browser] No diff but navigating with cached segments - rendering target route",
           );
@@ -276,21 +307,35 @@ export function createPartialUpdater(
             forceAwait: true,
           });
 
-          tx.commit(matchedIds, existingSegments);
+          const { scroll: commitScroll } = tx.commit(
+            matchedIds,
+            existingSegments,
+          );
+
+          // tx.commit() cached the source page's handleData because
+          // eventController hasn't been updated yet. Overwrite with the
+          // correct cached handleData to prevent cache corruption on
+          // subsequent navigations to this same URL.
+          if (mode.targetCacheHandleData) {
+            store.updateCacheHandleData(
+              store.getHistoryKey(),
+              mode.targetCacheHandleData,
+            );
+          }
 
           // 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!;
           const cachedUpdate = {
             root: newTree,
             metadata: {
               ...metadataWithoutHandles,
-              cachedHandleData: targetCacheHandleData,
+              cachedHandleData: mode.targetCacheHandleData,
             },
+            scroll: toScrollPayload(commitScroll),
           };
 
           const cachedHasTransition = existingSegments.some(
@@ -308,13 +353,11 @@ export function createPartialUpdater(
           }
 
           debugLog("[Browser] Navigation complete (rendered from cache)");
-          return streamComplete;
+          return;
         }
 
         // When leaving intercept, force re-render even with empty diff
-        // The matched segments are the non-intercept segments, which we need to render
-        // to remove the modal from the UI
-        if (leavingIntercept) {
+        if (mode.type === "leave-intercept") {
           debugLog(
             "[Browser] Leaving intercept - forcing re-render to remove modal",
           );
@@ -323,15 +366,19 @@ export function createPartialUpdater(
             forceAwait: true,
           });
 
-          tx.commit(matchedIds, existingSegments);
+          const { scroll: leaveScroll } = tx.commit(
+            matchedIds,
+            existingSegments,
+          );
 
           onUpdate({
             root: newTree,
             metadata: payload.metadata,
+            scroll: toScrollPayload(leaveScroll),
           });
 
           debugLog("[Browser] Navigation complete (left intercept)");
-          return streamComplete;
+          return;
         }
 
         // Same route revalidation with no changes - skip UI update
@@ -340,13 +387,15 @@ export function createPartialUpdater(
         );
         tx.commit(matchedIds, existingSegments);
         debugLog("[Browser] Navigation complete (no re-render)");
-        return streamComplete;
+        return;
       }
 
       // Reconcile server segments with cached segments (single source of truth)
       const matchedIds = matched || [];
       const actor: ReconcileActor =
-        staleRevalidation || isAction ? "stale-revalidation" : "navigation";
+        mode.type === "stale-revalidation" || mode.type === "action"
+          ? "stale-revalidation"
+          : "navigation";
 
       const reconciled = reconcileSegments({
         actor,
@@ -376,31 +425,28 @@ export function createPartialUpdater(
           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);
       }
 
       if (signal?.aborted) {
         debugLog("[Browser] Ignoring stale navigation (aborted before render)");
-        return streamComplete;
+        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:
           reconciled.interceptSegments.length > 0
             ? reconciled.interceptSegments
@@ -423,64 +469,68 @@ export function createPartialUpdater(
       // Final abort check before committing - another navigation may have started
       if (signal?.aborted) {
         debugLog("[Browser] Ignoring stale navigation (aborted before commit)");
-        return streamComplete;
+        return;
       }
 
       // Check if this is an intercept response (any slot is active)
-      // If so, disable scroll to keep the current scroll position
       const isInterceptResponse = hasActiveInterceptSlots(
         payload.metadata?.slots,
       );
 
-      // Track intercept context for action revalidation (only on navigation, not actions or stale revalidation)
-      if (!isAction && !staleRevalidation) {
+      // 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) {
-          // Save the source URL for action revalidation to maintain intercept context
-          store.setInterceptSourceUrl(segmentState.currentUrl);
+          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 reconciled.segments (which includes inserted diff segments) instead of matchedIds
-      const allSegmentIds = reconciled.segments.map((s) => s.id);
+      // Commit navigation - use server's matched as the authoritative segment ID list.
+      // reconciled.segments may be missing IDs (e.g., loader segments not in diff or cache)
+      // but the server's matched always includes all expected segment IDs.
+      const allSegmentIds = matchedIds;
       const serverLocationState = payload.metadata?.locationState;
       const overrides: CommitOverrides | undefined = isInterceptResponse
         ? {
             scroll: false,
             intercept: true,
-            interceptSourceUrl: segmentState.currentUrl,
+            interceptSourceUrl: effectiveInterceptSource,
             ...(serverLocationState && { serverState: serverLocationState }),
           }
         : serverLocationState
           ? { serverState: serverLocationState }
           : undefined;
-      tx.commit(allSegmentIds, reconciled.segments, overrides);
+      const { scroll: navScroll } = 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) {
           debugLog(
             `[Browser] Stale revalidation: history key changed (${historyKeyAtStart} -> ${historyKeyNow}), skipping UI update`,
           );
-          return streamComplete;
+          return;
         }
       }
 
       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
-      // For transitions: wrap in startTransition + addTransitionType for ViewTransition
+      // Emit update to trigger React render.
+      // Scroll info is included so NavigationProvider applies it after React commits.
       const hasTransition = reconciled.mainSegments.some((s) => s.transition);
+      const scrollPayload = toScrollPayload(navScroll);
 
-      if (isAction || staleRevalidation) {
+      if (mode.type === "action" || mode.type === "stale-revalidation") {
         startTransition(() => {
           if (hasTransition && addTransitionType) {
             addTransitionType("action");
@@ -488,6 +538,7 @@ export function createPartialUpdater(
           onUpdate({
             root: newTree,
             metadata: payload.metadata!,
+            scroll: scrollPayload,
           });
         });
       } else if (hasTransition) {
@@ -498,59 +549,52 @@ export function createPartialUpdater(
           onUpdate({
             root: newTree,
             metadata: payload.metadata!,
+            scroll: scrollPayload,
           });
         });
       } else {
         onUpdate({
           root: newTree,
           metadata: payload.metadata!,
+          scroll: scrollPayload,
         });
       }
 
       debugLog("[Browser] Navigation complete");
-      return streamComplete;
+      return;
     } else {
       // Full update (fallback)
-      // Reconstruct the tree client-side from segments via renderSegments
-      // to ensure consistent component references with action revalidation.
       console.warn(`[Browser] Full update (fallback)`);
 
       const segments = payload.metadata?.segments || [];
 
-      // Check if this navigation is stale (a newer one started)
       if (signal?.aborted) {
         debugLog("[Browser] Ignoring stale navigation (aborted)");
-        return streamComplete;
+        return;
       }
 
       const segmentIds = segments.map((s: ResolvedSegment) => s.id);
 
-      // Render on client for consistent component references
       const newTree = await renderSegments(segments);
 
-      // Final abort check before committing - another navigation may have started
       if (signal?.aborted) {
         debugLog("[Browser] Ignoring stale navigation (aborted before commit)");
-        return streamComplete;
+        return;
       }
 
-      // Commit navigation - transaction handles all store mutations atomically
       const fullUpdateServerState = payload.metadata?.locationState;
-      if (fullUpdateServerState) {
-        tx.commit(segmentIds, segments, { serverState: fullUpdateServerState });
-      } else {
-        tx.commit(segmentIds, segments);
-      }
+      const { scroll: fullScroll } = fullUpdateServerState
+        ? tx.commit(segmentIds, segments, {
+            serverState: fullUpdateServerState,
+          })
+        : 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
-      // For transitions: wrap in startTransition + addTransitionType
       const fullHasTransition = segments.some(
         (s: ResolvedSegment) => s.transition,
       );
+      const fullScrollPayload = toScrollPayload(fullScroll);
 
-      if (staleRevalidation) {
+      if (mode.type === "stale-revalidation") {
         await rawStreamComplete;
         startTransition(() => {
           if (fullHasTransition && addTransitionType) {
@@ -559,9 +603,10 @@ export function createPartialUpdater(
           onUpdate({
             root: newTree,
             metadata: payload.metadata!,
+            scroll: fullScrollPayload,
           });
         });
-      } else if (isAction) {
+      } else if (mode.type === "action") {
         startTransition(async () => {
           if (fullHasTransition && addTransitionType) {
             addTransitionType("action");
@@ -569,6 +614,7 @@ export function createPartialUpdater(
           onUpdate({
             root: newTree,
             metadata: payload.metadata!,
+            scroll: fullScrollPayload,
           });
         });
       } else if (fullHasTransition) {
@@ -579,16 +625,18 @@ export function createPartialUpdater(
           onUpdate({
             root: newTree,
             metadata: payload.metadata!,
+            scroll: fullScrollPayload,
           });
         });
       } else {
         onUpdate({
           root: newTree,
           metadata: payload.metadata!,
+          scroll: fullScrollPayload,
         });
       }
 
-      return streamComplete;
+      return;
     }
   }