Refactor GetSummary and AIChat components. Introduce GetState function for fetching brain state. Clean up unused imports and improve error handling in AIChat. Update layout and nav-user components for better functionality.

apps/mail/actions/getSummary.ts

@@ -1,10 +1,6 @@
 'use server';
 import { getAuthenticatedUserId } from '@/app/api/utils';
-import { connection, summary } from '@zero/db/schema';
-import { headers } from 'next/headers';
-import { and, eq } from 'drizzle-orm';
-import { auth } from '@/lib/auth';
-import { db } from '@zero/db';
+import { getActiveConnection } from './utils';
 import axios from 'axios';
 
 export const GetSummary = async (threadId: string) => {
@@ -14,11 +10,7 @@ export const GetSummary = async (threadId: string) => {
       return null;
     }
 
-    const response = await axios.get(process.env.BRAIN_URL + `/brain/thread/summary/${threadId}`, {
-      headers: {
-        // 'Authorization': `Bearer ${}`
-      },
-    });
+    const response = await axios.get(process.env.BRAIN_URL + `/brain/thread/summary/${threadId}`);
 
     return response.data ?? null;
   } catch (error) {
@@ -26,3 +18,21 @@ export const GetSummary = async (threadId: string) => {
     return null;
   }
 };
+
+export const GetState = async () => {
+  try {
+    if (!process.env.BRAIN_URL) {
+      return null;
+    }
+    const connection = await getActiveConnection();
+    if (!connection) {
+      return null;
+    }
+
+    const response = await axios.get(process.env.BRAIN_URL + `/limit/${connection.id}`);
+    return response.data ?? null;
+  } catch (error) {
+    console.error('Error getting summary:', error);
+    return null;
+  }
+};

apps/mail/app/(routes)/layout.tsx

