wip

backend/core/agent_crud.py

@@ -635,7 +635,6 @@ async def create_agent(
         logger.error(f"Error creating agent for user {user_id}: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Failed to create agent: {str(e)}")
 
-
 @router.post("/agents/generate-icon", response_model=AgentIconGenerationResponse)
 async def generate_agent_icon(
     request: AgentIconGenerationRequest,

frontend/src/components/thread/content/ThreadContent.tsx

@@ -1,4 +1,4 @@
-import React, { useRef, useState, useCallback, useEffect } from 'react';
+import React, { useRef, useState, useCallback, useEffect, useMemo } from 'react';
 import { CircleDashed, CheckCircle, AlertTriangle } from 'lucide-react';
 import { UnifiedMessage, ParsedContent, ParsedMetadata } from '@/components/thread/types';
 import { FileAttachmentGrid } from '@/components/thread/file-attachment';
@@ -19,6 +19,7 @@ import { ShowToolStream } from './ShowToolStream';
 import { ComposioUrlDetector } from './composio-url-detector';
 import { StreamingText } from './StreamingText';
 import { HIDE_STREAMING_XML_TAGS } from '@/components/thread/utils';
+import { useAgentsFromCache } from '@/hooks/react-query/agents/use-agents';
 
 
 // Helper function to render all attachments as standalone messages
@@ -425,6 +426,20 @@ export const ThreadContent: React.FC<ThreadContentProps> = ({
     const [shouldJustifyToTop, setShouldJustifyToTop] = useState(false);
     const { session } = useAuth();
 
+    // Collect unique agent IDs from messages to prefetch from cache
+    const agentIds = useMemo(() => {
+        const ids = new Set<string>();
+        messages.forEach(msg => {
+            if (msg.agent_id) {
+                ids.add(msg.agent_id);
+            }
+        });
+        return Array.from(ids);
+    }, [messages]);
+
+    // Get agents from cache (doesn't fetch, just reads from existing cache)
+    const agentsMap = useAgentsFromCache(agentIds);
+
     // React Query file preloader
     const { preloadFiles } = useFilePreloader();
 
@@ -815,21 +830,22 @@ export const ThreadContent: React.FC<ThreadContentProps> = ({
                                         // Get agent_id from the first assistant message in this group
                                         const firstAssistantMsg = group.messages.find(m => m.type === 'assistant');
                                         const groupAgentId = firstAssistantMsg?.agent_id;
+                                        const groupAgent = groupAgentId ? agentsMap.get(groupAgentId) : undefined;
 
                                         return (
                                             <div key={group.key} ref={groupIndex === groupedMessages.length - 1 ? latestMessageRef : null}>
                                                 <div className="flex flex-col gap-2">
                                                     <div className="flex items-center">
                                                         <div className="rounded-md flex items-center justify-center relative">
-                                                            {groupAgentId ? (
-                                                                <AgentAvatar agentId={groupAgentId} size={24} className="h-6 w-6" />
+                                                            {groupAgent || groupAgentId ? (
+                                                                <AgentAvatar agent={groupAgent} agentId={groupAgentId} size={24} className="h-6 w-6" />
                                                             ) : (
                                                                 getAgentInfo().avatar
                                                             )}
                                                         </div>
                                                         <p className='ml-2 text-sm text-muted-foreground'>
-                                                            {groupAgentId ? (
-                                                                <AgentName agentId={groupAgentId} fallback={getAgentInfo().name} />
+                                                            {groupAgent || groupAgentId ? (
+                                                                <AgentName agent={groupAgent} agentId={groupAgentId} fallback={getAgentInfo().name} />
                                                             ) : (
                                                                 getAgentInfo().name
                                                             )}

frontend/src/components/thread/content/agent-avatar.tsx

@@ -1,13 +1,17 @@
 'use client';
 
 import React from 'react';
-import { useAgent } from '@/hooks/react-query/agents/use-agents';
+import { useAgentFromCache } from '@/hooks/react-query/agents/use-agents';
 import { KortixLogo } from '@/components/sidebar/kortix-logo';
 import { DynamicIcon } from 'lucide-react/dynamic';
 import { cn } from '@/lib/utils';
+import type { Agent } from '@/hooks/react-query/agents/utils';
 
 interface AgentAvatarProps {
-  // For fetching agent by ID
+  // For passing agent data directly (preferred - no fetch)
+  agent?: Agent;
+  
+  // For fetching agent by ID (will use cache if available)
   agentId?: string;
   fallbackName?: string;
   
@@ -24,7 +28,8 @@ interface AgentAvatarProps {
 }
 
 export const AgentAvatar: React.FC<AgentAvatarProps> = ({ 
-  // Agent fetch props
+  // Agent data props
+  agent: propAgent,
   agentId, 
   fallbackName = "Suna",
   
@@ -39,7 +44,9 @@ export const AgentAvatar: React.FC<AgentAvatarProps> = ({
   size = 16, 
   className = ""
 }) => {
-  const { data: agent, isLoading } = useAgent(agentId || '');
+  // Try to get agent from cache if agentId is provided and agent prop is not
+  const cachedAgent = useAgentFromCache(!propAgent && agentId ? agentId : undefined);
+  const agent = propAgent || cachedAgent;
 
   // Determine values from props or agent data
   const iconName = propIconName ?? agent?.icon_name;
@@ -54,8 +61,8 @@ export const AgentAvatar: React.FC<AgentAvatarProps> = ({
     borderRadius: `${Math.min(size * 0.25, 16)}px` // 25% of size, max 16px
   };
 
-  // Show skeleton for loading state or when no data is available
-  if ((isLoading && agentId) || (!agent && !agentId && !propIconName && !propIsSunaDefault)) {
+  // Show skeleton when no data is available
+  if (!agent && !propIconName && !propIsSunaDefault && agentId) {
     return (
       <div 
         className={cn("bg-muted animate-pulse border", className)}
@@ -120,19 +127,18 @@ export const AgentAvatar: React.FC<AgentAvatarProps> = ({
 };
 
 interface AgentNameProps {
+  agent?: Agent;
   agentId?: string;
   fallback?: string;
 }
 
 export const AgentName: React.FC<AgentNameProps> = ({ 
+  agent: propAgent,
   agentId, 
   fallback = "Suna" 
 }) => {
-  const { data: agent, isLoading } = useAgent(agentId || '');
-
-  if (isLoading && agentId) {
-    return <span className="text-muted-foreground">Loading...</span>;
-  }
+  const cachedAgent = useAgentFromCache(!propAgent && agentId ? agentId : undefined);
+  const agent = propAgent || cachedAgent;
 
   return <span>{agent?.name || fallback}</span>;
 };

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

@@ -3,7 +3,7 @@ import { useQueryClient, type UseQueryOptions } from '@tanstack/react-query';
 import { toast } from 'sonner';
 import { agentKeys } from './keys';
 import { Agent, AgentUpdateRequest, AgentsParams, createAgent, deleteAgent, getAgent, getAgents, getThreadAgent, updateAgent } from './utils';
-import { useRef, useCallback, useState } from 'react';
+import { useRef, useCallback, useState, useMemo } from 'react';
 import { useRouter } from 'next/navigation';
 import { DEFAULT_AGENTPRESS_TOOLS } from '@/components/agents/tools';
 
@@ -213,4 +213,82 @@ export const useThreadAgent = (threadId: string) => {
       gcTime: 10 * 60 * 1000,
     }
   )();
+};
+
+/**
+ * Hook to get an agent from the cache without fetching.
+ * This checks all cached agent list queries to find the agent.
+ * Returns undefined if not found in cache.
+ */
+export const useAgentFromCache = (agentId: string | undefined): Agent | undefined => {
+  const queryClient = useQueryClient();
+  
+  return useMemo(() => {
+    if (!agentId) return undefined;
+
+    // First check if we have it in the detail cache
+    const cachedAgent = queryClient.getQueryData<Agent>(agentKeys.detail(agentId));
+    if (cachedAgent) return cachedAgent;
+
+    // Otherwise, search through all agent list caches
+    const allAgentLists = queryClient.getQueriesData<{ agents: Agent[] }>({ 
+      queryKey: agentKeys.lists() 
+    });
+
+    for (const [_, data] of allAgentLists) {
+      if (data?.agents) {
+        const found = data.agents.find(agent => agent.agent_id === agentId);
+        if (found) return found;
+      }
+    }
+
+    return undefined;
+  }, [agentId, queryClient]);
+};
+
+/**
+ * Hook to get multiple agents from cache by IDs.
+ * Returns a map of agentId -> Agent for quick lookup.
+ */
+export const useAgentsFromCache = (agentIds: string[]): Map<string, Agent> => {
+  const queryClient = useQueryClient();
+  
+  return useMemo(() => {
+    const agentsMap = new Map<string, Agent>();
+    
+    if (!agentIds || agentIds.length === 0) return agentsMap;
+
+    // Get all cached agent list queries
+    const allAgentLists = queryClient.getQueriesData<{ agents: Agent[] }>({ 
+      queryKey: agentKeys.lists() 
+    });
+
+    // Build a map of all cached agents
+    const allCachedAgents = new Map<string, Agent>();
+    for (const [_, data] of allAgentLists) {
+      if (data?.agents) {
+        data.agents.forEach(agent => {
+          allCachedAgents.set(agent.agent_id, agent);
+        });
+      }
+    }
+
+    // Also check individual agent caches
+    for (const agentId of agentIds) {
+      const cachedAgent = queryClient.getQueryData<Agent>(agentKeys.detail(agentId));
+      if (cachedAgent) {
+        allCachedAgents.set(agentId, cachedAgent);
+      }
+    }
+
+    // Return only the requested agents
+    for (const agentId of agentIds) {
+      const agent = allCachedAgents.get(agentId);
+      if (agent) {
+        agentsMap.set(agentId, agent);
+      }
+    }
+
+    return agentsMap;
+  }, [agentIds, queryClient]);
 };