From 537d40e27fa564444fdc295b7de10858ac444a26 Mon Sep 17 00:00:00 2001
From: Josh Larson <josh.larson@shopify.com>
Date: Fri, 18 Nov 2022 13:57:50 -0600
Subject: [PATCH 1/5] Introduce injected `session` and `sessionStorage` utils
 in user-code

---
 package-lock.json                             |  4 +--
 packages/hydrogen-remix/src/server.ts         |  2 ++
 templates/demo-store/app/data/index.ts        |  7 ++--
 .../demo-store/app/lib/session.server.ts      | 33 ++----------------
 templates/demo-store/app/root.tsx             |  8 ++---
 .../app/routes/__resources/cart/LinesAdd.tsx  |  9 ++---
 templates/demo-store/app/routes/account.tsx   |  4 +--
 .../routes/account/__private/address/$id.tsx  |  8 ++---
 .../app/routes/account/__private/edit.tsx     |  8 ++---
 .../app/routes/account/__private/logout.ts    | 10 +++---
 .../routes/account/__private/orders.$id.tsx   |  4 +--
 .../activate.$id.$activationToken.tsx         | 10 +++---
 .../app/routes/account/__public/login.tsx     | 14 +++-----
 .../app/routes/account/__public/recover.tsx   |  6 ++--
 .../app/routes/account/__public/register.tsx  | 14 +++-----
 .../__public/reset.$id.$resetToken.tsx        | 10 +++---
 templates/demo-store/app/routes/cart.tsx      |  8 ++---
 templates/demo-store/oxygen.ts                | 34 ++++++++++++++-----
 templates/demo-store/package.json             |  4 +--
 templates/demo-store/types/remix.env.d.ts     | 14 ++++++++
 templates/demo-store/types/worker.env.d.ts    |  4 ++-
 21 files changed, 93 insertions(+), 122 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index faa7aff1ad..207b74c5ed 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,5 +1,5 @@
 {
-  "name": "h2-demo-store",
+  "name": "h2",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
@@ -19057,7 +19057,7 @@
         "@remix-run/react": "0.0.0-experimental-e18af792a",
         "@shopify/eslint-plugin": "^42.0.1",
         "@shopify/hydrogen": "^2.0.0-alpha.1",
-        "@shopify/hydrogen-react": "^2022.10.1",
+        "@shopify/hydrogen-react": "^2022.10.3",
         "@shopify/oxygen-workers-types": "^3.17.2",
         "@shopify/prettier-config": "^1.1.2",
         "@tailwindcss/forms": "^0.5.3",
diff --git a/packages/hydrogen-remix/src/server.ts b/packages/hydrogen-remix/src/server.ts
index c7f28fb895..8a258a6695 100644
--- a/packages/hydrogen-remix/src/server.ts
+++ b/packages/hydrogen-remix/src/server.ts
@@ -26,6 +26,7 @@ export function createRequestHandler(
       ...options
     }: Omit<Parameters<typeof handleRequest>[1], 'loadContext'> &
       HydrogenHandlerParams,
+    customContext?: Record<string, any>,
   ) => {
     try {
       if (!cache && !!globalThis.caches) {
@@ -42,6 +43,7 @@ export function createRequestHandler(
           }),
           ...context,
           cache,
+          ...customContext,
         },
       });
 
diff --git a/templates/demo-store/app/data/index.ts b/templates/demo-store/app/data/index.ts
index f5018dbb52..0f047a1321 100644
--- a/templates/demo-store/app/data/index.ts
+++ b/templates/demo-store/app/data/index.ts
@@ -27,7 +27,7 @@ import {
 } from '~/lib/utils';
 import invariant from 'tiny-invariant';
 import {logout} from '~/routes/account/__private/logout';