@@ -12,7 +12,7 @@ export default function Layout({ children }: { children: React.ReactNode }) {
         <div className="relative flex max-h-screen w-full overflow-hidden">
           <SWRConfig
             value={{
-              //   provider: typeof window !== 'undefined' ? dexieStorageProvider : undefined,
+              provider: typeof window !== 'undefined' ? dexieStorageProvider : undefined,
               revalidateOnFocus: false,
               revalidateIfStale: false,
               shouldRetryOnError: false,

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

@@ -103,7 +103,7 @@ export async function POST(req: Request) {
             console.error('Error creating label:', error);
             throw new Error('Failed to create label');
           }
-          return { created: 'label.id' };
+          return { success: true };
         },
       },
       addLabelsToThreads: {

apps/mail/components/create/ai-chat.tsx

@@ -6,6 +6,7 @@ import { useSearchValue } from '@/hooks/use-search-value';
 import { useConnections } from '@/hooks/use-connections';
 import { useRef, useCallback, useEffect } from 'react';
 import { Markdown } from '@react-email/components';
+import { TextShimmer } from '../ui/text-shimmer';
 import { useThread } from '@/hooks/use-threads';
 import { useSession } from '@/lib/auth-client';
 import { cn, getEmailLogo } from '@/lib/utils';
@@ -102,26 +103,16 @@ const RenderThreads = ({
   return <div className="flex flex-col gap-2">{threads.map(renderThread)}</div>;
 };
 
-export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
-  const [value, setValue] = useState('');
+export function AIChat() {
   const [showVoiceChat, setShowVoiceChat] = useState(false);
-  const [expandedResults, setExpandedResults] = useState<Set<string>>(new Set());
-  const [searchValue, setSearchValue] = useSearchValue();
   const messagesEndRef = useRef<HTMLDivElement>(null);
   const messagesContainerRef = useRef<HTMLDivElement>(null);
-  const fileInputRef = useRef<HTMLInputElement>(null);
-  const pathname = usePathname();
-  const router = useRouter();
-  const searchParams = useSearchParams();
-  const { data: session } = useSession();
-  const { data: connections } = useConnections();
 
   const { messages, input, setInput, error, handleSubmit, status } = useChat({
     api: '/api/chat',
     maxSteps: 5,
   });
 
-  // Scroll to bottom function
   const scrollToBottom = useCallback(() => {
     if (messagesEndRef.current) {
       messagesEndRef.current.scrollIntoView({ behavior: 'smooth' });
@@ -134,134 +125,10 @@ export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
     // if (onMessagesChange) {
     //   onMessagesChange(messages);
     // }
-  }, [messages, onMessagesChange, scrollToBottom]);
-
-  useEffect(() => {
-    if (onReset) {
-      onReset();
-    }
-  }, [onReset]);
-
-  //   const handleSendMessage = async () => {
-  //     if (!value.trim() || isLoading) return;
-
-  //     const userMessage: Message = {
-  //       id: generateId(),
-  //       role: 'user',
-  //       content: value.trim(),
-  //       timestamp: new Date(),
-  //     };
-
-  //     setMessages((prev) => [...prev, userMessage]);
-  //     setValue('');
-  //     setIsLoading(true);
-
-  //     try {
-  //       if (!response.ok) {
-  //         throw new Error('Failed to get response');
-  //       }
-
-  //       const data = await response.json();
-
-  //       // Update the search value
-  //       setSearchValue({
-  //         value: data.searchQuery,
-  //         highlight: value.trim(),
-  //         isLoading: false,
-  //         isAISearching: false,
-  //         folder: searchValue.folder,
-  //       });
-
-  //       // Add assistant message with search results
-  //       const assistantMessage: Message = {
-  //         id: generateId(),
-  //         role: 'assistant',
-  //         content: data.content,
-  //         timestamp: new Date(),
-  //         type: 'search',
-  //         searchContent: {
-  //           searchDisplay: data.searchDisplay,
-  //           results: data.results,
-  //         },
-  //       };
-
-  //       setMessages((prev) => [...prev, assistantMessage]);
-  //     } catch (error) {
-  //       console.error('Error:', error);
-  //       toast.error('Failed to generate response. Please try again.');
-  //     } finally {
-  //       setIsLoading(false);
-  //     }
-  //   };
-
-  const handleAcceptSuggestion = (emailContent: { subject?: string; content: string }) => {
-    if (!editor) {
-      toast.error('Editor not found');
-      return;
-    }
-
-    try {
-      // Format the content to preserve line breaks
-      const formattedContent = emailContent.content
-        .split('\n')
-        .map((line) => `<p>${line}</p>`)
-        .join('');
-
-      // Set the content in the editor
-      editor.commands.setContent(formattedContent);
-
-      // Find the create-email component and update its content
-      const createEmailElement = document.querySelector('[data-create-email]');
-      if (createEmailElement) {
-        const handler = (createEmailElement as any).onContentGenerated;
-        if (handler && typeof handler === 'function') {
-          handler({ content: emailContent.content, subject: emailContent.subject });
-        }
-      }
-
-      toast.success('Email content applied successfully');
-    } catch (error) {
-      console.error('Error applying suggestion:', error);
-      toast.error('Failed to apply email content');
-    }
-  };
-
-  const handleRejectSuggestion = (messageId: string) => {
-    toast.info('Email suggestion rejected');
-  };
-
-  const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-    if (e.key === 'Enter' && (e.metaKey || e.ctrlKey)) {
-      e.preventDefault();
-      handleSubmit();
-    }
-  };
-
-  const generateId = () => nanoid();
-
-  const toggleExpandResults = (messageId: string) => {
-    setExpandedResults((prev) => {
-      const newSet = new Set(prev);
-      if (newSet.has(messageId)) {
-        newSet.delete(messageId);
-      } else {
-        newSet.add(messageId);
-      }
-      return newSet;
-    });
-  };
-
-  const sanitizeSnippet = (snippet: string) => {
-    return snippet
-      .replace(/<\/?[^>]+(>|$)/g, '') // Remove HTML tags
-      .replace(/&lt;/g, '<')
-      .replace(/&gt;/g, '>')
-      .replace(/&amp;/g, '&');
-  };
+  }, [messages, messagesEndRef]);
 
   return (
     <div className="flex h-full flex-col">
-      {/* Messages container */}
       <div className="flex-1 overflow-y-auto" ref={messagesContainerRef}>
         <div className="min-h-full space-y-4 px-4 py-4">
           {!messages.length ? (
@@ -325,10 +192,6 @@ export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
                     : 'overflow-wrap-anywhere mr-auto break-words bg-[#f0f0f0] p-2 dark:bg-[#313131]', // Assistant messages aligned to left
                 )}
               >
-                {/* <div className="prose dark:prose-invert overflow-wrap-anywhere break-words text-sm font-medium">
-                  {message.content}
-                </div> */}
-
                 {message.parts.map((part) => {
                   if (part.type === 'text') {
                     return <Markdown>{part.text}</Markdown>;
@@ -341,108 +204,22 @@ export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
                       ) : null)
                     );
                   }
-                  //   if (part.type === 'source') {
-                  //     return <p>Source: {part.source.title}</p>;
-                  //   }
-                  //   if (part.type === 'step-start') {
-                  //     return <ArrowDownCircle className="mx-auto h-4 w-4" />;
-                  //   }
-                  //   return <p>{part.type}</p>;
                 })}
