Implement logout functionality and improve error handling across the app

apps/mail/app/api/driver/connections/route.ts

@@ -1,4 +1,10 @@
-import { processIP, getRatelimitModule, checkRateLimit, getAuthenticatedUserId } from '../../utils';
+import {
+  processIP,
+  getRatelimitModule,
+  checkRateLimit,
+  getAuthenticatedUserId,
+  logoutUser,
+} from '../../utils';
 import { NextRequest, NextResponse } from 'next/server';
 import { Ratelimit } from '@upstash/ratelimit';
 import { connection } from '@zero/db/schema';
@@ -7,33 +13,39 @@ import { eq } from 'drizzle-orm';
 import { db } from '@zero/db';
 
 export const GET = async (req: NextRequest) => {
-  const userId = await getAuthenticatedUserId();
-  const finalIp = processIP(req);
-  const ratelimit = getRatelimitModule({
-    prefix: `ratelimit:get-connections-${userId}`,
-    limiter: Ratelimit.slidingWindow(60, '1m'),
-  });
-  const { success, headers } = await checkRateLimit(ratelimit, finalIp);
-  if (!success) {
-    return NextResponse.json(
-      { error: 'Too many requests. Please try again later.' },
-      { status: 429, headers },
-    );
-  }
+  try {
+    const userId = await getAuthenticatedUserId();
+    const finalIp = processIP(req);
+    const ratelimit = getRatelimitModule({
+      prefix: `ratelimit:get-connections-${userId}`,
+      limiter: Ratelimit.slidingWindow(60, '1m'),
+    });
+    const { success, headers } = await checkRateLimit(ratelimit, finalIp);
+    if (!success) {
+      return NextResponse.json(
+        { error: 'Too many requests. Please try again later.' },
+        { status: 429, headers },
+      );
+    }
 
-  const connections = (await db
-    .select({
-      id: connection.id,
-      email: connection.email,
-      name: connection.name,
-      picture: connection.picture,
-      createdAt: connection.createdAt,
-    })
-    .from(connection)
-    .where(eq(connection.userId, userId))) as IConnection[];
+    const connections = (await db
+      .select({
+        id: connection.id,
+        email: connection.email,
+        name: connection.name,
+        picture: connection.picture,
+        createdAt: connection.createdAt,
+      })
+      .from(connection)
+      .where(eq(connection.userId, userId))) as IConnection[];
 
-  return NextResponse.json(connections, {
-    status: 200,
-    headers,
-  });
+    return NextResponse.json(connections, {
+      status: 200,
+      headers,
+    });
+  } catch (error) {
+    console.warn('Error getting connections:', error);
+    await logoutUser();
+    return NextResponse.json([]);
+  }
 };

apps/mail/app/api/driver/count/route.ts

