Merge pull request #327 from escapade-mckv/frontend-refactor

chore(dev): frontend refactor - started with thread page

frontend/src/app/(dashboard)/agents/[threadId]/page.tsx

@@ -3,37 +3,19 @@
 import React, {
   useCallback,
   useEffect,
-  useMemo,
   useRef,
   useState,
 } from 'react';
-import Image from 'next/image';
 import { useRouter, useSearchParams } from 'next/navigation';
 import {
-  ArrowDown,
-  CheckCircle,
-  CircleDashed,
   AlertTriangle,
-  Info,
-  File,
-  ChevronRight,
 } from 'lucide-react';
 import {
-  addUserMessage,
+  BillingError,
-  startAgent,
-  stopAgent,
-  getAgentRuns,
-  getMessages,
-  getProject,
-  getThread,
-  updateProject,
   Project,
   Message as BaseApiMessageType,
-  BillingError,
-  checkBillingStatus,
 } from '@/lib/api';
 import { toast } from 'sonner';
-import { Skeleton } from '@/components/ui/skeleton';
 import { ChatInput } from '@/components/thread/chat-input/chat-input';
 import { FileViewerModal } from '@/components/thread/file-viewer-modal';
 import { SiteHeader } from '@/components/thread/thread-site-header';
@@ -58,7 +40,11 @@ import {
 import {
   safeJsonParse,
 } from '@/components/thread/utils';
-
+import { useThreadQuery } from '@/hooks/react-query/threads/use-threads';
+import { useAddUserMessageMutation, useMessagesQuery } from '@/hooks/react-query/threads/use-messages';
+import { useProjectQuery } from '@/hooks/react-query/threads/use-project';
+import { useAgentRunsQuery, useStartAgentMutation, useStopAgentMutation } from '@/hooks/react-query/threads/use-agent-run';
+import { useBillingStatusQuery } from '@/hooks/react-query/threads/use-billing-status';
 
 // Extend the base Message type with the expected database fields
 interface ApiMessageType extends BaseApiMessageType {
@@ -136,6 +122,18 @@ export default function ThreadPage({
   // Add debug mode state - check for debug=true in URL
   const [debugMode, setDebugMode] = useState(false);
 
+  const threadQuery = useThreadQuery(threadId);
+  const messagesQuery = useMessagesQuery(threadId);
+  const projectId = threadQuery.data?.project_id || '';
+  const projectQuery = useProjectQuery(projectId);
+  const agentRunsQuery = useAgentRunsQuery(threadId);
+  const billingStatusQuery = useBillingStatusQuery(); 
+
+  const addUserMessageMutation = useAddUserMessageMutation();
+  const startAgentMutation = useStartAgentMutation();
+  const stopAgentMutation = useStopAgentMutation();
+
+
   const handleProjectRenamed = useCallback((newName: string) => {
     setProjectName(newName);
   }, []);
@@ -356,89 +354,86 @@ export default function ThreadPage({
     messagesEndRef.current?.scrollIntoView({ behavior });
   };
 
+  // Effect to load initial data using React Query
   useEffect(() => {
     let isMounted = true;
 
-    async function loadData() {
+    async function initializeData() {
       if (!initialLoadCompleted.current) setIsLoading(true);
       setError(null);
 
       try {
         if (!threadId) throw new Error('Thread ID is required');
 
-        const threadData = await getThread(threadId).catch((err) => {
+        // Check if we have thread data
-          throw new Error('Failed to load thread data: ' + err.message);
+        if (threadQuery.isError) {
-        });
+          throw new Error('Failed to load thread data: ' + threadQuery.error);
+        }
 
         if (!isMounted) return;
 
-        if (threadData?.project_id) {
+        // Process project data when available
-          const projectData = await getProject(threadData.project_id);
+        if (projectQuery.data) {
-          if (isMounted && projectData) {
+          // Set project data
-            // Set project data
+          setProject(projectQuery.data);
-            setProject(projectData);
 
-            // Make sure sandbox ID is set correctly
+          // Make sure sandbox ID is set correctly
-            if (typeof projectData.sandbox === 'string') {
+          if (typeof projectQuery.data.sandbox === 'string') {
-              setSandboxId(projectData.sandbox);
+            setSandboxId(projectQuery.data.sandbox);
-            } else if (projectData.sandbox?.id) {
+          } else if (projectQuery.data.sandbox?.id) {
-              setSandboxId(projectData.sandbox.id);
+            setSandboxId(projectQuery.data.sandbox.id);
-            }
+          }
 
-            setProjectName(projectData.name || '');
+          setProjectName(projectQuery.data.name || '');
+        }
+
+        // Process messages data when available
+        if (messagesQuery.data && !messagesLoadedRef.current) {
+          // Map API message type to UnifiedMessage type
+          const unifiedMessages = (messagesQuery.data || [])
+            .filter((msg) => msg.type !== 'status')
+            .map((msg: ApiMessageType) => ({
+              message_id: msg.message_id || null,
+              thread_id: msg.thread_id || threadId,
+              type: (msg.type || 'system') as UnifiedMessage['type'],
+              is_llm_message: Boolean(msg.is_llm_message),
+              content: msg.content || '',
+              metadata: msg.metadata || '{}',
+              created_at: msg.created_at || new Date().toISOString(),
+              updated_at: msg.updated_at || new Date().toISOString(),
+            }));
+
+          setMessages(unifiedMessages);
+          console.log('[PAGE] Loaded Messages (excluding status, keeping browser_state):', unifiedMessages.length);
+          messagesLoadedRef.current = true;
+
+          if (!hasInitiallyScrolled.current) {
+            scrollToBottom('auto');
+            hasInitiallyScrolled.current = true;
           }
         }
 
-        if (!messagesLoadedRef.current) {
+        // Check for active agent runs
-          const messagesData = await getMessages(threadId);
+        if (agentRunsQuery.data && !agentRunsCheckedRef.current && isMounted) {
-          if (isMounted) {
+          console.log('[PAGE] Checking for active agent runs...');
-            // Map API message type to UnifiedMessage type
+          agentRunsCheckedRef.current = true;
-            const unifiedMessages = (messagesData || [])
-              .filter((msg) => msg.type !== 'status')
-              .map((msg: ApiMessageType) => ({
-                message_id: msg.message_id || null,
-                thread_id: msg.thread_id || threadId,
-                type: (msg.type || 'system') as UnifiedMessage['type'],
-                is_llm_message: Boolean(msg.is_llm_message),
-                content: msg.content || '',
-                metadata: msg.metadata || '{}',
-                created_at: msg.created_at || new Date().toISOString(),
-                updated_at: msg.updated_at || new Date().toISOString(),
-              }));
 
-            setMessages(unifiedMessages);
+          const activeRun = agentRunsQuery.data.find((run) => run.status === 'running');
-            console.log('[PAGE] Loaded Messages (excluding status, keeping browser_state):', unifiedMessages.length);
+          if (activeRun && isMounted) {
-            messagesLoadedRef.current = true;
+            console.log('[PAGE] Found active run on load:', activeRun.id);
-
+            setAgentRunId(activeRun.id);
-            if (!hasInitiallyScrolled.current) {
+          } else {
-              scrollToBottom('auto');
+            console.log('[PAGE] No active agent runs found');
-              hasInitiallyScrolled.current = true;
-            }
-          }
-        }
-
-        if (!agentRunsCheckedRef.current && isMounted) {
-          try {
-            console.log('[PAGE] Checking for active agent runs...');
-            const agentRuns = await getAgentRuns(threadId);
-            agentRunsCheckedRef.current = true;
-
-            const activeRun = agentRuns.find((run) => run.status === 'running');
-            if (activeRun && isMounted) {
-              console.log('[PAGE] Found active run on load:', activeRun.id);
-              setAgentRunId(activeRun.id);
-            } else {
-              console.log('[PAGE] No active agent runs found');
-              if (isMounted) setAgentStatus('idle');
-            }
-          } catch (err) {
-            console.error('[PAGE] Error checking for active runs:', err);
-            agentRunsCheckedRef.current = true;
             if (isMounted) setAgentStatus('idle');
           }
         }
 
-        initialLoadCompleted.current = true;
+        // Mark initialization as complete when we have the core data
+        if (threadQuery.data && messagesQuery.data && agentRunsQuery.data) {
+          initialLoadCompleted.current = true;
+          setIsLoading(false);
+        }
+
       } catch (err) {
         console.error('Error loading thread data:', err);
         if (isMounted) {
@@ -446,18 +441,27 @@ export default function ThreadPage({
             err instanceof Error ? err.message : 'Failed to load thread';
           setError(errorMessage);
           toast.error(errorMessage);
+          setIsLoading(false);
         }
-      } finally {
-        if (isMounted) setIsLoading(false);
       }
     }
 
-    loadData();
+    if (threadId) {
+      initializeData();
+    }
 
     return () => {
       isMounted = false;
     };
-  }, [threadId]);
+  }, [
+    threadId, 
+    threadQuery.data, 
+    threadQuery.isError, 
+    threadQuery.error,
+    projectQuery.data,
+    messagesQuery.data,
+    agentRunsQuery.data
+  ]);
 
   const handleSubmitMessage = useCallback(
     async (
@@ -483,10 +487,18 @@ export default function ThreadPage({
       scrollToBottom('smooth');
 
       try {
-        const results = await Promise.allSettled([
+        // Use React Query mutations instead of direct API calls
-          addUserMessage(threadId, message),
+        const messagePromise = addUserMessageMutation.mutateAsync({ 
-          startAgent(threadId, options),
+          threadId, 
-        ]);
+          message 
+        });
+        
+        const agentPromise = startAgentMutation.mutateAsync({ 
+          threadId, 
+          options 
+        });
+
+        const results = await Promise.allSettled([messagePromise, agentPromise]);
 
         // Handle failure to add the user message
         if (results[0].status === 'rejected') {
@@ -525,6 +537,11 @@ export default function ThreadPage({
         // If agent started successfully
         const agentResult = results[1].value;
         setAgentRunId(agentResult.agent_run_id);
+        
+        // Refresh queries after successful operations
+        messagesQuery.refetch();
+        agentRunsQuery.refetch();
+        
       } catch (err) {
         // Catch errors from addUserMessage or non-BillingError agent start errors
         console.error('Error sending message or starting agent:', err);
@@ -540,8 +557,8 @@ export default function ThreadPage({
         setIsSending(false);
       }
     },
-    [threadId, project?.account_id],
+    [threadId, project?.account_id, addUserMessageMutation, startAgentMutation, messagesQuery, agentRunsQuery],
-  ); // Ensure project.account_id is a dependency
+  );
 
   const handleStopAgent = useCallback(async () => {
     console.log(`[PAGE] Requesting agent stop via hook.`);
@@ -549,11 +566,18 @@ export default function ThreadPage({
 
     // First stop the streaming and let the hook handle refetching
     await stopStreaming();
-
+    
-    // We don't need to refetch messages here since the hook will do that
+    // Use React Query's stopAgentMutation if we have an agent run ID
-    // The centralizing of refetching in the hook simplifies this logic
+    if (agentRunId) {
-  }, [stopStreaming]);
+      try {
-
+        await stopAgentMutation.mutateAsync(agentRunId);
+        // Refresh agent runs after stopping
+        agentRunsQuery.refetch();
+      } catch (error) {
+        console.error('Error stopping agent:', error);
+      }
+    }
+  }, [stopStreaming, agentRunId, stopAgentMutation, agentRunsQuery]);
 
   useEffect(() => {
     const lastMsg = messages[messages.length - 1];
@@ -561,7 +585,7 @@ export default function ThreadPage({
     if ((isNewUserMessage || agentStatus === 'running') && !userHasScrolled) {
       scrollToBottom('smooth');
     }
-  }, [messages, agentStatus, userHasScrolled, scrollToBottom]);
+  }, [messages, agentStatus, userHasScrolled]);
 
   useEffect(() => {
     if (!latestMessageRef.current || messages.length === 0) return;
@@ -571,7 +595,7 @@ export default function ThreadPage({
     );
     observer.observe(latestMessageRef.current);
     return () => observer.disconnect();
-  }, [messages, streamingTextContent, streamingToolCall, setShowScrollButton]);
+  }, [messages, streamingTextContent, streamingToolCall]);
 
   useEffect(() => {
     console.log(`[PAGE] 🔄 Page AgentStatus: ${agentStatus}, Hook Status: ${streamHookStatus}, Target RunID: ${agentRunId || 'none'}, Hook RunID: ${currentHookRunId || 'none'}`);
@@ -930,57 +954,34 @@ export default function ThreadPage({
     }
   }, [projectName]);
 
-  // Add another useEffect to ensure messages are refreshed when agent status changes to idle
+  // Update messages when they change in the query
   useEffect(() => {
-    if (
+    if (messagesQuery.data && messagesQuery.status === 'success') {
-      agentStatus === 'idle' &&
+      // Only update if we're not in initial loading and the agent isn't running
-      streamHookStatus !== 'streaming' &&
+      if (!isLoading && agentStatus !== 'running' && agentStatus !== 'connecting') {
-      streamHookStatus !== 'connecting'
+        // Map API message type to UnifiedMessage type
-    ) {
+        const unifiedMessages = (messagesQuery.data || [])
-      console.log(
+          .filter((msg) => msg.type !== 'status')
-        '[PAGE] Agent status changed to idle, ensuring messages are up to date',
+          .map((msg: ApiMessageType) => ({
-      );
+            message_id: msg.message_id || null,
-      // Only do this if we're not in the initial loading state
+            thread_id: msg.thread_id || threadId,
-      if (!isLoading && initialLoadCompleted.current) {
+            type: (msg.type || 'system') as UnifiedMessage['type'],
-        // Double-check messages after a short delay to ensure we have latest content
+            is_llm_message: Boolean(msg.is_llm_message),
-        const timer = setTimeout(() => {
+            content: msg.content || '',
-          getMessages(threadId)
+            metadata: msg.metadata || '{}',
-            .then((messagesData) => {
+            created_at: msg.created_at || new Date().toISOString(),
-              if (messagesData) {
+            updated_at: msg.updated_at || new Date().toISOString(),
-                console.log(
+          }));
-                  `[PAGE] Backup refetch completed with ${messagesData.length} messages`,
-                );
-                // Map API message type to UnifiedMessage type
-                const unifiedMessages = (messagesData || [])
-                  .filter((msg) => msg.type !== 'status')
-                  .map((msg: ApiMessageType) => ({
-                    message_id: msg.message_id || null,
-                    thread_id: msg.thread_id || threadId,
-                    type: (msg.type || 'system') as UnifiedMessage['type'],
-                    is_llm_message: Boolean(msg.is_llm_message),
-                    content: msg.content || '',
-                    metadata: msg.metadata || '{}',
-                    created_at: msg.created_at || new Date().toISOString(),
-                    updated_at: msg.updated_at || new Date().toISOString(),
-                  }));
 
-                setMessages(unifiedMessages);
+        setMessages(unifiedMessages);
-                // Reset auto-opened panel to allow tool detection with fresh messages
+        // Reset auto-opened panel to allow tool detection with fresh messages
-                setAutoOpenedPanel(false);
+        setAutoOpenedPanel(false);
-                scrollToBottom('smooth');
+        scrollToBottom('smooth');
-              }
-            })
-            .catch((err) => {
-              console.error('Error in backup message refetch:', err);
-            });
-        }, 1000);
-
-        return () => clearTimeout(timer);
       }
     }
-  }, [agentStatus, threadId, isLoading, streamHookStatus]);
+  }, [messagesQuery.data, messagesQuery.status, isLoading, agentStatus, threadId]);
 
-  // Update the checkBillingStatus function
+  // Check billing status and handle billing limit
   const checkBillingLimits = useCallback(async () => {
     // Skip billing checks in local development mode
     if (isLocalMode()) {
@@ -991,9 +992,11 @@ export default function ThreadPage({
     }
 
     try {
-      const result = await checkBillingStatus();
+      // Use React Query to get billing status
+      await billingStatusQuery.refetch();
+      const result = billingStatusQuery.data;
 
-      if (!result.can_run) {
+      if (result && !result.can_run) {
         setBillingData({
           currentUsage: result.subscription?.minutes_limit || 0,
           limit: result.subscription?.minutes_limit || 0,
@@ -1008,9 +1011,9 @@ export default function ThreadPage({
       console.error('Error checking billing status:', err);
       return false;
     }
-  }, [project?.account_id]);
+  }, [project?.account_id, billingStatusQuery]);
 
-  // Update useEffect to use the renamed function
+  // Check billing when agent status changes
   useEffect(() => {
     const previousStatus = previousAgentStatus.current;
 
@@ -1023,7 +1026,7 @@ export default function ThreadPage({
     previousAgentStatus.current = agentStatus;
   }, [agentStatus, checkBillingLimits]);
 
-  // Update other useEffect to use the renamed function
+  // Check billing on initial load
   useEffect(() => {
     if (project?.account_id && initialLoadCompleted.current) {
       console.log('Checking billing status on page load');
@@ -1031,7 +1034,7 @@ export default function ThreadPage({
     }
   }, [project?.account_id, checkBillingLimits, initialLoadCompleted]);
 
-  // Update the last useEffect to use the renamed function
+  // Check billing after messages loaded
   useEffect(() => {
     if (messagesLoadedRef.current && project?.account_id && !isLoading) {
       console.log('Checking billing status after messages loaded');
@@ -1085,10 +1088,10 @@ export default function ThreadPage({
                   ? 'This thread either does not exist or you do not have access to it.'
                   : error
                 }
-              </p >
+              </p>
-            </div >
+            </div>
-          </div >
+          </div>
-        </div >
+        </div>
         <ToolCallSidePanel
           isOpen={isSidePanelOpen && initialLoadCompleted.current}
           onClose={() => {
@@ -1128,7 +1131,7 @@ export default function ThreadPage({
           onDismiss={() => setShowBillingAlert(false)}
           isOpen={showBillingAlert}
         />
-      </div >
+      </div>
     );
   } else {
     return (

frontend/src/hooks/react-query/threads/keys.ts

@@ -0,0 +1,10 @@
+import { createQueryKeys } from "@/hooks/use-query";
+
+export const threadKeys = createQueryKeys({
+  all: ['threads'] as const,
+  details: (threadId: string) => ['thread', threadId] as const,
+  messages: (threadId: string) => ['thread', threadId, 'messages'] as const,
+  project: (projectId: string) => ['project', projectId] as const,
+  agentRuns: (threadId: string) => ['thread', threadId, 'agent-runs'] as const,
+  billingStatus: ['billing', 'status'] as const,
+});

frontend/src/hooks/react-query/threads/use-agent-run.ts

@@ -0,0 +1,39 @@
+import { createMutationHook, createQueryHook } from "@/hooks/use-query";
+import { threadKeys } from "./keys";
+import { BillingError, getAgentRuns, startAgent, stopAgent } from "@/lib/api";
+
+export const useAgentRunsQuery = (threadId: string) =>
+  createQueryHook(
+    threadKeys.agentRuns(threadId),
+    () => getAgentRuns(threadId),
+    {
+      enabled: !!threadId,
+      retry: 1,
+    }
+  )();
+
+export const useStartAgentMutation = () =>
+  createMutationHook(
+    ({
+      threadId,
+      options,
+    }: {
+      threadId: string;
+      options?: {
+        model_name?: string;
+        enable_thinking?: boolean;
+        reasoning_effort?: string;
+        stream?: boolean;
+      };
+    }) => startAgent(threadId, options),
+    {
+      onError: (error) => {
+        if (!(error instanceof BillingError)) {
+          throw error;
+        }
+      },
+    }
+  )();
+
+export const useStopAgentMutation = () =>
+  createMutationHook((agentRunId: string) => stopAgent(agentRunId))();

frontend/src/hooks/react-query/threads/use-billing-status.ts

@@ -0,0 +1,14 @@
+import { createQueryHook } from "@/hooks/use-query";
+import { threadKeys } from "./keys";
+import { checkBillingStatus } from "@/lib/api";
+
+export const useBillingStatusQuery = (enabled = true) =>
+  createQueryHook(
+    threadKeys.billingStatus,
+    () => checkBillingStatus(),
+    {
+      enabled,
+      retry: 1,
+      staleTime: 1000 * 60 * 5,
+    }
+  )();

frontend/src/hooks/react-query/threads/use-messages.ts

@@ -0,0 +1,24 @@
+import { createMutationHook, createQueryHook } from "@/hooks/use-query";
+import { threadKeys } from "./keys";
+import { addUserMessage, getMessages } from "@/lib/api";
+
+export const useMessagesQuery = (threadId: string) =>
+  createQueryHook(
+    threadKeys.messages(threadId),
+    () => getMessages(threadId),
+    {
+      enabled: !!threadId,
+      retry: 1,
+    }
+  )();
+
+export const useAddUserMessageMutation = () =>
+  createMutationHook(
+    ({
+      threadId,
+      message,
+    }: {
+      threadId: string;
+      message: string;
+    }) => addUserMessage(threadId, message)
+  )();

frontend/src/hooks/react-query/threads/use-project.ts

@@ -0,0 +1,27 @@
+import { createMutationHook, createQueryHook } from "@/hooks/use-query";
+import { threadKeys } from "./keys";
+import { getProject, Project, updateProject } from "@/lib/api";
+
+export const useProjectQuery = (projectId: string | undefined) =>
+  createQueryHook(
+    threadKeys.project(projectId || ""),
+    () =>
+      projectId
+        ? getProject(projectId)
+        : Promise.reject("No project ID"),
+    {
+      enabled: !!projectId,
+      retry: 1,
+    }
+  )();
+
+export const useUpdateProjectMutation = () =>
+  createMutationHook(
+    ({
+      projectId,
+      data,
+    }: {
+      projectId: string;
+      data: Partial<Project>;
+    }) => updateProject(projectId, data)
+  )();

frontend/src/hooks/react-query/threads/use-threads.ts

@@ -0,0 +1,13 @@
+import { createQueryHook } from "@/hooks/use-query";
+import { threadKeys } from "./keys";
+import { getThread } from "@/lib/api";
+
+export const useThreadQuery = (threadId: string) =>
+  createQueryHook(
+    threadKeys.details(threadId),
+    () => getThread(threadId),
+    {
+      enabled: !!threadId,
+      retry: 1,
+    }
+  )();