-import type {HydrogenContext} from '@shopify/hydrogen-remix';
+import type {AppLoadContext, HydrogenContext} from '@shopify/hydrogen-remix';
 import {type Params} from '@remix-run/react';
 
 export interface CountriesData {
@@ -932,9 +932,8 @@ export async function getCustomerOrder(
 }
 
 export async function getCustomer(
-  context: HydrogenContext,
+  context: AppLoadContext & HydrogenContext,
   {
-    request,
     customerAccessToken,
     params,
   }: {
@@ -960,7 +959,7 @@ export async function getCustomer(
    * If the customer failed to load, we assume their access token is invalid.
    */
   if (!data || !data.customer) {
-    throw await logout(request, context, params);
+    throw logout(context, params);
   }
 
   return data.customer;
diff --git a/templates/demo-store/app/lib/session.server.ts b/templates/demo-store/app/lib/session.server.ts
index e87e904353..076f18f42a 100644
--- a/templates/demo-store/app/lib/session.server.ts
+++ b/templates/demo-store/app/lib/session.server.ts
@@ -6,21 +6,14 @@ import {
 
 let sessionStorage: SessionStorage;
 
-export async function getSession(
-  request: Request,
-  context: AppLoadContext & {env?: Record<string, string>},
-) {
-  if (!context.env?.ENCRYPTION_KEY) {
-    throw new Error('ENCRYPTION_KEY environment variable is not set');
-  }
-
+export async function getSession(request: Request, secrets: string[]) {
   sessionStorage ??= createCookieSessionStorage({
     cookie: {
       name: 'session',
       httpOnly: true,
       path: '/',
       sameSite: 'lax',
-      secrets: [context.env.ENCRYPTION_KEY],
+      secrets,
     },
   });
 
@@ -28,25 +21,5 @@ export async function getSession(
     request.headers.get('Cookie'),
   );
 
-  return {
-    async get(key: string): Promise<any> {
-      return await session.get(key);
-    },
-
-    set(key: string, value: any): void {
-      session.set(key, value);
-    },
-
-    flash(key: string, value: any): void {
-      session.flash(key, value);
-    },
-
-    unset(key: string): void {
-      session.unset(key);
-    },
-
-    async commit(): Promise<string> {
-      return await sessionStorage.commitSession(session);
-    },
-  };
+  return {session, sessionStorage};
 }
diff --git a/templates/demo-store/app/root.tsx b/templates/demo-store/app/root.tsx
index 879d984623..fc92aa4cd5 100644
--- a/templates/demo-store/app/root.tsx
+++ b/templates/demo-store/app/root.tsx
@@ -20,7 +20,6 @@ import {Layout} from '~/components';
 import {getLayoutData, getCountries} from '~/data';
 import {GenericError} from './components/GenericError';
 import {NotFound} from './components/NotFound';
-import {getSession} from './lib/session.server';
 import {Seo, Debugger} from './lib/seo';
 
 import styles from './styles/app.css';
@@ -59,9 +58,8 @@ export const meta: MetaFunction = () => ({
   viewport: 'width=device-width,initial-scale=1',
 });
 
-export async function loader({request, context, params}: LoaderArgs) {
-  const session = await getSession(request, context);
-  const cartId = await session.get('cartId');
+export async function loader({context, params}: LoaderArgs) {
+  const cartId = await context.session.get('cartId');
 
   return defer({
     layout: await getLayoutData(context, params),
@@ -105,7 +103,7 @@ export function CatchBoundary() {
         <Links />
       </head>
       <body>
-        <Layout layout={root.data.layout}>
+        <Layout layout={root.data?.layout}>
           {isNotFound ? (
             <NotFound type={caught.data?.pageType} />
           ) : (
diff --git a/templates/demo-store/app/routes/__resources/cart/LinesAdd.tsx b/templates/demo-store/app/routes/__resources/cart/LinesAdd.tsx
index 74bc021a92..85512180af 100644
--- a/templates/demo-store/app/routes/__resources/cart/LinesAdd.tsx
+++ b/templates/demo-store/app/routes/__resources/cart/LinesAdd.tsx
@@ -9,7 +9,6 @@ import {
 } from '@remix-run/react';
 import {useIsHydrated} from '~/hooks/useIsHydrated';
 import invariant from 'tiny-invariant';
-import {getSession} from '~/lib/session.server';
 import type {SerializeFrom} from '@remix-run/server-runtime';
 import {
   type ActionArgs,
@@ -97,12 +96,10 @@ const ACTION_PATH = '/cart/LinesAdd';
  * action that handles cart create (with lines) and lines add
  */
 async function action({request, context, params}: ActionArgs) {
+  const {session, sessionStorage} = context;
   const headers = new Headers();
 
-  const [session, formData] = await Promise.all([
-    getSession(request, context),
-    request.formData(),
-  ]);
+  const formData = await request.formData();
 
   const rawLines = formData.get('lines')
     ? (JSON.parse(String(formData.get('lines'))) as LinesAddLine[])
@@ -133,7 +130,7 @@ async function action({request, context, params}: ActionArgs) {
 
     // cart created - we only need a Set-Cookie header if we're creating
     session.set('cartId', cart.id);
-    headers.set('Set-Cookie', await session.commit());
+    headers.set('Set-Cookie', await sessionStorage.commitSession(session));
 
     return linesAddResponse({
       prevCart: null,
diff --git a/templates/demo-store/app/routes/account.tsx b/templates/demo-store/app/routes/account.tsx
index 6ef09c06f9..d727fae8ea 100644
--- a/templates/demo-store/app/routes/account.tsx
+++ b/templates/demo-store/app/routes/account.tsx
@@ -27,14 +27,12 @@ import {FeaturedCollections} from '~/components/FeaturedCollections';
 import {type LoaderArgs, redirect, json, defer} from '@shopify/hydrogen-remix';
 import {flattenConnection} from '@shopify/hydrogen-react';
 import {getCustomer} from '~/data';
-import {getSession} from '~/lib/session.server';
 import {getFeaturedData} from './featured-products';
 
 export async function loader({request, context, params}: LoaderArgs) {
   const {pathname} = new URL(request.url);
-  const session = await getSession(request, context);
   const lang = params.lang;
-  const customerAccessToken = await session.get('customerAccessToken');
+  const customerAccessToken = await context.session.get('customerAccessToken');
   const isAuthenticated = Boolean(customerAccessToken);
   const loginPath = lang ? `${lang}/account/login` : '/account/login';
 
diff --git a/templates/demo-store/app/routes/account/__private/address/$id.tsx b/templates/demo-store/app/routes/account/__private/address/$id.tsx
index 295a67eb64..9140b36470 100644
--- a/templates/demo-store/app/routes/account/__private/address/$id.tsx
+++ b/templates/demo-store/app/routes/account/__private/address/$id.tsx
@@ -16,7 +16,6 @@ import {
   updateCustomerAddress,
   updateCustomerDefaultAddress,
 } from '~/data';
-import {getSession} from '~/lib/session.server';
 import {getInputStyleClasses} from '~/lib/utils';
 import type {AccountOutletContext} from '../edit';
 
@@ -31,12 +30,9 @@ export const handle = {
 };
 
 export const action: ActionFunction = async ({request, context, params}) => {
-  const [formData, session] = await Promise.all([
-    request.formData(),
-    getSession(request, context),
-  ]);
+  const formData = await request.formData();
 
-  const customerAccessToken = await session.get('customerAccessToken');
+  const customerAccessToken = await context.session.get('customerAccessToken');
   invariant(customerAccessToken, 'You must be logged in to edit your account.');
 
   const addressId = formData.get('addressId');
diff --git a/templates/demo-store/app/routes/account/__private/edit.tsx b/templates/demo-store/app/routes/account/__private/edit.tsx
index bc70bc8c5e..b99d7e6858 100644
--- a/templates/demo-store/app/routes/account/__private/edit.tsx
+++ b/templates/demo-store/app/routes/account/__private/edit.tsx
@@ -13,7 +13,6 @@ import clsx from 'clsx';
 import invariant from 'tiny-invariant';
 import {Button, Text} from '~/components';
 import {getCustomer, updateCustomer} from '~/data';
-import {getSession} from '~/lib/session.server';
 import {getInputStyleClasses} from '~/lib/utils';
 
 export interface AccountOutletContext {
@@ -48,12 +47,9 @@ export const handle = {
 };
 
 export const action: ActionFunction = async ({request, context, params}) => {
-  const [formData, session] = await Promise.all([
-    request.formData(),
-    getSession(request, context),
-  ]);
+  const formData = await request.formData();
 
-  const customerAccessToken = await session.get('customerAccessToken');
+  const customerAccessToken = await context.session.get('customerAccessToken');
 
   invariant(
     customerAccessToken,
diff --git a/templates/demo-store/app/routes/account/__private/logout.ts b/templates/demo-store/app/routes/account/__private/logout.ts
index 6ef716ef1d..2d424e6fb5 100644
--- a/templates/demo-store/app/routes/account/__private/logout.ts
+++ b/templates/demo-store/app/routes/account/__private/logout.ts
@@ -4,21 +4,19 @@ import {
   redirect,
 } from '@shopify/hydrogen-remix';
 import {LoaderArgs} from '@remix-run/server-runtime';
-import {getSession} from '~/lib/session.server';
 
 export async function logout(
-  request: Request,
   context: AppLoadContext,
   params: LoaderArgs['params'],
 ) {
-  const session = await getSession(request, context);
+  const {session, sessionStorage} = context;
   session.unset('customerAccessToken');
 
   return redirect(
     params.lang ? `${params.lang}/account/login` : '/account/login',
     {
       headers: {
-        'Set-Cookie': await session.commit(),
+        'Set-Cookie': await sessionStorage.commitSession(session),
       },
     },
   );
@@ -28,6 +26,6 @@ export async function loader({params}: LoaderArgs) {
   return redirect(params.lang ? `${params.lang}/` : '/');
 }
 
-export const action: ActionFunction = async ({request, context, params}) => {
-  return logout(request, context, params);
+export const action: ActionFunction = async ({context, params}) => {
+  return logout(context, params);
 };
diff --git a/templates/demo-store/app/routes/account/__private/orders.$id.tsx b/templates/demo-store/app/routes/account/__private/orders.$id.tsx
index cc179eac30..8189e9c12f 100644
--- a/templates/demo-store/app/routes/account/__private/orders.$id.tsx
+++ b/templates/demo-store/app/routes/account/__private/orders.$id.tsx
@@ -16,7 +16,6 @@ import type {
 } from '@shopify/hydrogen-react/storefront-api-types';
 import {Link, Heading, PageHeader, Text} from '~/components';
 import {getCustomerOrder} from '~/data';
-import {getSession} from '~/lib/session.server';
 
 export const meta: MetaFunction = ({data}) => ({
   title: `Order ${data?.order?.name}`,
@@ -32,8 +31,7 @@ export async function loader({request, context, params}: LoaderArgs) {
 
   invariant(orderToken, 'Order token is required');
 
-  const session = await getSession(request, context);
-  const customerAccessToken = await session.get('customerAccessToken');
+  const customerAccessToken = await context.session.get('customerAccessToken');
 
   if (!customerAccessToken) {
     return redirect(
diff --git a/templates/demo-store/app/routes/account/__public/activate.$id.$activationToken.tsx b/templates/demo-store/app/routes/account/__public/activate.$id.$activationToken.tsx
index d9fc1c5120..c4c0846504 100644
--- a/templates/demo-store/app/routes/account/__public/activate.$id.$activationToken.tsx
+++ b/templates/demo-store/app/routes/account/__public/activate.$id.$activationToken.tsx
@@ -8,7 +8,6 @@ import {
 import {Form, useActionData} from '@remix-run/react';
 import {useRef, useState} from 'react';
 import {activateAccount} from '~/data';
-import {getSession} from '~/lib/session.server';
 import {getInputStyleClasses} from '~/lib/utils';
 
 type ActionData = {
@@ -37,10 +36,7 @@ export const action: ActionFunction = async ({
     });
   }
 
-  const [formData, session] = await Promise.all([
-    request.formData(),
-    getSession(request, context),
-  ]);
+  const formData = await request.formData();
 
   const password = formData.get('password');
   const passwordConfirm = formData.get('passwordConfirm');
@@ -64,11 +60,13 @@ export const action: ActionFunction = async ({
       password,
     });
 
+    const {session, sessionStorage} = context;
+
     session.set('customerAccessToken', accessToken);
 
     return redirect(lang ? `${lang}/account` : '/account', {
       headers: {
-        'Set-Cookie': await session.commit(),
+        'Set-Cookie': await sessionStorage.commitSession(session),
       },
     });
   } catch (error: any) {
diff --git a/templates/demo-store/app/routes/account/__public/login.tsx b/templates/demo-store/app/routes/account/__public/login.tsx
index 9ed68c7bcb..49820f5c67 100644
--- a/templates/demo-store/app/routes/account/__public/login.tsx
+++ b/templates/demo-store/app/routes/account/__public/login.tsx
@@ -9,7 +9,6 @@ import {
 import {Form, useActionData, useLoaderData} from '@remix-run/react';
 import {useState} from 'react';
 import {login} from '~/data';
-import {getSession} from '~/lib/session.server';
 import {getInputStyleClasses} from '~/lib/utils';
 import {Link} from '~/components';
 
@@ -17,9 +16,8 @@ export const handle = {
   isPublic: true,
 };
 
-export async function loader({request, context, params}: LoaderArgs) {
-  const session = await getSession(request, context);
-  const customerAccessToken = await session.get('customerAccessToken');
+export async function loader({context, params}: LoaderArgs) {
+  const customerAccessToken = await context.session.get('customerAccessToken');
 
   if (customerAccessToken) {
     return redirect(params.lang ? `${params.lang}/account` : '/account');
@@ -36,10 +34,7 @@ type ActionData = {
 const badRequest = (data: ActionData) => json(data, {status: 400});
 
 export const action: ActionFunction = async ({request, context, params}) => {
-  const [formData, session] = await Promise.all([
-    request.formData(),
-    getSession(request, context),
-  ]);
+  const formData = await request.formData();
 
   const email = formData.get('email');
   const password = formData.get('password');
@@ -57,11 +52,12 @@ export const action: ActionFunction = async ({request, context, params}) => {
 
   try {
     const customerAccessToken = await login(context, {email, password});
+    const {session, sessionStorage} = context;
     session.set('customerAccessToken', customerAccessToken);
 
     return redirect(params.lang ? `${params.lang}/account` : '/account', {
       headers: {
-        'Set-Cookie': await session.commit(),
+        'Set-Cookie': await sessionStorage.commitSession(session),
       },
     });
   } catch (error: any) {
diff --git a/templates/demo-store/app/routes/account/__public/recover.tsx b/templates/demo-store/app/routes/account/__public/recover.tsx
index af76855f22..2f1ffd4a83 100644
--- a/templates/demo-store/app/routes/account/__public/recover.tsx
+++ b/templates/demo-store/app/routes/account/__public/recover.tsx
@@ -8,13 +8,11 @@ import {
 import {Form, useActionData} from '@remix-run/react';
 import {useState} from 'react';
 import {sendPasswordResetEmail} from '~/data';
-import {getSession} from '~/lib/session.server';
 import {Link} from '~/components';
 import {getInputStyleClasses} from '~/lib/utils';
 
-export async function loader({request, context, params}: LoaderArgs) {
-  const session = await getSession(request, context);
-  const customerAccessToken = await session.get('customerAccessToken');
+export async function loader({context, params}: LoaderArgs) {
+  const customerAccessToken = await context.session.get('customerAccessToken');
 
   if (customerAccessToken) {
     return redirect(params.lang ? `${params.lang}/account` : '/account');
diff --git a/templates/demo-store/app/routes/account/__public/register.tsx b/templates/demo-store/app/routes/account/__public/register.tsx
index db9b6f3dfb..33876f30e6 100644
--- a/templates/demo-store/app/routes/account/__public/register.tsx
+++ b/templates/demo-store/app/routes/account/__public/register.tsx
@@ -9,13 +9,11 @@ import {
 import {Form, useActionData} from '@remix-run/react';
 import {useState} from 'react';
 import {login, registerCustomer} from '~/data';
-import {getSession} from '~/lib/session.server';
 import {getInputStyleClasses} from '~/lib/utils';
 import {Link} from '~/components';
 
-export async function loader({request, context, params}: LoaderArgs) {
-  const session = await getSession(request, context);
-  const customerAccessToken = await session.get('customerAccessToken');
+export async function loader({context, params}: LoaderArgs) {
+  const customerAccessToken = await context.session.get('customerAccessToken');
 
   if (customerAccessToken) {
     return redirect(params.lang ? `${params.lang}/account` : '/account');
@@ -31,10 +29,8 @@ type ActionData = {
 const badRequest = (data: ActionData) => json(data, {status: 400});
 
 export const action: ActionFunction = async ({request, context, params}) => {
-  const [formData, session] = await Promise.all([
-    request.formData(),
-    getSession(request, context),
-  ]);
+  const {session, sessionStorage} = context;
+  const formData = await request.formData();
 
   const email = formData.get('email');
   const password = formData.get('password');
@@ -57,7 +53,7 @@ export const action: ActionFunction = async ({request, context, params}) => {
 
     return redirect(params.lang ? `${params.lang}/account` : '/account', {
       headers: {
-        'Set-Cookie': await session.commit(),
+        'Set-Cookie': await sessionStorage.commitSession(session),
       },
     });
   } catch (error: any) {
diff --git a/templates/demo-store/app/routes/account/__public/reset.$id.$resetToken.tsx b/templates/demo-store/app/routes/account/__public/reset.$id.$resetToken.tsx
index ebb1bcc435..9238127ce9 100644
--- a/templates/demo-store/app/routes/account/__public/reset.$id.$resetToken.tsx
+++ b/templates/demo-store/app/routes/account/__public/reset.$id.$resetToken.tsx
@@ -8,7 +8,6 @@ import {
 import {Form, useActionData} from '@remix-run/react';
 import {useRef, useState} from 'react';
 import {resetPassword} from '~/data';
-import {getSession} from '~/lib/session.server';
 import {getInputStyleClasses} from '~/lib/utils';
 
 type ActionData = {
@@ -33,10 +32,7 @@ export const action: ActionFunction = async ({
     });
   }
 
-  const [formData, session] = await Promise.all([
-    request.formData(),
-    getSession(request, context),
-  ]);
+  const formData = await request.formData();
 
   const password = formData.get('password');
   const passwordConfirm = formData.get('passwordConfirm');
@@ -60,11 +56,13 @@ export const action: ActionFunction = async ({
       password,
     });
 
+    const {session, sessionStorage} = context;
+
     session.set('customerAccessToken', accessToken);
 
     return redirect(lang ? `${lang}/account` : '/account', {
       headers: {
-        'Set-Cookie': await session.commit(),
+        'Set-Cookie': await sessionStorage.commitSession(session),
       },
     });
   } catch (error: any) {
diff --git a/templates/demo-store/app/routes/cart.tsx b/templates/demo-store/app/routes/cart.tsx
index 4761f9c156..8b1807efef 100644
--- a/templates/demo-store/app/routes/cart.tsx
+++ b/templates/demo-store/app/routes/cart.tsx
@@ -6,7 +6,6 @@ import {
 } from '@shopify/hydrogen-remix';
 import invariant from 'tiny-invariant';
 import {getTopProducts, updateLineItem} from '~/data';
-import {getSession} from '~/lib/session.server';
 
 export async function loader({params, context}: LoaderArgs) {
   return defer(
@@ -24,17 +23,14 @@ export async function loader({params, context}: LoaderArgs) {
 export const action: ActionFunction = async ({request, context, params}) => {
   let cart;
 
-  const [session, formData] = await Promise.all([
-    getSession(request, context),
-    new URLSearchParams(await request.text()),
-  ]);
+  const formData = new URLSearchParams(await request.text());
 
   const redirectTo = formData.get('redirectTo');
   const intent = formData.get('intent');
   invariant(intent, 'Missing cart intent');
 
   // 1. Grab the cart ID from the session
-  const cartId = await session.get('cartId');
+  const cartId = await context.session.get('cartId');
 
   switch (intent) {
     case 'set-quantity': {
diff --git a/templates/demo-store/oxygen.ts b/templates/demo-store/oxygen.ts
index 476fed1957..04763dcfb9 100644
--- a/templates/demo-store/oxygen.ts
+++ b/templates/demo-store/oxygen.ts
@@ -1,6 +1,7 @@
 import {createRequestHandler} from '@shopify/hydrogen-remix';
 // The build remix app provided by remix build
 import * as remixBuild from 'remix-build';
+import {getSession} from '~/lib/session.server';
 
 declare const process: {env: {NODE_ENV: string}};
 
@@ -16,16 +17,33 @@ export default {
     env: Env,
     context: ExecutionContext,
   ): Promise<Response> {
+    if (!env?.SESSION_SECRET) {
+      // eslint-disable-next-line no-console
+      console.error('SESSION_SECRET environment variable is not set');
+      return new Response('Internal Server Error', {status: 500});
+    }
+
+    const {session, sessionStorage} = await getSession(request, [
+      env.SESSION_SECRET,
+    ]);
+
     try {
-      return await requestHandler(request, {
-        env,
-        context,
-        storefront: {
-          publicStorefrontToken: '3b580e70970c4528da70c98e097c2fa0',
-          storeDomain: 'hydrogen-preview',
-          storefrontApiVersion: '2022-10',
+      return await requestHandler(
+        request,
+        {
+          env,
+          context,
+          storefront: {
+            publicStorefrontToken: '3b580e70970c4528da70c98e097c2fa0',
+            storeDomain: 'hydrogen-preview',
+            storefrontApiVersion: '2022-10',
+          },
+        },
+        {
+          session,
+          sessionStorage,
         },
-      });
+      );
     } catch (error) {
       // eslint-disable-next-line no-console
       console.error(error);
diff --git a/templates/demo-store/package.json b/templates/demo-store/package.json
index a98f74d88f..8850ead969 100644
--- a/templates/demo-store/package.json
+++ b/templates/demo-store/package.json
@@ -4,9 +4,9 @@
   "scripts": {
     "build": "npm run build:css && npm run h2 build -- --entry oxygen.ts",
     "build:css": "postcss styles --base styles --dir app/styles --env production",
-    "dev": "ENCRYPTION_KEY=foobar run-p dev:css \"h2 dev -- --entry oxygen.ts\"",
+    "dev": "SESSION_SECRET=foobar run-p dev:css \"h2 dev -- --entry oxygen.ts\"",
     "dev:css": "postcss styles --base styles --dir app/styles -w",
-    "preview": "npm run build && ENCRYPTION_KEY=foobar npm run h2 preview",
+    "preview": "npm run build && SESSION_SECRET=foobar npm run h2 preview",
     "lint": "eslint --no-error-on-unmatched-pattern --ext .js,.ts,.jsx,.tsx .",
     "format": "prettier --write --ignore-unknown .",
     "typecheck": "tsc --noEmit",
diff --git a/templates/demo-store/types/remix.env.d.ts b/templates/demo-store/types/remix.env.d.ts
index d9f78fc474..c105529011 100644
--- a/templates/demo-store/types/remix.env.d.ts
+++ b/templates/demo-store/types/remix.env.d.ts
@@ -1,2 +1,16 @@
 /// <reference types="@remix-run/dev" />
 /// <reference types="@remix-run/oxygen" />
+
+import {Session, SessionStorage} from '@shopify/hydrogen-remix';
+import '@shopify/hydrogen-remix';
+
+/**
+ * Declare local additions to `AppLoadContext` to include the session utilities we injected in `oxygen.ts`.
+ * Note that additional `HydrogenContext` properties are already included via `@shopify/hydrogen-remix`.
+ */
+declare module '@shopify/hydrogen-remix' {
+  export interface AppLoadContext {
+    session: Session;
+    sessionStorage: SessionStorage;
+  }
+}
diff --git a/templates/demo-store/types/worker.env.d.ts b/templates/demo-store/types/worker.env.d.ts
index e6e2e2d0a7..6071ac27ec 100644
--- a/templates/demo-store/types/worker.env.d.ts
+++ b/templates/demo-store/types/worker.env.d.ts
@@ -1,3 +1,5 @@
 /// <reference types="@shopify/oxygen-workers-types" />
 
-interface Env {}
+interface Env {
+  SESSION_SECRET: string;
+}

From fca9488c894bbc7d4146e9eb0fa87f4d6aa848e9 Mon Sep 17 00:00:00 2001
From: Josh Larson <josh.larson@shopify.com>
Date: Fri, 18 Nov 2022 14:04:19 -0600
Subject: [PATCH 2/5] Revert my accidental await removal

---
 templates/demo-store/app/data/index.ts                      | 2 +-
 templates/demo-store/app/routes/account/__private/logout.ts | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/templates/demo-store/app/data/index.ts b/templates/demo-store/app/data/index.ts
index 0f047a1321..f6cd785f84 100644
--- a/templates/demo-store/app/data/index.ts
+++ b/templates/demo-store/app/data/index.ts
@@ -959,7 +959,7 @@ export async function getCustomer(
    * If the customer failed to load, we assume their access token is invalid.
    */
   if (!data || !data.customer) {
-    throw logout(context, params);
+    throw await logout(context, params);
   }
 
   return data.customer;
diff --git a/templates/demo-store/app/routes/account/__private/logout.ts b/templates/demo-store/app/routes/account/__private/logout.ts
index 2d424e6fb5..a98a73e158 100644
--- a/templates/demo-store/app/routes/account/__private/logout.ts
+++ b/templates/demo-store/app/routes/account/__private/logout.ts
@@ -27,5 +27,5 @@ export async function loader({params}: LoaderArgs) {
 }
 
 export const action: ActionFunction = async ({context, params}) => {
-  return logout(context, params);
+  return await logout(context, params);
 };

From 4e553893748e4adaf8c6b4c705d2e582e2fbf854 Mon Sep 17 00:00:00 2001
From: Josh Larson <josh.larson@shopify.com>
Date: Fri, 18 Nov 2022 14:05:08 -0600
Subject: [PATCH 3/5] Revert this change too

---
 templates/demo-store/app/root.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/templates/demo-store/app/root.tsx b/templates/demo-store/app/root.tsx
index fc92aa4cd5..6243b579eb 100644
--- a/templates/demo-store/app/root.tsx
+++ b/templates/demo-store/app/root.tsx
@@ -103,7 +103,7 @@ export function CatchBoundary() {
         <Links />
       </head>
       <body>
-        <Layout layout={root.data?.layout}>
+        <Layout layout={root.data.layout}>
           {isNotFound ? (
             <NotFound type={caught.data?.pageType} />
           ) : (

From 2cec0e5796d8972c40741ea73f60196f06991f4d Mon Sep 17 00:00:00 2001
From: Josh Larson <josh.larson@shopify.com>
Date: Fri, 18 Nov 2022 14:05:43 -0600
Subject: [PATCH 4/5] Nope

---
 templates/demo-store/app/routes/account/__private/logout.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/templates/demo-store/app/routes/account/__private/logout.ts b/templates/demo-store/app/routes/account/__private/logout.ts
index a98a73e158..2d424e6fb5 100644
--- a/templates/demo-store/app/routes/account/__private/logout.ts
+++ b/templates/demo-store/app/routes/account/__private/logout.ts
@@ -27,5 +27,5 @@ export async function loader({params}: LoaderArgs) {
 }
 
 export const action: ActionFunction = async ({context, params}) => {
-  return await logout(context, params);
+  return logout(context, params);
 };

From 356129de9cb4560aedd443c41261b574e0c73878 Mon Sep 17 00:00:00 2001
From: Josh Larson <josh.larson@shopify.com>
Date: Fri, 18 Nov 2022 15:53:27 -0600
Subject: [PATCH 5/5] Refactor into a custom `HydrogenSession` class

---
 .../demo-store/app/lib/session.server.ts      | 64 ++++++++++++++-----
 .../app/routes/__resources/cart/LinesAdd.tsx  |  4 +-
 .../app/routes/account/__private/logout.ts    |  4 +-
 .../activate.$id.$activationToken.tsx         |  4 +-
 .../app/routes/account/__public/login.tsx     |  4 +-
 .../app/routes/account/__public/register.tsx  |  4 +-
 .../__public/reset.$id.$resetToken.tsx        |  4 +-
 templates/demo-store/oxygen.ts                |  7 +-
 templates/demo-store/types/remix.env.d.ts     |  4 +-
 9 files changed, 64 insertions(+), 35 deletions(-)

diff --git a/templates/demo-store/app/lib/session.server.ts b/templates/demo-store/app/lib/session.server.ts
index 076f18f42a..060ca63c70 100644
--- a/templates/demo-store/app/lib/session.server.ts
+++ b/templates/demo-store/app/lib/session.server.ts
@@ -1,25 +1,57 @@
 import {
-  type AppLoadContext,
   createCookieSessionStorage,
   type SessionStorage,
+  type Session,
 } from '@shopify/hydrogen-remix';
 
-let sessionStorage: SessionStorage;
+/**
+ * This is a custom session implementation for your Hydrogen shop.
+ * Feel free to customize it to your needs, add helper methods, or
+ * swap out the cookie-based implementation with something else!
+ */
+export class HydrogenSession {
+  constructor(
+    private sessionStorage: SessionStorage,
+    private session: Session,
+  ) {}
 
-export async function getSession(request: Request, secrets: string[]) {
-  sessionStorage ??= createCookieSessionStorage({
-    cookie: {
-      name: 'session',
-      httpOnly: true,
-      path: '/',
-      sameSite: 'lax',
-      secrets,
-    },
-  });
+  static async init(request: Request, secrets: string[]) {
+    const storage = createCookieSessionStorage({
+      cookie: {
+        name: 'session',
+        httpOnly: true,
+        path: '/',
+        sameSite: 'lax',
+        secrets,
+      },
+    });
 
-  const session = await sessionStorage.getSession(
-    request.headers.get('Cookie'),
-  );
+    const session = await storage.getSession(request.headers.get('Cookie'));
 
-  return {session, sessionStorage};
+    return new this(storage, session);
+  }
+
+  get(key: string) {
+    return this.session.get(key);
+  }
+
+  destroy() {
+    return this.sessionStorage.destroySession(this.session);
+  }
+
+  flash(key: string, value: any) {
+    this.session.flash(key, value);
+  }
+
+  unset(key: string) {
+    this.session.unset(key);
+  }
+
+  set(key: string, value: any) {
+    this.session.set(key, value);
+  }
+
+  commit() {
+    return this.sessionStorage.commitSession(this.session);
+  }
 }
diff --git a/templates/demo-store/app/routes/__resources/cart/LinesAdd.tsx b/templates/demo-store/app/routes/__resources/cart/LinesAdd.tsx
index 85512180af..d57e907445 100644
--- a/templates/demo-store/app/routes/__resources/cart/LinesAdd.tsx
+++ b/templates/demo-store/app/routes/__resources/cart/LinesAdd.tsx
@@ -96,7 +96,7 @@ const ACTION_PATH = '/cart/LinesAdd';
  * action that handles cart create (with lines) and lines add
  */
 async function action({request, context, params}: ActionArgs) {
-  const {session, sessionStorage} = context;
+  const {session} = context;
   const headers = new Headers();
 
   const formData = await request.formData();
@@ -130,7 +130,7 @@ async function action({request, context, params}: ActionArgs) {
 
     // cart created - we only need a Set-Cookie header if we're creating
     session.set('cartId', cart.id);
-    headers.set('Set-Cookie', await sessionStorage.commitSession(session));
+    headers.set('Set-Cookie', await session.commit());
 
     return linesAddResponse({
       prevCart: null,
diff --git a/templates/demo-store/app/routes/account/__private/logout.ts b/templates/demo-store/app/routes/account/__private/logout.ts
index 2d424e6fb5..85b93afc4d 100644
--- a/templates/demo-store/app/routes/account/__private/logout.ts
+++ b/templates/demo-store/app/routes/account/__private/logout.ts
@@ -9,14 +9,14 @@ export async function logout(
   context: AppLoadContext,
   params: LoaderArgs['params'],
 ) {
-  const {session, sessionStorage} = context;
+  const {session} = context;
   session.unset('customerAccessToken');
 
   return redirect(
     params.lang ? `${params.lang}/account/login` : '/account/login',
     {
       headers: {
-        'Set-Cookie': await sessionStorage.commitSession(session),
+        'Set-Cookie': await session.commit(),
       },
     },
   );
diff --git a/templates/demo-store/app/routes/account/__public/activate.$id.$activationToken.tsx b/templates/demo-store/app/routes/account/__public/activate.$id.$activationToken.tsx
index c4c0846504..a695096729 100644
--- a/templates/demo-store/app/routes/account/__public/activate.$id.$activationToken.tsx
+++ b/templates/demo-store/app/routes/account/__public/activate.$id.$activationToken.tsx
@@ -60,13 +60,13 @@ export const action: ActionFunction = async ({
       password,
     });
 
-    const {session, sessionStorage} = context;
+    const {session} = context;
 
     session.set('customerAccessToken', accessToken);
 
     return redirect(lang ? `${lang}/account` : '/account', {
       headers: {
-        'Set-Cookie': await sessionStorage.commitSession(session),
+        'Set-Cookie': await session.commit(),
       },
     });
   } catch (error: any) {
diff --git a/templates/demo-store/app/routes/account/__public/login.tsx b/templates/demo-store/app/routes/account/__public/login.tsx
index 49820f5c67..2675e7ed19 100644
--- a/templates/demo-store/app/routes/account/__public/login.tsx
+++ b/templates/demo-store/app/routes/account/__public/login.tsx
@@ -52,12 +52,12 @@ export const action: ActionFunction = async ({request, context, params}) => {
 
   try {
     const customerAccessToken = await login(context, {email, password});
-    const {session, sessionStorage} = context;
+    const {session} = context;
     session.set('customerAccessToken', customerAccessToken);
 
     return redirect(params.lang ? `${params.lang}/account` : '/account', {
       headers: {
-        'Set-Cookie': await sessionStorage.commitSession(session),
+        'Set-Cookie': await session.commit(),
       },
     });
   } catch (error: any) {
diff --git a/templates/demo-store/app/routes/account/__public/register.tsx b/templates/demo-store/app/routes/account/__public/register.tsx
index 33876f30e6..724cec79e2 100644
--- a/templates/demo-store/app/routes/account/__public/register.tsx
+++ b/templates/demo-store/app/routes/account/__public/register.tsx
@@ -29,7 +29,7 @@ type ActionData = {
 const badRequest = (data: ActionData) => json(data, {status: 400});
 
 export const action: ActionFunction = async ({request, context, params}) => {
-  const {session, sessionStorage} = context;
+  const {session} = context;
   const formData = await request.formData();
 
   const email = formData.get('email');
@@ -53,7 +53,7 @@ export const action: ActionFunction = async ({request, context, params}) => {
 
     return redirect(params.lang ? `${params.lang}/account` : '/account', {
       headers: {
-        'Set-Cookie': await sessionStorage.commitSession(session),
+        'Set-Cookie': await session.commit(),
       },
     });
   } catch (error: any) {
diff --git a/templates/demo-store/app/routes/account/__public/reset.$id.$resetToken.tsx b/templates/demo-store/app/routes/account/__public/reset.$id.$resetToken.tsx
index 9238127ce9..3cbc002673 100644
--- a/templates/demo-store/app/routes/account/__public/reset.$id.$resetToken.tsx
+++ b/templates/demo-store/app/routes/account/__public/reset.$id.$resetToken.tsx
@@ -56,13 +56,13 @@ export const action: ActionFunction = async ({
       password,
     });
 
-    const {session, sessionStorage} = context;
+    const {session} = context;
 
     session.set('customerAccessToken', accessToken);
 
     return redirect(lang ? `${lang}/account` : '/account', {
       headers: {
-        'Set-Cookie': await sessionStorage.commitSession(session),
+        'Set-Cookie': await session.commit(),
       },
     });
   } catch (error: any) {
diff --git a/templates/demo-store/oxygen.ts b/templates/demo-store/oxygen.ts
index 04763dcfb9..87081ea968 100644
--- a/templates/demo-store/oxygen.ts
+++ b/templates/demo-store/oxygen.ts
@@ -1,7 +1,7 @@
 import {createRequestHandler} from '@shopify/hydrogen-remix';
 // The build remix app provided by remix build
 import * as remixBuild from 'remix-build';
-import {getSession} from '~/lib/session.server';
+import {HydrogenSession} from '~/lib/session.server';
 
 declare const process: {env: {NODE_ENV: string}};
 
@@ -23,9 +23,7 @@ export default {
       return new Response('Internal Server Error', {status: 500});
     }
 
-    const {session, sessionStorage} = await getSession(request, [
-      env.SESSION_SECRET,
-    ]);
+    const session = await HydrogenSession.init(request, [env.SESSION_SECRET]);
 
     try {
       return await requestHandler(
@@ -41,7 +39,6 @@ export default {
         },
         {
           session,
-          sessionStorage,
         },
       );
     } catch (error) {
diff --git a/templates/demo-store/types/remix.env.d.ts b/templates/demo-store/types/remix.env.d.ts
index c105529011..ee932e0386 100644
--- a/templates/demo-store/types/remix.env.d.ts
+++ b/templates/demo-store/types/remix.env.d.ts
@@ -3,6 +3,7 @@
 
 import {Session, SessionStorage} from '@shopify/hydrogen-remix';
 import '@shopify/hydrogen-remix';
+import {HydrogenSession} from '~/lib/session.server';
 
 /**
  * Declare local additions to `AppLoadContext` to include the session utilities we injected in `oxygen.ts`.
@@ -10,7 +11,6 @@ import '@shopify/hydrogen-remix';
  */
 declare module '@shopify/hydrogen-remix' {
   export interface AppLoadContext {
-    session: Session;
-    sessionStorage: SessionStorage;
+    session: HydrogenSession;
   }
 }