fix: rework sessions

packages/auth/client/index.ts

@@ -1,10 +1,12 @@
 import type { ClientResponse, InferRequestType } from 'hono/client';
 import { hc } from 'hono/client';
+import superjson from 'superjson';
 
 import { NEXT_PUBLIC_WEBAPP_URL } from '@documenso/lib/constants/app';
 import { AppError } from '@documenso/lib/errors/app-error';
 
 import type { AuthAppType } from '../server';
+import type { SessionValidationResult } from '../server/lib/session/session';
 import { handleSignInRedirect } from '../server/lib/utils/redirect';
 import type {
   TDisableTwoFactorRequestSchema,
@@ -45,8 +47,14 @@ export class AuthClient {
     window.location.href = redirectPath ?? this.signOutredirectPath;
   }
 
-  public async session() {
-    return this.client.session.$get();
+  public async getSession() {
+    const response = await this.client['session-json'].$get();
+
+    await this.handleError(response);
+
+    const result = await response.json();
+
+    return superjson.deserialize<SessionValidationResult>(result);
   }
 
   private async handleError<T>(response: ClientResponse<T>): Promise<void> {

packages/auth/server/routes/session.ts

@@ -1,10 +1,17 @@
 import { Hono } from 'hono';
+import superjson from 'superjson';
 
 import type { SessionValidationResult } from '../lib/session/session';
 import { getOptionalSession } from '../lib/utils/get-session';
 
-export const sessionRoute = new Hono().get('/session', async (c) => {
-  const session: SessionValidationResult = await getOptionalSession(c);
+export const sessionRoute = new Hono()
+  .get('/session', async (c) => {
+    const session: SessionValidationResult = await getOptionalSession(c);
 
-  return c.json(session);
-});
+    return c.json(session);
+  })
+  .get('/session-json', async (c) => {
+    const session: SessionValidationResult = await getOptionalSession(c);
+
+    return c.json(superjson.serialize(session));
+  });

packages/lib/client-only/providers/session.tsx

@@ -1,25 +1,31 @@
-import { createContext, useContext } from 'react';
+import { createContext, useCallback, useContext, useEffect, useState } from 'react';
 import React from 'react';
 
-import type { SessionUser } from '@documenso/auth/server/lib/session/session';
-import type { Session } from '@documenso/prisma/client';
+import { useLocation } from 'react-router';
 
-import type { TGetTeamByUrlResponse } from '../../server-only/team/get-team';
-import type { TGetTeamsResponse } from '../../server-only/team/get-teams';
+import { authClient } from '@documenso/auth/client';
+import type { SessionUser } from '@documenso/auth/server/lib/session/session';
+import { type TGetTeamsResponse } from '@documenso/lib/server-only/team/get-teams';
+import type { Session } from '@documenso/prisma/client';
+import { trpc } from '@documenso/trpc/client';
 
 export type AppSession = {
   session: Session;
   user: SessionUser;
-  currentTeam: TGetTeamByUrlResponse | null;
   teams: TGetTeamsResponse;
 };
 
 interface SessionProviderProps {
   children: React.ReactNode;
-  session: AppSession | null;
+  initialSession: AppSession | null;
 }
 
-const SessionContext = createContext<AppSession | null>(null);
+interface SessionContextValue {
+  sessionData: AppSession | null;
+  refresh: () => Promise<void>;
+}
+
+const SessionContext = createContext<SessionContextValue | null>(null);
 
 export const useSession = () => {
   const context = useContext(SessionContext);
@@ -28,18 +34,78 @@ export const useSession = () => {
     throw new Error('useSession must be used within a SessionProvider');
   }
 
-  return context;
+  if (!context.sessionData) {
+    throw new Error('Session not found');
+  }
+
+  return {
+    ...context.sessionData,
+    refresh: context.refresh,
+  };
 };
 
 export const useOptionalSession = () => {
-  return (
-    useContext(SessionContext) || {
-      user: null,
-      session: null,
-    }
-  );
+  const context = useContext(SessionContext);
+
+  if (!context) {
+    throw new Error('useOptionalSession must be used within a SessionProvider');
+  }
+
+  return context;
 };
 
-export const SessionProvider = ({ children, session }: SessionProviderProps) => {
-  return <SessionContext.Provider value={session}>{children}</SessionContext.Provider>;
+export const SessionProvider = ({ children, initialSession }: SessionProviderProps) => {
+  const [session, setSession] = useState<AppSession | null>(initialSession);
+
+  const location = useLocation();
+
+  const refreshSession = useCallback(async () => {
+    const newSession = await authClient.getSession();
+
+    if (!newSession.isAuthenticated) {
+      setSession(null);
+      return;
+    }
+
+    const teams = await trpc.team.getTeams.query().catch(() => {
+      // Todo: Log
+      return [];
+    });
+
+    setSession({
+      session: newSession.session,
+      user: newSession.user,
+      teams,
+    });
+  }, []);
+
+  useEffect(() => {
+    const onFocus = () => {
+      void refreshSession();
+    };
+
+    window.addEventListener('focus', onFocus);
+
+    return () => {
+      window.removeEventListener('focus', onFocus);
+    };
+  }, [refreshSession]);
+
+  /**
+   * Refresh session in background on navigation.
+   */
+  useEffect(() => {
+    void refreshSession();
+  }, [location.pathname]);
+
+  return (
+    <SessionContext.Provider
+      value={{
+        sessionData: session,
+        refresh: refreshSession,
+      }}
+    >
+      {children}
+    </SessionContext.Provider>
+  );
 };

packages/lib/server-only/team/get-teams.ts

@@ -1,6 +1,7 @@
 import type { z } from 'zod';
 
 import { prisma } from '@documenso/prisma';
+import { SubscriptionSchema } from '@documenso/prisma/generated/zod/modelSchema/SubscriptionSchema';
 import { TeamMemberSchema } from '@documenso/prisma/generated/zod/modelSchema/TeamMemberSchema';
 import { TeamSchema } from '@documenso/prisma/generated/zod/modelSchema/TeamSchema';
 
@@ -12,6 +13,9 @@ export const ZGetTeamsResponseSchema = TeamSchema.extend({
   currentTeamMember: TeamMemberSchema.pick({
     role: true,
   }),
+  subscription: SubscriptionSchema.pick({
+    status: true,
+  }).nullable(),
 }).array();
 
 export type TGetTeamsResponse = z.infer<typeof ZGetTeamsResponseSchema>;
@@ -26,6 +30,11 @@ export const getTeams = async ({ userId }: GetTeamsOptions): Promise<TGetTeamsRe
       },
     },
     include: {
+      subscription: {
+        select: {
+          status: true,
+        },
+      },
       members: {
         where: {
           userId,

packages/trpc/react/index.tsx

@@ -1,4 +1,4 @@
-import { useState } from 'react';
+import { useMemo, useState } from 'react';
 
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { httpBatchLink, httpLink, splitLink } from '@trpc/client';
@@ -7,7 +7,6 @@ import SuperJSON from 'superjson';
 
 import { getBaseUrl } from '@documenso/lib/universal/get-base-url';
 
-// import { getBaseUrl } from '@documenso/lib/universal/get-base-url';
 import type { AppRouter } from '../server/router';
 
 export { getQueryKey } from '@trpc/react-query';
@@ -39,24 +38,27 @@ export interface TrpcProviderProps {
 export function TrpcProvider({ children, headers }: TrpcProviderProps) {
   const [queryClient] = useState(() => new QueryClient());
 
-  const [trpcClient] = useState(() =>
-    trpc.createClient({
-      links: [
-        splitLink({
-          condition: (op) => op.context.skipBatch === true,
-          true: httpLink({
-            url: `${getBaseUrl()}/api/trpc`,
-            headers,
-            transformer: SuperJSON,
+  // May cause remounting issues.
+  const trpcClient = useMemo(
+    () =>
+      trpc.createClient({
+        links: [
+          splitLink({
+            condition: (op) => op.context.skipBatch === true,
+            true: httpLink({
+              url: `${getBaseUrl()}/api/trpc`,
+              headers,
+              transformer: SuperJSON,
+            }),
+            false: httpBatchLink({
+              url: `${getBaseUrl()}/api/trpc`,
+              headers,
+              transformer: SuperJSON,
+            }),
           }),
-          false: httpBatchLink({
-            url: `${getBaseUrl()}/api/trpc`,
-            headers,
-            transformer: SuperJSON,
-          }),
-        }),
-      ],
-    }),
+        ],
+      }),
+    [headers],
   );
 
   return (