@@ -1,26 +1,38 @@
-import { checkRateLimit, getAuthenticatedUserId, getRatelimitModule, processIP } from '../../utils';
+import {
+  checkRateLimit,
+  getAuthenticatedUserId,
+  getRatelimitModule,
+  logoutUser,
+  processIP,
+} from '../../utils';
 import { type NextRequest, NextResponse } from 'next/server';
 import { getActiveDriver } from '@/actions/utils';
 import { Ratelimit } from '@upstash/ratelimit';
 
 export const GET = async (req: NextRequest) => {
-  const userId = await getAuthenticatedUserId();
-  const finalIp = processIP(req);
-  const ratelimit = getRatelimitModule({
-    prefix: `ratelimit:get-count-${userId}`,
-    limiter: Ratelimit.slidingWindow(60, '1m'),
-  });
-  const { success, headers } = await checkRateLimit(ratelimit, finalIp);
-  if (!success) {
-    return NextResponse.json(
-      { error: 'Too many requests. Please try again later.' },
-      { status: 429, headers },
-    );
+  try {
+    const userId = await getAuthenticatedUserId();
+    const finalIp = processIP(req);
+    const ratelimit = getRatelimitModule({
+      prefix: `ratelimit:get-count-${userId}`,
+      limiter: Ratelimit.slidingWindow(60, '1m'),
+    });
+    const { success, headers } = await checkRateLimit(ratelimit, finalIp);
+    if (!success) {
+      return NextResponse.json(
+        { error: 'Too many requests. Please try again later.' },
+        { status: 429, headers },
+      );
+    }
+    const driver = await getActiveDriver();
+    const count = await driver.count();
+    return NextResponse.json(count, {
+      status: 200,
+      headers,
+    });
+  } catch (error) {
+    console.warn('Error getting count:', error);
+    await logoutUser();
+    return NextResponse.json({});
   }
-  const driver = await getActiveDriver();
-  const count = await driver.count();
-  return NextResponse.json(count, {
-    status: 200,
-    headers,
-  });
 };

apps/mail/app/api/driver/google.ts

@@ -129,22 +129,12 @@ export const driver = async (config: IConfig): Promise<MailManager> => {
     process.env.GOOGLE_REDIRECT_URI as string,
   );
 
-  const getScope = () =>
-    [
-      'https://www.googleapis.com/auth/gmail.modify',
-      'https://www.googleapis.com/auth/userinfo.profile',
-      'https://www.googleapis.com/auth/userinfo.email',
-    ].join(' ');
+  const getScope = () => ['https://www.googleapis.com/auth/gmail.modify'].join(' ');
   if (config.auth) {
     auth.setCredentials({
       refresh_token: config.auth.refresh_token,
       scope: getScope(),
     });
-    if (process.env.NODE_ENV === 'production') {
-      EnableBrain()
-        .then(() => console.log('✅ Driver: Enabled'))
-        .catch(() => console.log('✅ Driver: Enabled'));
-    }
   }
   const parse = ({
     id,
@@ -535,18 +525,15 @@ export const driver = async (config: IConfig): Promise<MailManager> => {
       const { folder: normalizedFolder, q: normalizedQ } = normalizeSearch(folder, q ?? '');
       const labelIds = [..._labelIds];
       if (normalizedFolder) labelIds.push(normalizedFolder.toUpperCase());
-
-      return withExponentialBackoff(async () => {
-        const res = await gmail.users.threads.list({
-          userId: 'me',
-          q: normalizedQ ? normalizedQ : undefined,
-          labelIds: folder === 'inbox' ? labelIds : [],
-          maxResults,
-          pageToken: pageToken ? pageToken : undefined,
-          quotaUser: config.auth?.email,
-        });
-        return { ...res.data, threads: res.data.threads } as any;
+      const res = await gmail.users.threads.list({
+        userId: 'me',
+        q: normalizedQ ? normalizedQ : undefined,
+        labelIds: folder === 'inbox' ? labelIds : [],
+        maxResults,
+        pageToken: pageToken ? pageToken : undefined,
+        quotaUser: config.auth?.email,
       });
+      return { ...res.data, threads: res.data.threads } as any;
     },
     get: async (id: string) => {
       return withExponentialBackoff(async () => {

apps/mail/app/api/driver/notes/route.ts

@@ -1,33 +1,45 @@
-import { processIP, getRatelimitModule, checkRateLimit, getAuthenticatedUserId } from '../../utils';
+import {
+  processIP,
+  getRatelimitModule,
+  checkRateLimit,
+  getAuthenticatedUserId,
+  logoutUser,
+} from '../../utils';
 import { NextRequest, NextResponse } from 'next/server';
 import { fetchThreadNotes } from '@/actions/notes';
 import { Ratelimit } from '@upstash/ratelimit';
 import { notesManager } from '../../notes/db';
 
 export const GET = async (req: NextRequest) => {
-  const userId = await getAuthenticatedUserId();
-  const finalIp = processIP(req);
-  const ratelimit = getRatelimitModule({
-    prefix: `ratelimit:get-thread-notes-${userId}`,
-    limiter: Ratelimit.slidingWindow(60, '1m'),
-  });
-  const { success, headers } = await checkRateLimit(ratelimit, finalIp);
-  if (!success) {
-    return NextResponse.json(
-      { error: 'Too many requests. Please try again later.' },
-      { status: 429, headers },
-    );
-  }
-  const searchParams = req.nextUrl.searchParams;
+  try {
+    const userId = await getAuthenticatedUserId();
+    const finalIp = processIP(req);
+    const ratelimit = getRatelimitModule({
+      prefix: `ratelimit:get-thread-notes-${userId}`,
+      limiter: Ratelimit.slidingWindow(60, '1m'),
+    });
+    const { success, headers } = await checkRateLimit(ratelimit, finalIp);
+    if (!success) {
+      return NextResponse.json(
+        { error: 'Too many requests. Please try again later.' },
+        { status: 429, headers },
+      );
+    }
+    const searchParams = req.nextUrl.searchParams;
 
-  if (!searchParams.get('threadId')) {
-    return NextResponse.json({ error: 'Missing threadId' }, { status: 400 });
-  }
+    if (!searchParams.get('threadId')) {
+      return NextResponse.json({ error: 'Missing threadId' }, { status: 400 });
+    }
 
-  const notes = await notesManager.getThreadNotes(userId, searchParams.get('threadId')!);
+    const notes = await notesManager.getThreadNotes(userId, searchParams.get('threadId')!);
 
-  return NextResponse.json(notes, {
-    status: 200,
-    headers,
-  });
+    return NextResponse.json(notes, {
+      status: 200,
+      headers,
+    });
+  } catch (error) {
+    console.warn('Error getting thread notes:', error);
+    await logoutUser();
+    return NextResponse.json([]);
+  }
 };

apps/mail/app/api/driver/route.ts

@@ -45,6 +45,6 @@ export const GET = async (req: NextRequest) => {
   } catch (error) {
     console.warn('Error getting threads:', error);
     await logoutUser();
-    // return NextResponse.redirect('/login?error=unauthorized');
+    return NextResponse.json({ messages: [], nextPageToken: null });
   }
 };

apps/mail/app/api/utils.ts

@@ -1,4 +1,5 @@
 import { Ratelimit, Algorithm, RatelimitConfig } from '@upstash/ratelimit';
+import { deleteActiveConnection } from '@/actions/utils';
 import { redirect } from 'next/navigation';
 import { NextRequest } from 'next/server';
 import { headers } from 'next/headers';
@@ -42,7 +43,9 @@ export async function getAuthenticatedUserId(): Promise<string> {
   return session.user.id;
 }
 
+// Forcefully logout the user, this will delete the active connection
 export async function logoutUser() {
+  await deleteActiveConnection();
   const headersList = await headers();
   await auth.api.signOut({ headers: headersList });
 }

apps/mail/app/api/v1/hotkeys/route.ts

@@ -1,4 +1,6 @@
+import { checkRateLimit, getAuthenticatedUserId, getRatelimitModule } from '../../utils';
 import type { Shortcut } from '@/config/shortcuts';
+import { Ratelimit } from '@upstash/ratelimit';
 import { userHotkeys } from '@zero/db/schema';
 import { NextResponse } from 'next/server';
 import { headers } from 'next/headers';
@@ -7,33 +9,48 @@ import { eq } from 'drizzle-orm';
 import { db } from '@zero/db';
 
 export async function GET() {
-  const session = await auth.api.getSession({
-    headers: await headers(),
+  const userId = await getAuthenticatedUserId();
+
+  const ratelimit = getRatelimitModule({
+    prefix: 'ratelimit:hotkeys',
+    limiter: Ratelimit.slidingWindow(60, '1m'),
   });
-  if (!session) {
-    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
+
+  const { success, headers } = await checkRateLimit(ratelimit, userId);
+  if (!success) {
+    return NextResponse.json(
+      { error: 'Too many requests. Please try again later.' },
+      { status: 429, headers },
+    );
   }
 
-  const result = await db.select().from(userHotkeys).where(eq(userHotkeys.userId, session.user.id));
+  const result = await db.select().from(userHotkeys).where(eq(userHotkeys.userId, userId));
 
   return NextResponse.json(result[0]?.shortcuts || []);
 }
 
 export async function POST(request: Request) {
-  const session = await auth.api.getSession({
-    headers: await headers(),
+  const userId = await getAuthenticatedUserId();
+
+  const ratelimit = getRatelimitModule({
+    prefix: 'ratelimit:hotkeys-post',
+    limiter: Ratelimit.slidingWindow(60, '1m'),
   });
-  if (!session) {
-    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
-  }
 
+  const { success, headers } = await checkRateLimit(ratelimit, userId);
+  if (!success) {
+    return NextResponse.json(
+      { error: 'Too many requests. Please try again later.' },
+      { status: 429, headers },
+    );
+  }
   const shortcuts = (await request.json()) as Shortcut[];
   const now = new Date();
 
   await db
     .insert(userHotkeys)
     .values({
-      userId: session.user.id,
+      userId: userId,
       shortcuts,
       createdAt: now,
       updatedAt: now,
@@ -50,17 +67,24 @@ export async function POST(request: Request) {
 }
 
 export async function PUT(request: Request) {
-  const session = await auth.api.getSession({
-    headers: await headers(),
+  const userId = await getAuthenticatedUserId();
+
+  const ratelimit = getRatelimitModule({
+    prefix: 'ratelimit:hotkeys-put',
+    limiter: Ratelimit.slidingWindow(60, '1m'),
   });
-  if (!session) {
-    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
-  }
 
+  const { success, headers } = await checkRateLimit(ratelimit, userId);
+  if (!success) {
+    return NextResponse.json(
+      { error: 'Too many requests. Please try again later.' },
+      { status: 429, headers },
+    );
+  }
   const shortcut = (await request.json()) as Shortcut;
   const now = new Date();
 
-  const result = await db.select().from(userHotkeys).where(eq(userHotkeys.userId, session.user.id));
+  const result = await db.select().from(userHotkeys).where(eq(userHotkeys.userId, userId));
 
   const existingShortcuts = (result[0]?.shortcuts || []) as Shortcut[];
   const updatedShortcuts = existingShortcuts.map((s: Shortcut) =>
@@ -74,7 +98,7 @@ export async function PUT(request: Request) {
   await db
     .insert(userHotkeys)
     .values({
-      userId: session.user.id,
+      userId,
       shortcuts: updatedShortcuts,
       createdAt: now,
       updatedAt: now,