fix(auth): improve auth error handling and fix Link CORS issue

webapp/src/app/(root)/page.tsx

@@ -1,12 +1,12 @@
 import { prisma } from '@/lib/prisma';
-import { getSession } from '@/lib/auth';
+import { getAuthSession } from '@/lib/auth';
 import TodoItemComponent from './components/TodoItem';
 import CreateTodoForm from './components/CreateTodoForm';
 import { TodoItemStatus } from '@prisma/client';
 import Header from '@/components/Header';
 
 export default async function Home() {
-  const { userId } = await getSession();
+  const { userId } = await getAuthSession();
 
   const todos = await prisma.todoItem.findMany({
     where: {

webapp/src/app/api/cognito-token/route.ts

@@ -1,21 +1,15 @@
 import { NextResponse } from 'next/server';
-import { cookies } from 'next/headers';
-import { fetchAuthSession } from 'aws-amplify/auth/server';
-import { runWithAmplifyServerContext } from '@/lib/amplifyServerUtils';
+import { tryGetAuthSession } from '@/lib/auth';
 
 export async function GET() {
   try {
-    const session = await runWithAmplifyServerContext({
-      nextServerContext: { cookies },
-      operation: (contextSpec) => fetchAuthSession(contextSpec),
-    });
-
-    if (session.tokens?.accessToken == null) {
+    const session = await tryGetAuthSession();
+    if (!session) {
       return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
     }
 
     return NextResponse.json({
-      accessToken: session.tokens.accessToken.toString(),
+      accessToken: session.accessToken,
     });
   } catch (error) {
     console.error('Error fetching Cognito token:', error);

webapp/src/app/auth-callback/page.tsx

@@ -1,25 +1,16 @@
 import { redirect } from 'next/navigation';
-import { getSession, UserNotCreatedError } from '@/lib/auth';
+import { getAuthSession } from '@/lib/auth';
 import { prisma } from '@/lib/prisma';
 
 export const dynamic = 'force-dynamic';
 
 export default async function AuthCallbackPage() {
-  try {
-    await getSession();
-  } catch (e) {
-    console.log(e);
-    if (e instanceof UserNotCreatedError) {
-      const userId = e.userId;
-      console.log(userId);
-      await prisma.user.create({
-        data: {
-          id: userId,
-        },
-      });
-    } else {
-      throw e;
-    }
+  const { userId } = await getAuthSession();
+
+  const user = await prisma.user.findUnique({ where: { id: userId } });
+  if (user == null) {
+    await prisma.user.create({ data: { id: userId } });
   }
+
   redirect('/');
 }

webapp/src/app/sign-in/page.tsx

@@ -1,5 +1,3 @@
-import Link from 'next/link';
-
 export default function SignInPage() {
   return (
     <div className="min-h-screen bg-gray-50 flex flex-col justify-center py-12 sm:px-6 lg:px-8">
@@ -15,15 +13,18 @@ export default function SignInPage() {
               Please sign in with your Cognito account to continue
             </p>
 
-            <Link
+            {/* Use <a> instead of <Link> to trigger a full-page navigation.
+                The sign-in route returns a 302 redirect to Cognito, which
+                would cause a CORS error if fetched via client-side navigation. */}
+            {/* eslint-disable-next-line @next/next/no-html-link-for-pages */}
+            <a
               href="/api/auth/sign-in"
               // you can add a query string to change the locale of cognito managed login page.
               // href="/api/auth/sign-in?lang=ja"
               className="w-full flex justify-center py-3 px-4 border border-transparent rounded-md shadow-sm text-sm font-medium text-white bg-indigo-600 hover:bg-indigo-700 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-indigo-500"
-              prefetch={false} // prevent CORS error
             >
               Sign in with Cognito
-            </Link>
+            </a>
           </div>
         </div>
       </div>

webapp/src/components/Header.tsx

@@ -1,11 +1,6 @@
-'use client';
-
 import Link from 'next/link';
-import { useRouter } from 'next/navigation';
 
 export default function Header() {
-  const router = useRouter();
-
   return (
     <header className="bg-indigo-600 text-white shadow-md">
       <div className="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8">
@@ -16,13 +11,16 @@ export default function Header() {
             </Link>
           </div>
           <div>
-            <Link
+            {/* Use <a> instead of <Link> to trigger a full-page navigation.
+                The sign-out route returns a 302 redirect to Cognito, which
+                would cause a CORS error if fetched via client-side navigation. */}
+            {/* eslint-disable-next-line @next/next/no-html-link-for-pages */}
+            <a
               href="/api/auth/sign-out"
               className="inline-flex items-center px-4 py-2 border border-transparent text-sm font-medium rounded-md text-indigo-600 bg-white hover:bg-gray-100 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-indigo-500"
-              prefetch={false} // prevent CORS error
             >
               Sign Out
-            </Link>
+            </a>
           </div>
         </div>
       </div>

webapp/src/lib/auth.ts

@@ -1,38 +1,58 @@
+import { cache } from 'react';
 import { cookies } from 'next/headers';
 import { fetchAuthSession } from 'aws-amplify/auth/server';
 import { runWithAmplifyServerContext } from '@/lib/amplifyServerUtils';
 import { prisma } from '@/lib/prisma';
 
-export class UserNotCreatedError {
-  constructor(public readonly userId: string) {}
-}
-
-export async function getSession() {
+/**
+ * Get the authenticated session without DB access.
+ * Use when only userId/email/accessToken is needed.
+ * Memoized per request via React cache().
+ */
+export const getAuthSession = cache(async () => {
   const session = await runWithAmplifyServerContext({
     nextServerContext: { cookies },
     operation: (contextSpec) => fetchAuthSession(contextSpec),
   });
   if (session.userSub == null || session.tokens?.idToken == null || session.tokens?.accessToken == null) {
     throw new Error('session not found');
   }
-  const userId = session.userSub;
   const email = session.tokens.idToken.payload.email;
   if (typeof email != 'string') {
-    throw new Error(`invalid email ${userId}.`);
-  }
-  const user = await prisma.user.findUnique({
-    where: {
-      id: userId,
-    },
-  });
-  if (user == null) {
-    throw new UserNotCreatedError(userId);
+    throw new Error(`invalid email ${session.userSub}.`);
   }
-
   return {
-    userId: user.id,
+    userId: session.userSub,
     email,
     accessToken: session.tokens.accessToken.toString(),
-    user,
   };
+});
+
+/**
+ * Try to get the authenticated session, returning null on failure.
+ * Use in API Routes to avoid try/catch boilerplate for auth checks.
+ */
+export async function tryGetAuthSession() {
+  try {
+    return await getAuthSession();
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Get the authenticated session with the User record from DB.
+ * Memoized per request via React cache().
+ */
+export const getSessionWithUser = cache(async () => {
+  const auth = await getAuthSession();
+  const user = await prisma.user.findUnique({ where: { id: auth.userId } });
+  if (user == null) {
+    throw new UserNotFoundError(auth.userId);
+  }
+  return { ...auth, user };
+});
+
+export class UserNotFoundError {
+  constructor(public readonly userId: string) {}
 }

webapp/src/lib/safe-action.ts

@@ -1,8 +1,5 @@
-import { prisma } from '@/lib/prisma';
-import { runWithAmplifyServerContext } from '@/lib/amplifyServerUtils';
-import { getCurrentUser } from 'aws-amplify/auth/server';
+import { getSessionWithUser } from '@/lib/auth';
 import { createSafeActionClient, DEFAULT_SERVER_ERROR_MESSAGE } from 'next-safe-action';
-import { cookies } from 'next/headers';
 
 export class MyCustomError extends Error {
   constructor(message: string) {
@@ -28,24 +25,6 @@ const actionClient = createSafeActionClient({
 });
 
 export const authActionClient = actionClient.use(async ({ next }) => {
-  const currentUser = await runWithAmplifyServerContext({
-    nextServerContext: { cookies },
-    operation: (contextSpec) => getCurrentUser(contextSpec),
-  });
-
-  if (!currentUser) {
-    throw new Error('Session is not valid!');
-  }
-
-  const user = await prisma.user.findUnique({
-    where: {
-      id: currentUser.userId,
-    },
-  });
-
-  if (user == null) {
-    throw new Error('user not found');
-  }
-
+  const { user } = await getSessionWithUser();
   return next({ ctx: { userId: user.id } });
 });