From eae5d77d1136cb158eb88bce981bc94ffc8dab9b Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Fri, 23 Sep 2022 11:35:36 -0700
Subject: [PATCH 01/10] 1st attempt at porting PR #713

---
 .../src/ssr/browser/main.jsx                  | 28 ++++++------
 .../src/ssr/server/react-rendering.js         | 31 +++++++++----
 .../components/with-react-query/index.js      | 14 +++++-
 .../{device-context.js => contexts.js}        |  8 ++--
 .../src/ssr/universal/hooks.js                | 44 +++++++++++++++++++
 .../app/pages/home.tsx                        | 11 +++++
 6 files changed, 110 insertions(+), 26 deletions(-)
 rename packages/pwa-kit-react-sdk/src/ssr/universal/{device-context.js => contexts.js} (51%)
 create mode 100644 packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js

diff --git a/packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx b/packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx
index 3d5ec96a5e..6277c68ca1 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx
+++ b/packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx
@@ -8,7 +8,7 @@
 import React from 'react'
 import ReactDOM from 'react-dom'
 import {BrowserRouter as Router} from 'react-router-dom'
-import DeviceContext from '../universal/device-context'
+import {DeviceContext, ServerContext} from '../universal/contexts'
 import App from '../universal/components/_app'
 import {getAppConfig} from '../universal/compatibility'
 import Switch from '../universal/components/switch'