-
-                {/* {message.type === 'search' &&
-                message.searchContent &&
-                message.searchContent.results.length > 0 && (
-                  <div className="bg-muted space-y-4 rounded-lg px-4 pt-3">
-                    {(expandedResults.has(message.id)
-                      ? message.searchContent.results
-                      : message.searchContent.results.slice(0, 5)
-                    ).map((result: any, i: number) => (
-                      <div key={i} className="border-t pt-4 first:border-t-0 first:pt-0">
-                        <div className="font-medium">
-                          <p className="max-w-sm truncate text-sm">
-                            {result.subject.toLowerCase().includes('meeting') ? (
-                              <span className="text-blue-500">📅 {result.subject}</span>
-                            ) : (
-                              result.subject || 'No subject'
-                            )}
-                          </p>
-                          <span className="text-muted-foreground text-sm">
-                            from {result.from || 'Unknown sender'}
-                          </span>
-                        </div>
-                        <div className="text-muted-foreground mt-1 line-clamp-2 text-xs">
-                          {sanitizeSnippet(result.snippet)}
-                        </div>
-                        <div className="text-muted-foreground mt-1 text-sm">
-                          <button
-                            onClick={() => handleThreadClick(result.id)}
-                            className="cursor-pointer border-none bg-transparent p-0 text-blue-500 hover:underline"
-                          >
-                            Open email
-                          </button>
-                        </div>
-                      </div>
-                    ))}
-                    {message.searchContent.results.length > 5 && (
-                      <Button
-                        variant="ghost"
-                        className="text-muted-foreground hover:text-foreground w-full"
-                        onClick={() => toggleExpandResults(message.id)}
-                      >
-                        {expandedResults.has(message.id)
-                          ? `Show less (${message.searchContent.results.length - 5} fewer results)`
-                          : `Show more (${message.searchContent.results.length - 5} more results)`}
-                      </Button>
-                    )}
-                  </div>
-                )}
-
-              {message.type === 'email' && message.emailContent && (
-                <div className="bg-background mt-4 rounded border p-4 font-mono text-sm">
-                  {message.emailContent.subject && (
-                    <div className="mb-2 text-blue-500">
-                      Subject: {message.emailContent.subject}
-                    </div>
-                  )}
-                  <div className="whitespace-pre-wrap">{message.emailContent.content}</div>
-                  <div className="mt-4 flex gap-2">
-                    <Button
-                      variant="outline"
-                      size="sm"
-                      className="h-8 border-green-500/20 hover:bg-green-500/10 hover:text-green-500"
-                      onClick={() => handleAcceptSuggestion(message.emailContent!)}
-                    >
-                      <CheckIcon className="mr-1 h-4 w-4" />
-                      Accept
-                    </Button>
-                    <Button
-                      variant="outline"
-                      size="sm"
-                      className="border-destructive/20 hover:bg-destructive/10 hover:text-destructive h-8"
-                      onClick={() => handleRejectSuggestion(message.id)}
-                    >
-                      <XIcon className="mr-1 h-4 w-4" />
-                      Reject
-                    </Button>
-                  </div>
-                </div>
-              )} */}
               </div>
             ))
           )}
-          {/* Invisible element to scroll to */}
           <div ref={messagesEndRef} />
 
-          {JSON.stringify(error)}
-          {/* Loading indicator */}
           {status === 'submitted' && (
             <div className="flex flex-col gap-2 rounded-lg">
               <div className="flex items-center gap-2">
-                <span className="text-muted-foreground text-sm">zero is thinking...</span>
+                <TextShimmer className="text-muted-foreground text-sm">
+                  zero is thinking...
+                </TextShimmer>
               </div>
             </div>
           )}
