From 739da33c638e627f3334355c1f3633d913b7b198 Mon Sep 17 00:00:00 2001
From: Tobbe Lundberg <tobbe@tlundberg.com>
Date: Fri, 24 May 2024 15:41:19 +0200
Subject: [PATCH] RSC: client side navigation (#10684)

---
 packages/vite/src/ClientRouter.tsx | 33 ++++++++++++++----------------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/packages/vite/src/ClientRouter.tsx b/packages/vite/src/ClientRouter.tsx
index 9a3b46628958..e114aaa205c8 100644
--- a/packages/vite/src/ClientRouter.tsx
+++ b/packages/vite/src/ClientRouter.tsx
@@ -3,7 +3,7 @@
 // what to do about rscFetch here vs renderFromRscServer and see if maybe that
 // one should live somewhere else where @redwoodjs/router can import from
 
-import React, { useMemo } from 'react'
+import React, { useEffect, useMemo } from 'react'
 
 import type { Options } from 'react-server-dom-webpack/client'
 import { createFromFetch, encodeReply } from 'react-server-dom-webpack/client'
@@ -62,8 +62,6 @@ function rscFetch(rscId: string, props: Record<string, unknown> = {}) {
   return createFromFetch<never, React.ReactElement>(response, options)
 }
 
-let routes: Thenable<React.ReactElement> | null = null
-
 export const Router = ({ paramTypes, children }: RouterProps) => {
   return (
     // Wrap it in the provider so that useLocation can be used
@@ -75,33 +73,32 @@ export const Router = ({ paramTypes, children }: RouterProps) => {
   )
 }
 
+let routes: Thenable<React.ReactElement> | null = null
+
 const LocationAwareRouter = ({ paramTypes, children }: RouterProps) => {
-  const location = useLocation()
+  const { pathname, search } = useLocation()
+  const [, setRenderCount] = React.useState(0)
 
   const { namedRoutesMap } = useMemo(() => {
     return analyzeRoutes(children, {
-      currentPathName: location.pathname,
-      // @TODO We haven't handled this with SSR/Streaming yet.
-      // May need a babel plugin to extract userParamTypes from Routes.tsx
+      currentPathName: pathname,
       userParamTypes: paramTypes,
     })
-  }, [location.pathname, children, paramTypes])
+  }, [pathname, children, paramTypes])
 
   // Assign namedRoutes so it can be imported like import {routes} from 'rwjs/router'
   // Note that the value changes at runtime
   Object.assign(namedRoutes, namedRoutesMap)
 
-  // TODO (RSC): Refetch when the location changes
-  // It currently works because we always do a full page refresh, but that's
-  // not what we really want to do)
+  useEffect(() => {
+    // Force re-render
+    setRenderCount((rc) => rc + 1)
+
+    routes = rscFetch('__rwjs__Routes', { location: { pathname, search } })
+  }, [pathname, search])
+
   if (!routes) {
-    routes = rscFetch('__rwjs__Routes', {
-      // All we need right now is the pathname. Plus, `location` is a URL
-      // object, and it doesn't JSON.stringify well. Basically all you end up
-      // with is the href. That's why we manually construct the object here
-      // instead of just passing `location`.
-      location: { pathname: location.pathname },
-    })
+    routes = rscFetch('__rwjs__Routes', { location: { pathname, search } })
   }
 
   return routes