@@ -74,18 +74,20 @@ export const start = () => {
         .then(() => new Promise((resolve) => loadableReady(resolve)))
         .then(() => {
             ReactDOM.hydrate(
-                <Router>
-                    <DeviceContext.Provider value={{type: window.__DEVICE_TYPE__}}>
-                        <AppConfig locals={locals}>
-                            <Switch
-                                error={error}
-                                appState={window.__PRELOADED_STATE__}
-                                routes={routes}
-                                App={WrappedApp}
-                            />
-                        </AppConfig>
-                    </DeviceContext.Provider>
-                </Router>,
+                <ServerContext.Provider value={{}}>
+                    <Router>
+                        <DeviceContext.Provider value={{type: window.__DEVICE_TYPE__}}>
+                            <AppConfig locals={locals}>
+                                <Switch
+                                    error={error}
+                                    appState={window.__PRELOADED_STATE__}
+                                    routes={routes}
+                                    App={WrappedApp}
+                                />
+                            </AppConfig>
+                        </DeviceContext.Provider>
+                    </Router>
+                </ServerContext.Provider>,
                 rootEl,
                 () => {
                     window.__HYDRATING__ = false
diff --git a/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.js b/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.js
index a5f2ca2fee..cbe9ced494 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.js
@@ -18,7 +18,7 @@ import {StaticRouter as Router, matchPath} from 'react-router-dom'
 import serialize from 'serialize-javascript'
 
 import {getAssetUrl} from '../universal/utils'
-import DeviceContext from '../universal/device-context'
+import {DeviceContext, ServerContext} from '../universal/contexts'
 
 import Document from '../universal/components/_document'
 import App from '../universal/components/_app'
@@ -191,20 +191,33 @@ export const render = async (req, res, next) => {
     }
 }
 
-const OuterApp = ({res, error, App, appState, routes, routerContext, location, deviceType}) => {
+const OuterApp = ({
+    req,
+    res,
+    error,
+    App,
+    appState,
+    routes,
+    routerContext,
+    location,
+    deviceType
+}) => {
     const AppConfig = getAppConfig()
     return (
-        <Router location={location} context={routerContext}>
-            <DeviceContext.Provider value={{type: deviceType}}>
-                <AppConfig locals={res.locals}>
-                    <Switch error={error} appState={appState} routes={routes} App={App} />
-                </AppConfig>
-            </DeviceContext.Provider>
-        </Router>
+        <ServerContext.Provider value={{req, res}}>
+            <Router location={location} context={routerContext}>
+                <DeviceContext.Provider value={{type: deviceType}}>
+                    <AppConfig locals={res.locals}>
+                        <Switch error={error} appState={appState} routes={routes} App={App} />
+                    </AppConfig>
+                </DeviceContext.Provider>
+            </Router>
+        </ServerContext.Provider>
     )
 }
 
 OuterApp.propTypes = {
+    req: PropTypes.object,
     res: PropTypes.object,
     error: PropTypes.object,
     App: PropTypes.elementType,
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
index 8d992d13a5..6118815946 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
@@ -9,10 +9,22 @@ import {FetchStrategy} from '../fetch-strategy'
 import React from 'react'
 import {dehydrate, Hydrate, QueryClient, QueryClientProvider} from '@tanstack/react-query'
 import ssrPrepass from 'react-ssr-prepass'
+import {IsPrePassContext} from '../../contexts'
 
 const isServerSide = typeof window === 'undefined'
 const STATE_KEY = '__reactQuery'
 
+const WithPrepassContext = (Component) => {
+    const WrappedComponent = (props) => {
+        return (
+            <IsPrePassContext.Provider value={true}>
+                <Component {...props} />
+            </IsPrePassContext.Provider>
+        )
+    }
+    return WrappedComponent
+}
+
 export const withReactQuery = (Wrapped) => {
     /* istanbul ignore next */
     const wrappedComponentName = Wrapped.displayName || Wrapped.name
@@ -39,7 +51,7 @@ export const withReactQuery = (Wrapped) => {
             const queryClient = (res.locals.__queryClient =
                 res.locals.__queryClient || new QueryClient())
 
-            await ssrPrepass(appJSX)
+            await ssrPrepass(WithPrepassContext(appJSX))
 
             const queryCache = queryClient.getQueryCache()
             const queries = queryCache.getAll().filter((q) => q.options.enabled !== false)
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/device-context.js b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
similarity index 51%
rename from packages/pwa-kit-react-sdk/src/ssr/universal/device-context.js
rename to packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
index 80f96309dd..3859008380 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/device-context.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
@@ -1,9 +1,11 @@
 /*
- * Copyright (c) 2021, salesforce.com, inc.
+ * Copyright (c) 2022, Salesforce, Inc.
  * All rights reserved.
  * SPDX-License-Identifier: BSD-3-Clause
  * For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
  */
-
 import React from 'react'
-export default React.createContext()
+
+export const DeviceContext = React.createContext()
+export const ServerContext = React.createContext()
+export const IsPrePassContext = React.createContext(false)
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
new file mode 100644
index 0000000000..131ec01672
--- /dev/null
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2022, Salesforce, Inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ * For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+ */
+
+import {useContext} from 'react'
+import {IsPrePassContext, ServerContext} from './contexts'
+
+/**
+ * Server context
+ * @typedef {Object} ServerContext
+ * @property {Object} req - Request object
+ * @property {Object} res - Response object
+ */
+
+/**
+ * @callback serverCtxCallback
+ * @param {ServerContext} serverContext
+ */
+
+/**
+ * Execute the given function within the server context of a server-rendered page
+ *
+ * @param {serverCtxCallback} fn - Function will be called when the server has received the request but the response is not sent yet
+ * @example
+ * useServerContext(({req, res}) => { res.status(404) })
+ */
+export const useServerContext = (fn) => {
+    const serverContext = useContext(ServerContext)
+
+    const isOnServer = Boolean(serverContext.req)
+    const isPrePass = useIsPrePass()
+
+    if (isOnServer && !isPrePass) {
+        fn(serverContext)
+    }
+}
+
+/**
+ * @private
+ */
+export const useIsPrePass = () => useContext(IsPrePassContext)
diff --git a/packages/template-typescript-minimal/app/pages/home.tsx b/packages/template-typescript-minimal/app/pages/home.tsx
index cbdeffdba7..7b3e041f12 100644
--- a/packages/template-typescript-minimal/app/pages/home.tsx
+++ b/packages/template-typescript-minimal/app/pages/home.tsx
@@ -10,6 +10,8 @@ import {useQuery} from '@tanstack/react-query'
 import HelloTS from '../components/hello-typescript'
 import HelloJS from '../components/hello-javascript'
 
+import {useServerContext} from 'pwa-kit-react-sdk/ssr/universal/hooks'
+
 interface Props {
     value: number
 }
@@ -82,6 +84,8 @@ h1 {
 `
 
 const Home = ({value}: Props) => {
+    // TODO: why is this log only once? I thought twice.
+    console.log('--- FOO')
     const [counter, setCounter] = useState(0)
 
     useEffect(() => {
@@ -101,6 +105,13 @@ const Home = ({value}: Props) => {
             })
     )
 
+    useServerContext(({req, res}) => {
+        console.log('--- useServerContext')
+        if (query.error) {
+            res.status(404)
+        }
+    })
+
     return (
         <div>
             <style dangerouslySetInnerHTML={{__html: style}} />

From 2e379840c8b48e2b3d88f4ad117cbe62b04acd89 Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Mon, 26 Sep 2022 10:52:36 -0700
Subject: [PATCH 02/10] Clarify React component vs element

---
 .../components/with-react-query/index.js       | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
index 6118815946..f5b09131e3 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
@@ -14,17 +14,6 @@ import {IsPrePassContext} from '../../contexts'
 const isServerSide = typeof window === 'undefined'
 const STATE_KEY = '__reactQuery'
 
-const WithPrepassContext = (Component) => {
-    const WrappedComponent = (props) => {
-        return (
-            <IsPrePassContext.Provider value={true}>
-                <Component {...props} />
-            </IsPrePassContext.Provider>
-        )
-    }
-    return WrappedComponent
-}
-
 export const withReactQuery = (Wrapped) => {
     /* istanbul ignore next */
     const wrappedComponentName = Wrapped.displayName || Wrapped.name
@@ -51,7 +40,12 @@ export const withReactQuery = (Wrapped) => {
             const queryClient = (res.locals.__queryClient =
                 res.locals.__queryClient || new QueryClient())
 
-            await ssrPrepass(WithPrepassContext(appJSX))
+            const withPrepassContext = React.createElement(
+                IsPrePassContext.Provider,
+                {value: true},
+                appJSX
+            )
+            await ssrPrepass(withPrepassContext)
 
             const queryCache = queryClient.getQueryCache()
             const queries = queryCache.getAll().filter((q) => q.options.enabled !== false)

From 980a1474e48399368949b735d5b733b78990f5cc Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Mon, 26 Sep 2022 11:07:23 -0700
Subject: [PATCH 03/10] Simplify the hook interface

---
 .../src/ssr/universal/hooks.js                | 22 ++++++++-----------
 .../app/pages/home.tsx                        |  9 ++++----
 2 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
index 131ec01672..2506084538 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
@@ -13,28 +13,24 @@ import {IsPrePassContext, ServerContext} from './contexts'
  * @typedef {Object} ServerContext
  * @property {Object} req - Request object
  * @property {Object} res - Response object
+ * @property {boolean} isServerSide
  */
 
 /**
- * @callback serverCtxCallback
- * @param {ServerContext} serverContext
- */
-
-/**
- * Execute the given function within the server context of a server-rendered page
+ * Get the server context
+ * @returns {ServerContext} ServerContext object
  *
- * @param {serverCtxCallback} fn - Function will be called when the server has received the request but the response is not sent yet
  * @example
- * useServerContext(({req, res}) => { res.status(404) })
+ * const {res, isServerSide} = useServerContext()
+ * if (isServerSide && query.error) { res.status(404) }
  */
-export const useServerContext = (fn) => {
+export const useServerContext = () => {
     const serverContext = useContext(ServerContext)
-
-    const isOnServer = Boolean(serverContext.req)
     const isPrePass = useIsPrePass()
 
-    if (isOnServer && !isPrePass) {
-        fn(serverContext)
+    return {
+        ...serverContext,
+        isServerSide: Boolean(serverContext.req) && !isPrePass
     }
 }
 
diff --git a/packages/template-typescript-minimal/app/pages/home.tsx b/packages/template-typescript-minimal/app/pages/home.tsx
index 7b3e041f12..9571e11052 100644
--- a/packages/template-typescript-minimal/app/pages/home.tsx
+++ b/packages/template-typescript-minimal/app/pages/home.tsx
@@ -105,12 +105,11 @@ const Home = ({value}: Props) => {
             })
     )
 
-    useServerContext(({req, res}) => {
+    const {res, isServerSide} = useServerContext()
+    if (isServerSide && query.data) {
         console.log('--- useServerContext')
-        if (query.error) {
-            res.status(404)
-        }
-    })
+        res.status(404)
+    }
 
     return (
         <div>

From b6c8d53c193019757758eaa6bcb74924995cf41b Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Mon, 26 Sep 2022 11:24:06 -0700
Subject: [PATCH 04/10] Some clean-up

---
 .../ssr/universal/components/with-react-query/index.js | 10 +++-------
 .../template-typescript-minimal/app/pages/home.tsx     |  2 --
 2 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
index f5b09131e3..8fa22789ef 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
@@ -36,16 +36,12 @@ export const withReactQuery = (Wrapped) => {
         /**
          * @private
          */
-        static async doInitAppState({res, appJSX}) {
+        static async doInitAppState({res, appJSX: _appJSX}) {
             const queryClient = (res.locals.__queryClient =
                 res.locals.__queryClient || new QueryClient())
 
-            const withPrepassContext = React.createElement(
-                IsPrePassContext.Provider,
-                {value: true},
-                appJSX
-            )
-            await ssrPrepass(withPrepassContext)
+            const appJSX = React.createElement(IsPrePassContext.Provider, {value: true}, _appJSX)
+            await ssrPrepass(appJSX)
 
             const queryCache = queryClient.getQueryCache()
             const queries = queryCache.getAll().filter((q) => q.options.enabled !== false)
diff --git a/packages/template-typescript-minimal/app/pages/home.tsx b/packages/template-typescript-minimal/app/pages/home.tsx
index 9571e11052..ab8502829a 100644
--- a/packages/template-typescript-minimal/app/pages/home.tsx
+++ b/packages/template-typescript-minimal/app/pages/home.tsx
@@ -84,8 +84,6 @@ h1 {
 `
 
 const Home = ({value}: Props) => {
-    // TODO: why is this log only once? I thought twice.
-    console.log('--- FOO')
     const [counter, setCounter] = useState(0)
 
     useEffect(() => {

From d9aa427fdbd9feb19077f5782e3f7492272551ea Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Wed, 28 Sep 2022 18:22:02 -0700
Subject: [PATCH 05/10] Isolate mention of prepass to within `with-react-query`

---
 .../universal/components/with-react-query/index.js    |  6 +++---
 .../pwa-kit-react-sdk/src/ssr/universal/contexts.js   |  1 -
 packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js | 11 +++--------
 3 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
index 8fa22789ef..fd731c96fb 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
@@ -9,7 +9,6 @@ import {FetchStrategy} from '../fetch-strategy'
 import React from 'react'
 import {dehydrate, Hydrate, QueryClient, QueryClientProvider} from '@tanstack/react-query'
 import ssrPrepass from 'react-ssr-prepass'
-import {IsPrePassContext} from '../../contexts'
 
 const isServerSide = typeof window === 'undefined'
 const STATE_KEY = '__reactQuery'
@@ -36,12 +35,13 @@ export const withReactQuery = (Wrapped) => {
         /**
          * @private
          */
-        static async doInitAppState({res, appJSX: _appJSX}) {
+        static async doInitAppState({res, appJSX}) {
             const queryClient = (res.locals.__queryClient =
                 res.locals.__queryClient || new QueryClient())
 
-            const appJSX = React.createElement(IsPrePassContext.Provider, {value: true}, _appJSX)
+            res.locals.isServerContextReady = false
             await ssrPrepass(appJSX)
+            delete res.locals.isServerContextReady
 
             const queryCache = queryClient.getQueryCache()
             const queries = queryCache.getAll().filter((q) => q.options.enabled !== false)
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
index 3859008380..4968fd8f3c 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
@@ -8,4 +8,3 @@ import React from 'react'
 
 export const DeviceContext = React.createContext()
 export const ServerContext = React.createContext()
-export const IsPrePassContext = React.createContext(false)
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
index 2506084538..45f3779f70 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
@@ -6,7 +6,7 @@
  */
 
 import {useContext} from 'react'
-import {IsPrePassContext, ServerContext} from './contexts'
+import {ServerContext} from './contexts'
 
 /**
  * Server context
@@ -26,15 +26,10 @@ import {IsPrePassContext, ServerContext} from './contexts'
  */
 export const useServerContext = () => {
     const serverContext = useContext(ServerContext)
-    const isPrePass = useIsPrePass()
+    const res = serverContext.res
 
     return {
         ...serverContext,
-        isServerSide: Boolean(serverContext.req) && !isPrePass
+        isServerSide: Boolean(res) && res.locals.isServerContextReady !== false
     }
 }
-
-/**
- * @private
- */
-export const useIsPrePass = () => useContext(IsPrePassContext)

From 4c33040b6eb408b036bb3eab4587f71ffef02184 Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Thu, 29 Sep 2022 12:42:50 -0700
Subject: [PATCH 06/10] Not touching `res.locals`

---
 .../ssr/universal/components/with-react-query/index.js    | 8 +++++---
 packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js     | 3 +--
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
index fd731c96fb..c174107da0 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/components/with-react-query/index.js
@@ -39,9 +39,11 @@ export const withReactQuery = (Wrapped) => {
             const queryClient = (res.locals.__queryClient =
                 res.locals.__queryClient || new QueryClient())
 
-            res.locals.isServerContextReady = false
-            await ssrPrepass(appJSX)
-            delete res.locals.isServerContextReady
+            // Without the request object, our useServerContext hook would be able tell whether on prepass
+            const withoutReq = React.cloneElement(appJSX, {
+                req: undefined
+            })
+            await ssrPrepass(withoutReq)
 
             const queryCache = queryClient.getQueryCache()
             const queries = queryCache.getAll().filter((q) => q.options.enabled !== false)
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
index 45f3779f70..4bdfd86236 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
@@ -26,10 +26,9 @@ import {ServerContext} from './contexts'
  */
 export const useServerContext = () => {
     const serverContext = useContext(ServerContext)
-    const res = serverContext.res
 
     return {
         ...serverContext,
-        isServerSide: Boolean(res) && res.locals.isServerContextReady !== false
+        isServerSide: Boolean(serverContext.req)
     }
 }

From 4ccae9cc26d05135400cc57ebe845c982d16e641 Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Thu, 29 Sep 2022 13:06:35 -0700
Subject: [PATCH 07/10] Merge some files together

For example: merging contexts.js and contexts/index.js
---
 .../src/ssr/universal/contexts.js             | 10 ------
 .../src/ssr/universal/contexts/index.js       |  5 ++-
 .../src/ssr/universal/hooks.js                | 34 -------------------
 .../src/ssr/universal/hooks/index.js          | 27 ++++++++++++++-
 4 files changed, 30 insertions(+), 46 deletions(-)
 delete mode 100644 packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
 delete mode 100644 packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js

diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
deleted file mode 100644
index 4968fd8f3c..0000000000
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/contexts.js
+++ /dev/null
@@ -1,10 +0,0 @@
-/*
- * Copyright (c) 2022, Salesforce, Inc.
- * All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
- * For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
- */
-import React from 'react'
-
-export const DeviceContext = React.createContext()
-export const ServerContext = React.createContext()
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/contexts/index.js b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts/index.js
index 2dad91b0c0..6a0d31addb 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/contexts/index.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts/index.js
@@ -9,6 +9,9 @@ import React, {useEffect, useRef} from 'react'
 import PropTypes from 'prop-types'
 import {useLocation} from 'react-router-dom'
 
+const DeviceContext = React.createContext()
+const ServerContext = React.createContext()
+
 const CorrelationIdContext = React.createContext()
 
 /**
@@ -60,4 +63,4 @@ CorrelationIdProvider.propTypes = {
     location: PropTypes.object
 }
 
-export {CorrelationIdContext, CorrelationIdProvider}
+export {CorrelationIdContext, CorrelationIdProvider, DeviceContext, ServerContext}
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
deleted file mode 100644
index 4bdfd86236..0000000000
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks.js
+++ /dev/null
@@ -1,34 +0,0 @@
-/*
- * Copyright (c) 2022, Salesforce, Inc.
- * All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
- * For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
- */
-
-import {useContext} from 'react'
-import {ServerContext} from './contexts'
-
-/**
- * Server context
- * @typedef {Object} ServerContext
- * @property {Object} req - Request object
- * @property {Object} res - Response object
- * @property {boolean} isServerSide
- */
-
-/**
- * Get the server context
- * @returns {ServerContext} ServerContext object
- *
- * @example
- * const {res, isServerSide} = useServerContext()
- * if (isServerSide && query.error) { res.status(404) }
- */
-export const useServerContext = () => {
-    const serverContext = useContext(ServerContext)
-
-    return {
-        ...serverContext,
-        isServerSide: Boolean(serverContext.req)
-    }
-}
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks/index.js b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks/index.js
index 46d4e5bc9e..897b958911 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/hooks/index.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/hooks/index.js
@@ -7,7 +7,7 @@
 /* istanbul ignore file */
 
 import React, {useContext} from 'react'
-import {CorrelationIdContext} from '../contexts'
+import {CorrelationIdContext, ServerContext} from '../contexts'
 
 /**
  * Use this hook to get the correlation id value of the closest CorrelationIdProvider component.
@@ -21,3 +21,28 @@ export const useCorrelationId = () => {
     }
     return context
 }
+
+/**
+ * Server context
+ * @typedef {Object} ServerContext
+ * @property {Object} req - Request object
+ * @property {Object} res - Response object
+ * @property {boolean} isServerSide
+ */
+
+/**
+ * Get the server context
+ * @returns {ServerContext} ServerContext object
+ *
+ * @example
+ * const {res, isServerSide} = useServerContext()
+ * if (isServerSide && query.error) { res.status(404) }
+ */
+export const useServerContext = () => {
+    const serverContext = useContext(ServerContext)
+
+    return {
+        ...serverContext,
+        isServerSide: Boolean(serverContext.req)
+    }
+}

From 3944f41cfa23bdd5cdb04954940b28e969ff5cd0 Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Thu, 29 Sep 2022 16:06:20 -0700
Subject: [PATCH 08/10] Revert changes to device-context.js

---
 packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx      | 3 ++-
 .../pwa-kit-react-sdk/src/ssr/server/react-rendering.js  | 3 ++-
 .../src/ssr/universal/contexts/index.js                  | 6 ++----
 .../src/ssr/universal/device-context.js                  | 9 +++++++++
 4 files changed, 15 insertions(+), 6 deletions(-)
 create mode 100644 packages/pwa-kit-react-sdk/src/ssr/universal/device-context.js

diff --git a/packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx b/packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx
index e17d0a110f..e5de27c9df 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx
+++ b/packages/pwa-kit-react-sdk/src/ssr/browser/main.jsx
@@ -8,7 +8,8 @@
 import React from 'react'
 import ReactDOM from 'react-dom'
 import {BrowserRouter as Router} from 'react-router-dom'
-import {DeviceContext, ServerContext, CorrelationIdProvider} from '../universal/contexts'
+import DeviceContext from '../universal/device-context'
+import {ServerContext, CorrelationIdProvider} from '../universal/contexts'
 import App from '../universal/components/_app'
 import {getAppConfig} from '../universal/compatibility'
 import Switch from '../universal/components/switch'
diff --git a/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.js b/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.js
index 5d5b84203d..2a3defc7fb 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.js
@@ -18,7 +18,8 @@ import {StaticRouter as Router, matchPath} from 'react-router-dom'
 import serialize from 'serialize-javascript'
 
 import {getAssetUrl} from '../universal/utils'
-import {DeviceContext, ServerContext, CorrelationIdProvider} from '../universal/contexts'
+import DeviceContext from '../universal/device-context'
+import {ServerContext, CorrelationIdProvider} from '../universal/contexts'
 
 import Document from '../universal/components/_document'
 import App from '../universal/components/_app'
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/contexts/index.js b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts/index.js
index 6a0d31addb..23240730e9 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/universal/contexts/index.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/contexts/index.js
@@ -9,10 +9,8 @@ import React, {useEffect, useRef} from 'react'
 import PropTypes from 'prop-types'
 import {useLocation} from 'react-router-dom'
 
-const DeviceContext = React.createContext()
-const ServerContext = React.createContext()
-
 const CorrelationIdContext = React.createContext()
+const ServerContext = React.createContext()
 
 /**
  * This provider initializes the correlation id,
@@ -63,4 +61,4 @@ CorrelationIdProvider.propTypes = {
     location: PropTypes.object
 }
 
-export {CorrelationIdContext, CorrelationIdProvider, DeviceContext, ServerContext}
+export {CorrelationIdContext, CorrelationIdProvider, ServerContext}
diff --git a/packages/pwa-kit-react-sdk/src/ssr/universal/device-context.js b/packages/pwa-kit-react-sdk/src/ssr/universal/device-context.js
new file mode 100644
index 0000000000..dcece9703c
--- /dev/null
+++ b/packages/pwa-kit-react-sdk/src/ssr/universal/device-context.js
@@ -0,0 +1,9 @@
+/*
+ * Copyright (c) 2022, Salesforce, Inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ * For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+ */
+
+import React from 'react'
+export default React.createContext()

From f258e9f941dd31afcf7c918f8118e9784299391d Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Thu, 29 Sep 2022 19:50:29 -0700
Subject: [PATCH 09/10] Add test

---
 .../src/ssr/server/react-rendering.test.js    | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.test.js b/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.test.js
index 0060661550..c88ecb5639 100644
--- a/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.test.js
+++ b/packages/pwa-kit-react-sdk/src/ssr/server/react-rendering.test.js
@@ -59,6 +59,7 @@ jest.mock('../universal/routes', () => {
     const {Redirect} = require('react-router-dom')
     const {Helmet} = require('react-helmet')
     const {useQuery} = require('@tanstack/react-query')
+    const {useServerContext} = require('../universal/hooks')
 
     // Test utility to exercise paths that work with @loadable/component.
     const fakeLoadable = (Wrapped) => {
@@ -262,6 +263,15 @@ jest.mock('../universal/routes', () => {
         return <div>{isLoading ? 'loading' : data.prop}</div>
     }
 
+    const GetServerContext = () => {
+        const {res, isServerSide} = useServerContext()
+        if (isServerSide) {
+            console.log('--- isServerSide')
+            res.status(404)
+        }
+        return <div />
+    }
+
     GetPropsReturnsObject.propTypes = {
         prop: PropTypes.node
     }
@@ -324,6 +334,10 @@ jest.mock('../universal/routes', () => {
             {
                 path: '/disabled-use-query-isnt-resolved/',
                 component: DisabledUseQueryIsntResolved
+            },
+            {
+                path: '/server-context',
+                component: GetServerContext
             }
         ]
     }
@@ -692,6 +706,20 @@ describe('The Node SSR Environment', () => {
                 const html = res.text
                 expect(html).toEqual(expect.stringContaining('<div>loading</div>'))
             }
+        },
+        {
+            description: 'Get the server context and set the response status to 404',
+            req: {url: '/server-context'},
+            mocks: () => {
+                jest.spyOn(console, 'log')
+            },
+            assertions: (res) => {
+                expect(res.statusCode).toBe(404)
+
+                // Expect the console.log to be called only once, even though the component
+                // is going to be rendered twice on the server (on prepass and then 2nd pass)
+                expect(console.log).toHaveBeenCalledTimes(1)
+            }
         }
     ]
 

From 6e2e7bb19b83e15b329dc5a63625d569fe84535e Mon Sep 17 00:00:00 2001
From: Vincent Marta <vmarta@salesforce.com>
Date: Thu, 29 Sep 2022 19:53:14 -0700
Subject: [PATCH 10/10] Remove the example code

---
 packages/template-typescript-minimal/app/pages/home.tsx | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/packages/template-typescript-minimal/app/pages/home.tsx b/packages/template-typescript-minimal/app/pages/home.tsx
index ab8502829a..cbdeffdba7 100644
--- a/packages/template-typescript-minimal/app/pages/home.tsx
+++ b/packages/template-typescript-minimal/app/pages/home.tsx
@@ -10,8 +10,6 @@ import {useQuery} from '@tanstack/react-query'
 import HelloTS from '../components/hello-typescript'
 import HelloJS from '../components/hello-javascript'
 
-import {useServerContext} from 'pwa-kit-react-sdk/ssr/universal/hooks'
-
 interface Props {
     value: number
 }
@@ -103,12 +101,6 @@ const Home = ({value}: Props) => {
             })
     )
 
-    const {res, isServerSide} = useServerContext()
-    if (isServerSide && query.data) {
-        console.log('--- useServerContext')
-        res.status(404)
-    }
-
     return (
         <div>
             <style dangerouslySetInnerHTML={{__html: style}} />