Fix issue with surfacing entire state through use optimistic

Author: brophdawg11

packages/react-router/lib/components.tsx

@@ -86,13 +86,14 @@ const USE_OPTIMISTIC = "useOptimistic";
 // @ts-expect-error Needs React 19 types but we develop against 18
 const useOptimisticImpl = React[USE_OPTIMISTIC];
 
-function useOptimisticSafe(val: unknown) {
+function useOptimisticSafe<T>(
+  val: T,
+): [T, React.Dispatch<React.SetStateAction<T>>] {
   if (useOptimisticImpl) {
     // eslint-disable-next-line react-hooks/rules-of-hooks
     return useOptimisticImpl(val);
   } else {
-    // eslint-disable-next-line react-hooks/rules-of-hooks
-    return React.useState(val);
+    return [val, () => undefined];
   }
 }
 
@@ -645,7 +646,7 @@ export function RouterProvider({
         } else {
           React.startTransition(() => {
             if (unstable_transitions === true) {
-              setOptimisticState(newState);
+              setOptimisticState((s) => getOptimisticRouterState(s, newState));
             }
             logErrorsAndSetState(newState);
           });
@@ -747,7 +748,7 @@ export function RouterProvider({
         } else {
           React.startTransition(() => {
             if (unstable_transitions === true) {
-              setOptimisticState(newState);
+              setOptimisticState((s) => getOptimisticRouterState(s, newState));
             }
             logErrorsAndSetState(newState);
           });
@@ -865,6 +866,33 @@ export function RouterProvider({
   );
 }
 
+function getOptimisticRouterState(
+  currentState: RouterState,
+  newState: RouterState,
+): RouterState {
+  return {
+    // Don't surface "current location specific" stuff mid-navigation
+    // (historyAction, location, matches, loaderData, errors, initialized,
+    // restoreScroll, preventScrollReset, blockers, etc.)
+    ...currentState,
+    // Only surface "pending/in-flight stuff"
+    // (navigation, revalidation, actionData, fetchers, )
+    navigation:
+      newState.navigation.state !== "idle"
+        ? newState.navigation
+        : currentState.navigation,
+    revalidation:
+      newState.revalidation !== "idle"
+        ? newState.revalidation
+        : currentState.revalidation,
+    actionData:
+      newState.navigation.state !== "submitting"
+        ? newState.actionData
+        : currentState.actionData,
+    fetchers: newState.fetchers,
+  };
+}
+
 // Memoize to avoid re-renders when updating `ViewTransitionContext`
 const MemoizedDataRoutes = React.memo(DataRoutes);
 

playground/react-transitions/app/routes/transitions.tsx

@@ -80,61 +80,59 @@ export default function Transitions() {
               </button>
               <ul>
                 <li>
-                  In the current state, <code>useNavigate</code> navigations are
-                  not wrapped in <code>startTransition</code>, so they don't
-                  play nice with other transition-aware state updates
+                  In the current state, <code>useNavigate</code> doesn't wrap
+                  the navigation in <code>startTransition</code>, so they don't
+                  play nice with other transition-aware state updates (try
+                  updating the transition-aware counter mid-navigation)
                 </li>
                 <li>
-                  Fixed by{" "}
-                  <code>
-                    &lt;HydrateRouter unstable_transitions={"{"}true{"}"} /&gt;
-                  </code>
+                  With the new flag, they are wrapped in{" "}
+                  <code>startTransition</code>
                 </li>
               </ul>
             </li>
 
             <li>
               <button
                 onClick={() =>
-                  navigate("/transitions/slow", { flushSync: true })
+                  // @ts-expect-error Needs React 19 types
+                  startTransition(() => navigate("/transitions/slow"))
                 }
               >
                 <code>
-                  navigate("/transitions/slow", {"{"} flushSync: true {"}"})
+                  startTransition(() =&gt; navigate("/transitions/slow")
                 </code>
               </button>
               <ul>
                 <li>
-                  With the new flag, useNavigate automatically wraps the
-                  navigation in <code>React.startTransition</code>. Passing the{" "}
-                  <code>flushSync</code> option will opt out of that and apply
-                  <code>React.flushSync</code> to the underlying state update
+                  If you wrap them in <code>startTransition</code> manually,
+                  they play nicely with those updates but they prevent our
+                  internal mid-navigation state updates from surfacing
+                </li>
+                <li>
+                  That can be fixed by enabling <code>useOptimistic</code>{" "}
+                  inside <code>&lt;RouterProvider&gt;</code>
                 </li>
               </ul>
             </li>
 
             <li>
               <button
                 onClick={() =>
-                  // @ts-expect-error Needs React 19 types
-                  startTransition(() => navigate("/transitions/slow"))
+                  navigate("/transitions/slow", { flushSync: true })
                 }
               >
                 <code>
-                  startTransition(() =&gt; navigate("/transitions/slow")
+                  navigate("/transitions/slow", {"{"} flushSync: true {"}"})
                 </code>
               </button>
               <ul>
                 <li>
-                  Once you wrap them in <code>startTransition</code>, they play
-                  nicely with those updates but they prevent our internal
-                  mid-navigation state updates from surfacing
-                </li>
-                <li>
-                  Fixed by{" "}
-                  <code>
-                    &lt;HydrateRouter unstable_transitions={"{"}true{"}"} /&gt;
-                  </code>
+                  Once <code>useNavigate</code> is wrapped automatically,
+                  passing the
+                  <code>flushSync</code> option will opt out of{" "}
+                  <code>startTransition</code> and apply
+                  <code>React.flushSync</code> to the underlying state update
                 </li>
               </ul>
             </li>
@@ -146,25 +144,36 @@ export default function Transitions() {
               <ul>
                 <li>
                   In the current state, <code>&lt;Link&gt;</code> navigations
-                  are not wrapped in startTransition, so they don't play nice
-                  with other transition-aware state updates
+                  are not wrapped in <code>startTransition</code>, so they don't
+                  play nice with other transition-aware state updates
                 </li>
                 <li>
-                  Fixed by{" "}
-                  <code>
-                    &lt;HydrateRouter unstable_transitions={"{"}true{"}"} /&gt;
-                  </code>
+                  With the new flag, they are wrapped in{" "}
+                  <code>startTransition</code>
                 </li>
               </ul>
             </li>
 
             <li>
-              <Link to="/transitions/parent">/transitions/parent</Link>
-            </li>
-            <li>
-              <Link to="/transitions/parent/child">
-                /transitions/parent/child
-              </Link>
+              Nested Parent/Child Await:
+              <ul>
+                <li>
+                  Should not re-fallback when going from parent -&gt; child
+                </li>
+                <li>
+                  <span style={{ color: "red" }}>
+                    But for some reason they are when we enable the flag?
+                  </span>
+                </li>
+                <li>
+                  <Link to="/transitions/parent">/transitions/parent</Link>
+                </li>
+                <li>
+                  <Link to="/transitions/parent/child">
+                    /transitions/parent/child
+                  </Link>
+                </li>
+              </ul>
             </li>
           </ul>
         </div>