Refactor mail API and components by reorganizing imports, enhancing error handling, and improving label management. Removed unused thread-labels route and updated label handling in various components for better consistency and performance.

2026-06-29 23:37:05 +00:00 · 2025-04-24 18:05:03 -07:00
parent 45e92440ca
commit fc6d6376a8
12 changed files with 114 additions and 222 deletions
--- a/apps/mail/app/api/ai-search/route.ts
+++ b/apps/mail/app/api/ai-search/route.ts
@@ -1,10 +1,10 @@
-import { generateText, tool } from 'ai';
-import { z } from 'zod';
-import { auth } from '@/lib/auth';
-import { headers } from 'next/headers';
-import { openai } from '@ai-sdk/openai';
 import { getActiveDriver } from '@/actions/utils';
 import { type gmail_v1 } from 'googleapis';
+import { openai } from '@ai-sdk/openai';
+import { generateText, tool } from 'ai';
+import { headers } from 'next/headers';
+import { auth } from '@/lib/auth';
+import { z } from 'zod';

 // Define our email search tool
 const emailSearchTool = tool({
@@ -74,7 +74,7 @@ For sender searches, use the exact name/email provided by the user.
 For domain searches, search both the from: field and general content.
 For subject/content searches, include relevant synonyms and related terms.

-Important: This is a search-only assistant. Do not generate email content or handle email composition requests.`
+Important: This is a search-only assistant. Do not generate email content or handle email composition requests.`,
    };

    const { text, steps } = await generateText({
@@ -89,8 +89,8 @@ Important: This is a search-only assistant. Do not generate email content or han

    // Extract the search query and explanation from the tool call
    const toolCall = steps
-      .flatMap(step => step.toolCalls)
-      .find(call => call.toolName === 'emailSearch');
+      .flatMap((step) => step.toolCalls)
+      .find((call) => call.toolName === 'emailSearch');

    if (!toolCall?.args?.query) {
      throw new Error('Failed to generate search query');
@@ -104,46 +104,47 @@ Important: This is a search-only assistant. Do not generate email content or han
    const results = await driver.list('', searchQuery, 20);

    // Process the results - use the raw response from Gmail API
-    const processResultPromises = results?.threads?.map(async thread => {
-      const rawThread = thread as gmail_v1.Schema$Thread;
-      
-      try {
-        // Get the thread data using our existing driver
-        const threadData = await driver.get(rawThread.id!);
-        const firstMessage = threadData.messages[0];
-        
-        if (!firstMessage) {
-          throw new Error('No messages found in thread');
-        }
+    const processResultPromises =
+      results?.threads?.map(async (thread) => {
+        const rawThread = thread as gmail_v1.Schema$Thread;

-        return {
-          id: rawThread.id!,
-          snippet: rawThread.snippet || '',
-          historyId: rawThread.historyId,
-          subject: firstMessage.subject || 'No subject',
-          from: firstMessage.sender.email || firstMessage.sender.name || 'Unknown sender',
-        };
-      } catch (error) {
-        console.error('Error processing thread:', error);
-        return {
-          id: rawThread.id!,
-          snippet: rawThread.snippet || '',
-          historyId: rawThread.historyId,
-          subject: 'Error loading subject',
-          from: 'Error loading sender',
-        };
-      }
-    }) || [];
+        try {
+          // Get the thread data using our existing driver
+          const threadData = await driver.get(rawThread.id!);
+          const firstMessage = threadData.messages[0];
+
+          if (!firstMessage) {
+            throw new Error('No messages found in thread');
+          }
+
+          return {
+            id: rawThread.id!,
+            snippet: rawThread.snippet || '',
+            historyId: rawThread.historyId,
+            subject: firstMessage.subject || 'No subject',
+            from: firstMessage.sender.email || firstMessage.sender.name || 'Unknown sender',
+          };
+        } catch (error) {
+          console.error('Error processing thread:', error);
+          return {
+            id: rawThread.id!,
+            snippet: rawThread.snippet || '',
+            historyId: rawThread.historyId,
+            subject: 'Error loading subject',
+            from: 'Error loading sender',
+          };
+        }
+      }) || [];

    // Resolve all promises
    const resolvedResults = await Promise.all(processResultPromises);

    // Create a natural response using the AI's text and search results
    const hasResults = resolvedResults.length > 0;
-    
+
    // Let the AI's response text lead the way
    let summary = text;
-    
+
    // Add result information
    if (hasResults) {
      summary += `\n\nI found ${resolvedResults.length} email${resolvedResults.length === 1 ? '' : 's'} ${searchExplanation}. Here ${resolvedResults.length === 1 ? 'it is' : 'they are'}:`;
@@ -151,20 +152,22 @@ Important: This is a search-only assistant. Do not generate email content or han
      summary += `\n\nI couldn't find any emails ${searchExplanation}. Would you like to try a different search?`;
    }

-    return new Response(JSON.stringify({
-      content: summary,
-      searchQuery,
-      searchDisplay: `Searched for "${searchQuery}"`,
-      results: resolvedResults,
-    }), {
-      headers: { 'Content-Type': 'application/json' },
-    });
-
+    return new Response(
+      JSON.stringify({
+        content: summary,
+        searchQuery,
+        searchDisplay: `Searched for "${searchQuery}"`,
+        results: resolvedResults,
+      }),
+      {
+        headers: { 'Content-Type': 'application/json' },
+      },
+    );
  } catch (error) {
    console.error('AI Search error:', error);
    return new Response(JSON.stringify({ error: 'Failed to process search request' }), {
-      status: 500,
+      status: 400,
      headers: { 'Content-Type': 'application/json' },
    });
  }
-} 
+}
--- a/apps/mail/app/api/chat/route.ts
+++ b/apps/mail/app/api/chat/route.ts
@@ -6,9 +6,10 @@ export async function POST(req: Request) {
    const { messages, context } = await req.json();

    const lastMessage = messages[messages.length - 1].content;
-    
-    let systemPrompt = 'You are a helpful AI assistant. Provide clear, concise, and accurate responses.';
-    
+
+    let systemPrompt =
+      'You are a helpful AI assistant. Provide clear, concise, and accurate responses.';
+
    // If this is an email request, modify the system prompt
    if (context?.isEmailRequest) {
      systemPrompt = `You are an email writing assistant. Generate professional, well-structured emails.
@@ -30,12 +31,12 @@ Output format:
    const { completion } = await generateCompletions({
      model: 'llama3-8b-8192',
      systemPrompt,
-      prompt: context?.isEmailRequest 
+      prompt: context?.isEmailRequest
        ? `Generate a professional email for the following request: ${lastMessage}`
        : lastMessage,
      temperature: 0.7,
      max_tokens: 500,
-      userName: 'User'
+      userName: 'User',
    });

    // If this was an email request, try to parse the JSON response
@@ -52,9 +53,6 @@ Output format:
    return NextResponse.json({ content: completion });
  } catch (error) {
    console.error('Chat API Error:', error);
-    return NextResponse.json(
-      { error: 'Failed to generate response' },
-      { status: 500 }
-    );
+    return NextResponse.json({ error: 'Failed to generate response' }, { status: 400 });
  }
-} 
+}
--- a/apps/mail/app/api/driver/google.ts
+++ b/apps/mail/app/api/driver/google.ts
@@ -260,7 +260,7 @@ export const driver = async (config: IConfig): Promise<MailManager> => {
      threadId: threadId || '',
      title: snippet ? he.decode(snippet).trim() : 'ERROR',
      tls: wasSentWithTLS(receivedHeaders) || !!hasTLSReport,
-      tags: labelIds || [],
+      tags: labelIds?.map((l) => ({ id: l, name: l })) || [],
      listUnsubscribe,
      listUnsubscribePost,
      replyTo,
@@ -978,7 +978,18 @@ export const driver = async (config: IConfig): Promise<MailManager> => {
      const res = await gmail.users.labels.list({
        userId: 'me',
      });
-      return res.data.labels;
+      // wtf google, null values for EVERYTHING?
+      return (
+        res.data.labels?.map((label) => ({
+          id: label.id ?? '',
+          name: label.name ?? '',
+          type: label.type ?? '',
+          color: {
+            backgroundColor: label.color?.backgroundColor ?? '',
+            textColor: label.color?.textColor ?? '',
+          },
+        })) ?? []
+      );
    },
    getLabel: async (labelId: string) => {
      const res = await gmail.users.labels.get({
--- a/apps/mail/app/api/driver/types.ts
+++ b/apps/mail/app/api/driver/types.ts
@@ -1,4 +1,5 @@
 import { type IOutgoingMessage, type InitialThread, type ParsedMessage } from '@/types';
+import { Label } from '@/hooks/use-labels';

 export interface IGetThreadResponse {
  messages: ParsedMessage[];
@@ -48,7 +49,7 @@ export interface MailManager {
    options: { addLabels: string[]; removeLabels: string[] },
  ): Promise<void>;
  getAttachment(messageId: string, attachmentId: string): Promise<string | undefined>;
-  getUserLabels(): Promise<any>;
+  getUserLabels(): Promise<Label[]>;
  getLabel: (labelId: string) => Promise<any>;
  createLabel(label: {
    name: string;
--- a/apps/mail/app/api/v1/labels/route.ts
+++ b/apps/mail/app/api/v1/labels/route.ts
@@ -2,15 +2,7 @@ import { processIP, getRatelimitModule, checkRateLimit, getAuthenticatedUserId }
 import { NextRequest, NextResponse } from 'next/server';
 import { getActiveDriver } from '@/actions/utils';
 import { Ratelimit } from '@upstash/ratelimit';
-
-interface Label {
-  name: string;
-  color?: {
-    backgroundColor: string;
-    textColor: string;
-  };
-  type?: 'user' | 'system';
-}
+import { Label } from '@/hooks/use-labels';

 export async function GET(req: NextRequest) {
  const userId = await getAuthenticatedUserId();
@@ -30,17 +22,14 @@ export async function GET(req: NextRequest) {

  try {
    const driver = await getActiveDriver();
-    if (!driver) {
-      return NextResponse.json({ error: 'Email driver not configured' }, { status: 500 });
-    }
    const labels = await driver.getUserLabels();
    if (!labels) {
      return NextResponse.json([], { status: 200 });
    }
-    return NextResponse.json(labels.filter((label: Label) => label.type === 'user'));
+    return NextResponse.json(labels.filter((label) => label.type === 'user'));
  } catch (error) {
    console.error('Error fetching labels:', error);
-    return NextResponse.json({ error: 'Failed to fetch labels' }, { status: 500 });
+    return NextResponse.json({ error: 'Failed to fetch labels' }, { status: 400 });
  }
 }

@@ -67,11 +56,11 @@ export async function POST(req: NextRequest) {
      type: 'user',
    };
    const driver = await getActiveDriver();
-    const result = await driver?.createLabel(label);
+    const result = await driver.createLabel(label);
    return NextResponse.json(result);
  } catch (error) {
    console.error('Error creating label:', error);
-    return NextResponse.json({ error: 'Failed to create label' }, { status: 500 });
+    return NextResponse.json({ error: 'Failed to create label' }, { status: 400 });
  }
 }

@@ -94,11 +83,11 @@ export async function PATCH(req: NextRequest) {
  try {
    const { id, ...label } = (await req.json()) as Label & { id: string } & { type: string };
    const driver = await getActiveDriver();
-    const result = await driver?.updateLabel(id, label);
+    const result = await driver.updateLabel(id, label);
    return NextResponse.json(result);
  } catch (error) {
    console.error('Error updating label:', error);
-    return NextResponse.json({ error: 'Failed to update label' }, { status: 500 });
+    return NextResponse.json({ error: 'Failed to update label' }, { status: 400 });
  }
 }

@@ -121,10 +110,10 @@ export async function DELETE(req: NextRequest) {
  try {
    const { id } = (await req.json()) as { id: string };
    const driver = await getActiveDriver();
-    await driver?.deleteLabel(id);
+    await driver.deleteLabel(id);
    return NextResponse.json({ success: true });
  } catch (error) {
    console.error('Error deleting label:', error);
-    return NextResponse.json({ error: 'Failed to delete label' }, { status: 500 });
+    return NextResponse.json({ error: 'Failed to delete label' }, { status: 400 });
  }
 }
--- a/apps/mail/app/api/v1/thread-labels/route.ts
+++ b/apps/mail/app/api/v1/thread-labels/route.ts
@@ -1,51 +0,0 @@
-import { processIP, getRatelimitModule, checkRateLimit, getAuthenticatedUserId } from '../../utils';
-import { NextRequest, NextResponse } from 'next/server';
-import { getActiveDriver } from '@/actions/utils';
-import { Ratelimit } from '@upstash/ratelimit';
-import { Label } from '@/hooks/use-labels';
-
-export async function GET(req: NextRequest) {
-  const userId = await getAuthenticatedUserId();
-  const finalIp = processIP(req);
-  const ratelimit = getRatelimitModule({
-    prefix: `ratelimit:get-thread-labels-${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 { searchParams } = new URL(req.url);
-    const ids = searchParams.get('ids');
-
-    if (!ids) {
-      return NextResponse.json({ error: 'Thread IDs are required' }, { status: 400 });
-    }
-
-    const threadIds = ids.split(',');
-    const driver = await getActiveDriver();
-
-    const labels = await Promise.all(threadIds.map(async (id) => await driver.getLabel(id)));
-
-    const userLabels: Label[] = labels
-      .filter((label): label is Label => {
-        return label && typeof label === 'object' && label.type === 'user';
-      })
-      .map((label) => ({
-        id: label.id,
-        name: label.name,
-        type: label.type,
-        color: label.color,
-      }));
-
-    return NextResponse.json(userLabels);
-  } catch (error) {
-    console.error('Error fetching thread labels:', error);
-    return NextResponse.json({ error: 'Failed to fetch thread labels' }, { status: 500 });
-  }
-}
--- a/apps/mail/components/mail/mail-display.tsx
+++ b/apps/mail/components/mail/mail-display.tsx
@@ -325,7 +325,7 @@ const MailDisplay = ({ emailData, isMuted, index, totalEmails, demo }: Props) =>
                </span>
              </div>
              <div className="flex items-center gap-4">
-                <MailDisplayLabels labels={emailData?.tags || []} />
+                <MailDisplayLabels labels={emailData?.tags.map((t) => t.name) || []} />
                <div className="bg-iconLight dark:bg-iconDark/20 relative h-3 w-0.5 rounded-full" />
                <div className="flex items-center gap-2 text-sm text-[#6D6D6D] dark:text-[#8C8C8C]">
                  {(() => {
--- a/apps/mail/components/mail/mail-list.tsx
+++ b/apps/mail/components/mail/mail-list.tsx
@@ -18,6 +18,7 @@ import { preloadThread, useThread, useThreads } from '@/hooks/use-threads';
 import { ThreadContextMenu } from '@/components/context/thread-context';
 import { Avatar, AvatarImage, AvatarFallback } from '../ui/avatar';
 import { useMailNavigation } from '@/hooks/use-mail-navigation';
+import { Label, useThreadLabels } from '@/hooks/use-labels';
 import { useSearchValue } from '@/hooks/use-search-value';
 import { markAsRead, markAsUnread } from '@/actions/mail';
 import { ScrollArea } from '@/components/ui/scroll-area';
@@ -32,7 +33,6 @@ import { RenderLabels } from './render-labels';
 import { Badge } from '@/components/ui/badge';
 import { useDraft } from '@/hooks/use-drafts';
 import { useTranslations } from 'next-intl';
-import { Label } from '@/hooks/use-labels';
 import { Button } from '../ui/button';
 import { useQueryState } from 'nuqs';
 import { Categories } from './mail';
@@ -156,6 +156,10 @@ const Thread = memo(
    const latestMessage = demo ? demoMessage : getThreadData?.latest;
    const emailContent = demo ? demoMessage?.body : getThreadData?.latest?.body;

+    const { labels: threadLabels } = useThreadLabels(
+      latestMessage ? latestMessage.tags?.map((t) => t.id!) : [],
+    );
+
    const mainSearchTerm = useMemo(() => {
      if (!searchValue.highlight) return '';
      return getMainSearchTerm(searchValue.highlight);
@@ -186,11 +190,6 @@ const Thread = memo(

    const isMailBulkSelected = mail.bulkSelected.includes(latestMessage?.threadId ?? message.id);

-    const threadLabels = useMemo(() => {
-      if (!latestMessage) return [];
-      return [...(latestMessage.tags || [])];
-    }, [latestMessage]);
-
    const isFolderInbox = folder === FOLDERS.INBOX || !folder;
    const isFolderSpam = folder === FOLDERS.SPAM;
    const isFolderSent = folder === FOLDERS.SENT;
@@ -439,7 +438,7 @@ const Thread = memo(
                          )}
                        </span>{' '}
                        <span className="flex items-center space-x-2">
-                          <RenderLabels ids={threadLabels} />
+                          <RenderLabels labels={threadLabels} />
                        </span>
                      </span>
                      {getThreadData.totalReplies > 1 ? (
@@ -781,13 +780,13 @@ export const MailList = memo(({ isCompact }: MailListProps) => {
 MailList.displayName = 'MailList';

 export const MailLabels = memo(
-  ({ labels }: { labels: string[] }) => {
+  ({ labels }: { labels: Label[] }) => {
    const t = useTranslations();

    if (!labels.length) return null;

    const visibleLabels = labels.filter(
-      (label) => !['unread', 'inbox'].includes(label.toLowerCase()),
+      (label) => !['unread', 'inbox'].includes(label.name.toLowerCase()),
    );

    if (!visibleLabels.length) return null;
@@ -795,13 +794,13 @@ export const MailLabels = memo(
    return (
      <div className={cn('flex select-none items-center')}>
        {visibleLabels.map((label) => {
-          const style = getDefaultBadgeStyle(label);
-          if (label.toLowerCase() === 'notes') {
+          const style = getDefaultBadgeStyle(label.name);
+          if (label.name.toLowerCase() === 'notes') {
            return (
-              <Tooltip key={label}>
+              <Tooltip key={label.id}>
                <TooltipTrigger asChild>
                  <Badge className="rounded-md bg-amber-100 p-1 text-amber-700 hover:bg-amber-200 dark:bg-amber-900/30 dark:text-amber-400">
-                    {getLabelIcon(label)}
+                    {getLabelIcon(label.name)}
                  </Badge>
                </TooltipTrigger>
                <TooltipContent className="hidden px-1 py-0 text-xs">
@@ -814,7 +813,7 @@ export const MailLabels = memo(
          // Skip rendering if style is "secondary" (default case)
          if (style === 'secondary') return null;

-          const normalizedLabel = getNormalizedLabelKey(label);
+          const normalizedLabel = getNormalizedLabelKey(label.name);

          let labelContent;
          switch (normalizedLabel) {
@@ -844,8 +843,8 @@ export const MailLabels = memo(
          }

          return (
-            <Badge key={label} className="rounded-md p-1" variant={style}>
-              {getLabelIcon(label)}
+            <Badge key={label.id} className="rounded-md p-1" variant={style}>
+              {getLabelIcon(label.name)}
            </Badge>
          );
        })}
--- a/apps/mail/components/mail/render-labels.tsx
+++ b/apps/mail/components/mail/render-labels.tsx
@@ -2,14 +2,13 @@

 import { PopoverContent, PopoverTrigger } from '@radix-ui/react-popover';
 import { Tooltip, TooltipContent, TooltipTrigger } from '../ui/tooltip';
-import { Label, useThreadLabels } from '@/hooks/use-labels';
 import { useSearchValue } from '@/hooks/use-search-value';
+import { Label } from '@/hooks/use-labels';
 import { Popover } from '../ui/popover';
 import { cn } from '@/lib/utils';
 import * as React from 'react';

-export const RenderLabels = ({ ids }: { ids: string[] }) => {
-  const { data: labels = [] } = useThreadLabels(ids);
+export const RenderLabels = ({ labels }: { labels: Label[] }) => {
  const [searchValue, setSearchValue] = useSearchValue();
  const label = React.useMemo(() => labels[0], [labels]);

--- a/apps/mail/hooks/use-labels.ts
+++ b/apps/mail/hooks/use-labels.ts
@@ -1,3 +1,4 @@
+import { useMemo } from 'react';
 import { toast } from 'sonner';
 import axios from 'axios';
 import useSWR from 'swr';
@@ -84,49 +85,19 @@ export function useLabels() {
    }
  };

-  const getThreadLabels = async (ids: string[]) => {
-    const { data } = await useThreadLabels(ids);
-    return data || [];
-  };
-
  return {
    labels: labels || [],
    isLoading,
    error,
-    createLabel,
-    updateLabel,
-    deleteLabel,
-    getThreadLabels,
-    refresh: () => mutate(),
  };
 }

 export function useThreadLabels(ids: string[]) {
-  const key = ids.length > 0 ? `/api/v1/thread-labels?ids=${ids.join(',')}` : null;
+  const { labels } = useLabels();
+  const threadLabels = useMemo(() => {
+    if (!labels) return [];
+    return labels.filter((label) => (label.id ? ids.includes(label.id) : false));
+  }, [labels, ids]);

-  return useSWR<Label[]>(
-    key,
-    async (url) => {
-      try {
-        const response = await fetch(url, {
-          method: 'GET',
-          headers: { 'Content-Type': 'application/json' },
-        });
-
-        if (!response.ok) {
-          throw new Error('Failed to fetch thread labels');
-        }
-
-        return response.json();
-      } catch (error) {
-        toast.error('Failed to fetch thread labels');
-        throw error;
-      }
-    },
-    {
-      revalidateOnFocus: false,
-      revalidateOnReconnect: false,
-      dedupingInterval: 5000,
-    },
-  );
+  return { labels: threadLabels };
 }
--- a/apps/mail/hooks/use-threads.ts
+++ b/apps/mail/hooks/use-threads.ts
@@ -6,6 +6,7 @@ import type { InitialThread, ParsedMessage } from '@/types';
 import { useSearchValue } from '@/hooks/use-search-value';
 import { useSession } from '@/lib/auth-client';
 import { defaultPageSize } from '@/lib/utils';
+import { useAtom, useAtomValue } from 'jotai';
 import { Label } from '@/hooks/use-labels';
 import useSWRInfinite from 'swr/infinite';
 import useSWR, { preload } from 'swr';
@@ -13,7 +14,6 @@ import { useQueryState } from 'nuqs';
 import { useMemo } from 'react';
 import { toast } from 'sonner';
 import axios from 'axios';
-import { useAtom, useAtomValue } from 'jotai';

 export const preloadThread = async (userId: string, threadId: string, connectionId: string) => {
  console.log(`🔄 Prefetching email ${threadId}...`);
@@ -154,36 +154,6 @@ export const useThreads = () => {
  };
 };

-export function useThreadLabels(ids: string[]) {
-  const key = ids.length > 0 ? `/api/v1/thread-labels?ids=${ids.join(',')}` : null;
-
-  return useSWR<Label[]>(
-    key,
-    async (url) => {
-      try {
-        const response = await fetch(url, {
-          method: 'GET',
-          headers: { 'Content-Type': 'application/json' },
-        });
-
-        if (!response.ok) {
-          throw new Error('Failed to fetch thread labels');
-        }
-
-        return response.json();
-      } catch (error) {
-        toast.error('Failed to fetch thread labels');
-        throw error;
-      }
-    },
-    {
-      revalidateOnFocus: false,
-      revalidateOnReconnect: false,
-      dedupingInterval: 5000,
-    },
-  );
-}
-
 export const useThread = (threadId: string | null) => {
  const { data: session } = useSession();
  const [_threadId] = useQueryState('threadId');
--- a/apps/mail/types/index.ts
+++ b/apps/mail/types/index.ts
@@ -1,3 +1,5 @@
+import { Label } from '@/hooks/use-labels';
+
 export interface User {
  name: string;
  email: string;
@@ -51,7 +53,7 @@ export interface ParsedMessage {
  connectionId?: string;
  title: string;
  subject: string;
-  tags: string[];
+  tags: Label[];
  sender: Sender;
  to: Sender[];
  cc: Sender[] | null;