+          {status === 'error' && <div className="text-red-500">Error, please try again later</div>}
         </div>
       </div>
 
@@ -469,7 +246,7 @@ export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
                   form="ai-chat-form"
                   type="submit"
                   className="border-border/50 inline-flex h-7 cursor-pointer items-center justify-center gap-1.5 overflow-hidden rounded-md border bg-white pl-1.5 pr-1 dark:bg-[#262626]"
-                  disabled={!input.trim()}
+                  disabled={!input.trim() || status !== 'ready'}
                 >
                   <div className="flex items-center justify-center gap-2.5 pl-0.5">
                     <div className="justify-start text-center text-sm leading-none text-black dark:text-white">

apps/mail/components/ui/nav-user.tsx

@@ -31,6 +31,7 @@ import { signOut, useSession } from '@/lib/auth-client';
 import { AddConnectionDialog } from '../connection/add';
 import { putConnection } from '@/actions/connections';
 import { useSidebar } from '@/components/ui/sidebar';
+import { useBrainState } from '@/hooks/use-summary';
 import { dexieStorageProvider } from '@/lib/idb';
 import { EnableBrain } from '@/actions/brain';
 import { useRouter } from 'next/navigation';
@@ -101,6 +102,8 @@ export function NavUser() {
     );
   };
 
+  const { data: brainState } = useBrainState();
+
   const handleThemeToggle = () => {
     setTheme(theme === 'dark' ? 'light' : 'dark');
   };
@@ -266,14 +269,6 @@ export function NavUser() {
                       Terms
                     </a>
                   </div>
-                  <DropdownMenuSeparator className="mt-1" />
-                  <p className="text-muted-foreground px-2 py-1 text-[11px] font-medium">Debug</p>
-                  <DropdownMenuItem onClick={handleClearCache}>
-                    <div className="flex items-center gap-2">
-                      <HelpCircle size={16} className="opacity-60" />
-                      <p className="text-[13px] opacity-60">Clear Local Cache</p>
-                    </div>
-                  </DropdownMenuItem>
                 </>
               </DropdownMenuContent>
             </DropdownMenu>
@@ -394,12 +389,14 @@ export function NavUser() {
                       <p className="text-[13px] opacity-60">Clear Local Cache</p>
                     </div>
                   </DropdownMenuItem>
-                  {/* <DropdownMenuItem onClick={handleEnableBrain}>
-                    <div className="flex items-center gap-2">
-                      <BrainIcon size={16} className="opacity-60" />
-                      <p className="text-[13px] opacity-60">Enable Brain Activity</p>
-                    </div>
-                  </DropdownMenuItem> */}
+                  {!brainState?.enabled ? (
+                    <DropdownMenuItem onClick={handleEnableBrain}>
+                      <div className="flex items-center gap-2">
+                        <BrainIcon size={16} className="opacity-60" />
+                        <p className="text-[13px] opacity-60">Enable Brain Activity</p>
+                      </div>
+                    </DropdownMenuItem>
+                  ) : null}
                 </DropdownMenuContent>
               </DropdownMenu>
             </div>

apps/mail/hooks/use-summary.ts

@@ -1,11 +1,12 @@
 'use client';
-import { GetSummary } from '@/actions/getSummary';
-import { string } from 'zod';
+import { GetState, GetSummary } from '@/actions/getSummary';
+import { useSession } from '@/lib/auth-client';
 import useSWR from 'swr';
 
 export const useSummary = (threadId: string | null) => {
+  const { data: session } = useSession();
   const { data, isLoading } = useSWR<{ short: string; long: string } | null>(
-    threadId ? `ai:summary:${threadId}` : null,
+    session && threadId ? `ai:summary:${threadId}` : null,
     async () => {
       if (!threadId) return null;
       return await GetSummary(threadId);
@@ -14,3 +15,15 @@ export const useSummary = (threadId: string | null) => {
 
   return { data, isLoading };
 };
+
+export const useBrainState = () => {
+  const { data: session } = useSession();
+  const { data, isLoading } = useSWR<{ enabled: boolean } | null>(
+    session ? `brain:state:${session?.connectionId}` : null,
+    async () => {
+      return await GetState();
+    },
+  );
+
+  return { data, isLoading };
+};