diff --git a/agentex-ui/app/agentex-ui-root.tsx b/agentex-ui/app/agentex-ui-root.tsx index 0beaac3a..64c8e735 100644 --- a/agentex-ui/app/agentex-ui-root.tsx +++ b/agentex-ui/app/agentex-ui-root.tsx @@ -7,7 +7,6 @@ import { ToastContainer } from 'react-toastify'; import { TaskSidebar } from '@/components/agentex/task-sidebar'; import { TracesSidebar } from '@/components/agentex/traces-sidebar'; import { AgentexProvider } from '@/components/providers'; -import { QueryProvider } from '@/components/providers/query-provider'; import { useLocalStorageState } from '@/hooks/use-local-storage-state'; import { SearchParamKey, @@ -65,7 +64,7 @@ export function AgentexUIRoot({ }, [handleSelectTask]); return ( - + <> - + ); } diff --git a/agentex-ui/components/agentex/task-sidebar.tsx b/agentex-ui/components/agentex/task-sidebar.tsx index 67c7cecd..0a3b80c4 100644 --- a/agentex-ui/components/agentex/task-sidebar.tsx +++ b/agentex-ui/components/agentex/task-sidebar.tsx @@ -144,12 +144,10 @@ export function TaskSidebar() { const [isCollapsed, setIsCollapsed] = useState(false); const scrollContainerRef = useRef(null); - // Flatten all pages into a single array of tasks const tasks = useMemo(() => { - return data?.pages.flatMap(page => page) ?? []; + return data?.pages?.flatMap(page => page) ?? []; }, [data]); - // Scroll detection for infinite loading useEffect(() => { const scrollContainer = scrollContainerRef.current; if (!scrollContainer) return; diff --git a/agentex-ui/components/providers/query-provider.tsx b/agentex-ui/components/providers/query-provider.tsx index 6dfe7581..8a1efab6 100644 --- a/agentex-ui/components/providers/query-provider.tsx +++ b/agentex-ui/components/providers/query-provider.tsx @@ -10,8 +10,7 @@ export function QueryProvider({ children }: { children: React.ReactNode }) { new QueryClient({ defaultOptions: { queries: { - staleTime: 60 * 1000, // 1 minute - refetchOnWindowFocus: false, + staleTime: Infinity, }, }, }) diff --git a/agentex-ui/components/providers/task-provider.tsx b/agentex-ui/components/providers/task-provider.tsx index de359621..08a0add5 100644 --- a/agentex-ui/components/providers/task-provider.tsx +++ b/agentex-ui/components/providers/task-provider.tsx @@ -23,7 +23,8 @@ export function TaskProvider({ useTaskSubscription({ agentexClient, taskId, - enabled: !!taskId && agent?.acp_type === 'agentic', + agentName: agentName || '', + enabled: !!taskId && (agent?.acp_type === 'agentic' || !agentName), }); return <>{children}; diff --git a/agentex-ui/hooks/use-agents.ts b/agentex-ui/hooks/use-agents.ts index 2c697577..062bd1df 100644 --- a/agentex-ui/hooks/use-agents.ts +++ b/agentex-ui/hooks/use-agents.ts @@ -3,24 +3,16 @@ import { useQuery } from '@tanstack/react-query'; import type AgentexSDK from 'agentex'; import type { Agent } from 'agentex/resources'; -/** - * Query key factory for agents - */ export const agentsKeys = { all: ['agents'] as const, }; -/** - * Fetches the list of agents - * This replaces useFetchAgents and can be renamed to that if needed - */ export function useAgents(agentexClient: AgentexSDK) { return useQuery({ queryKey: agentsKeys.all, queryFn: async (): Promise => { return agentexClient.agents.list(); }, - staleTime: 5 * 60 * 1000, // 5 minutes - agents don't change often refetchOnWindowFocus: false, }); } diff --git a/agentex-ui/hooks/use-create-task.ts b/agentex-ui/hooks/use-create-task.ts index bdc1b0a7..6e767990 100644 --- a/agentex-ui/hooks/use-create-task.ts +++ b/agentex-ui/hooks/use-create-task.ts @@ -1,23 +1,74 @@ 'use client'; -import { useMutation, useQueryClient } from '@tanstack/react-query'; +import { + InfiniteData, + useMutation, + useQueryClient, +} from '@tanstack/react-query'; import { agentRPCNonStreaming } from 'agentex/lib'; import { toast } from '@/components/agentex/toast'; - -import { tasksKeys } from './use-tasks'; +import { tasksKeys } from '@/hooks/use-tasks'; import type AgentexSDK from 'agentex'; -import type { Task } from 'agentex/resources'; +import type { + Agent, + Task, + TaskListResponse, + TaskRetrieveResponse, +} from 'agentex/resources'; + +export function updateTaskInInfiniteQuery( + task: Task, + agentName: string, + data: InfiniteData | undefined +): InfiniteData | undefined { + if (!data) return undefined; + + if (data.pages.some(page => page.some(t => t.id === task.id))) { + return { + pages: data.pages.map(page => + page.map(t => + t.id === task.id ? { ...task, agents: t.agents || null } : t + ) + ), + pageParams: data.pageParams, + }; + } + + // Create a dummy agent to add to the task + const agent: Agent = { + id: '1', + name: agentName, + description: '', + acp_type: 'agentic', + created_at: new Date().toISOString(), + updated_at: new Date().toISOString(), + }; + const taskWithAgentName: TaskListResponse.TaskListResponseItem = { + ...task, + agents: [agent], + }; + + // Add the new task to the top of the first page + const newPages = [...data.pages]; + if (newPages.length > 0 && newPages[0]) { + newPages[0] = [taskWithAgentName, ...newPages[0]]; + } else { + newPages[0] = [taskWithAgentName]; + } + + return { + pages: newPages, + pageParams: data.pageParams, + }; +} type CreateTaskParams = { agentName: string; params?: Record; }; -/** - * Creates a new task for a given agent - */ export function useCreateTask({ agentexClient, }: { @@ -45,11 +96,19 @@ export function useCreateTask({ return response.result; }, - onSuccess: (_, variables) => { - queryClient.invalidateQueries({ queryKey: tasksKeys.all }); - queryClient.invalidateQueries({ - queryKey: tasksKeys.byAgentName(variables.agentName), - }); + onSuccess: (task, variables) => { + queryClient.setQueryData( + tasksKeys.individualById(task.id), + task + ); + queryClient.setQueryData>( + tasksKeys.all, + data => updateTaskInInfiniteQuery(task, variables.agentName, data) + ); + queryClient.setQueryData>( + tasksKeys.byAgentName(variables.agentName), + data => updateTaskInInfiniteQuery(task, variables.agentName, data) + ); }, onError: error => { toast.error({ diff --git a/agentex-ui/hooks/use-spans.ts b/agentex-ui/hooks/use-spans.ts index c118167d..ca77abbe 100644 --- a/agentex-ui/hooks/use-spans.ts +++ b/agentex-ui/hooks/use-spans.ts @@ -6,6 +6,12 @@ import { useAgentexClient } from '@/components/providers'; import type { Span } from 'agentex/resources'; +export const spansKeys = { + all: ['spans'] as const, + byTraceId: (traceId: string | null) => + traceId ? ([...spansKeys.all, traceId] as const) : spansKeys.all, +}; + type UseSpansState = { spans: Span[]; isLoading: boolean; @@ -16,7 +22,7 @@ export function useSpans(traceId: string | null): UseSpansState { const { agentexClient } = useAgentexClient(); const { data, isLoading, error } = useQuery({ - queryKey: ['spans', traceId], + queryKey: spansKeys.byTraceId(traceId), queryFn: async ({ signal }) => { if (!traceId) { return []; diff --git a/agentex-ui/hooks/use-task-messages.ts b/agentex-ui/hooks/use-task-messages.ts index 4bb8f829..977eae85 100644 --- a/agentex-ui/hooks/use-task-messages.ts +++ b/agentex-ui/hooks/use-task-messages.ts @@ -50,7 +50,6 @@ export function useTaskMessages({ }; }, enabled: !!taskId, - staleTime: Infinity, refetchOnMount: false, refetchOnWindowFocus: false, }); diff --git a/agentex-ui/hooks/use-task-subscription.ts b/agentex-ui/hooks/use-task-subscription.ts index 90b144bb..e6e43904 100644 --- a/agentex-ui/hooks/use-task-subscription.ts +++ b/agentex-ui/hooks/use-task-subscription.ts @@ -1,22 +1,25 @@ import { useEffect } from 'react'; -import { useQueryClient } from '@tanstack/react-query'; +import { InfiniteData, useQueryClient } from '@tanstack/react-query'; import { subscribeTaskState } from 'agentex/lib'; -import { taskMessagesKeys } from './use-task-messages'; -import { tasksKeys } from './use-tasks'; +import { updateTaskInInfiniteQuery } from '@/hooks/use-create-task'; +import { taskMessagesKeys } from '@/hooks/use-task-messages'; +import type { TaskMessagesData } from '@/hooks/use-task-messages'; +import { tasksKeys } from '@/hooks/use-tasks'; -import type { TaskMessagesData } from './use-task-messages'; import type AgentexSDK from 'agentex'; -import type { Task } from 'agentex/resources'; +import type { TaskListResponse, TaskRetrieveResponse } from 'agentex/resources'; export function useTaskSubscription({ agentexClient, taskId, + agentName, enabled = true, }: { agentexClient: AgentexSDK; taskId: string; + agentName: string; enabled?: boolean; }) { const queryClient = useQueryClient(); @@ -58,11 +61,18 @@ export function useTaskSubscription({ }, onAgentsChange() {}, onTaskChange(task) { - queryClient.setQueryData(tasksKeys.byId(taskId), task); - queryClient.setQueryData(tasksKeys.all, old => { - if (!old) return [task]; - return old.map(t => (t.id === task.id ? task : t)); - }); + queryClient.setQueryData( + tasksKeys.individualById(taskId), + task + ); + queryClient.setQueryData>( + tasksKeys.all, + data => updateTaskInInfiniteQuery(task, agentName, data) + ); + queryClient.setQueryData>( + tasksKeys.byAgentName(agentName), + data => updateTaskInInfiniteQuery(task, agentName, data) + ); }, onStreamStatusChange() {}, onError() { @@ -82,5 +92,5 @@ export function useTaskSubscription({ return () => { abortController.abort(); }; - }, [agentexClient, taskId, enabled, queryClient]); + }, [agentexClient, taskId, enabled, queryClient, agentName]); } diff --git a/agentex-ui/hooks/use-tasks.ts b/agentex-ui/hooks/use-tasks.ts index db293e1b..0f626924 100644 --- a/agentex-ui/hooks/use-tasks.ts +++ b/agentex-ui/hooks/use-tasks.ts @@ -7,21 +7,15 @@ import type { TaskRetrieveResponse, } from 'agentex/resources'; -/** - * Query key factory for tasks - */ export const tasksKeys = { all: ['tasks'] as const, + individual: ['task'] as const, byAgentName: (agentName: string | undefined) => - agentName - ? ([...tasksKeys.all, 'agent', agentName] as const) - : tasksKeys.all, - byId: (taskId: string) => [...tasksKeys.all, taskId] as const, + agentName ? ([...tasksKeys.all, agentName] as const) : tasksKeys.all, + individualById: (taskId: string) => + [...tasksKeys.individual, taskId] as const, }; -/** - * Fetches a single task by ID - */ export function useTask({ agentexClient, taskId, @@ -30,20 +24,16 @@ export function useTask({ taskId: string; }) { return useQuery({ - queryKey: tasksKeys.byId(taskId), + queryKey: tasksKeys.individualById(taskId), queryFn: async (): Promise => { return agentexClient.tasks.retrieve(taskId, { relationships: ['agents'], }); }, enabled: !!taskId, - staleTime: 30 * 1000, }); } -/** - * useQuery hook for infinite scrolling tasks - */ export function useInfiniteTasks( agentexClient: AgentexSDK, options?: { agentName?: string; limit?: number } @@ -62,13 +52,11 @@ export function useInfiniteTasks( return agentexClient.tasks.list(params); }, getNextPageParam: (lastPage, allPages) => { - if (!lastPage || lastPage.length < limit) { + if (lastPage.length < limit) { return undefined; } return allPages.length + 1; }, initialPageParam: 1, - staleTime: 30 * 1000, - refetchOnWindowFocus: true, }); } diff --git a/agentex/src/domain/repositories/task_repository.py b/agentex/src/domain/repositories/task_repository.py index 40c8781e..a228501e 100644 --- a/agentex/src/domain/repositories/task_repository.py +++ b/agentex/src/domain/repositories/task_repository.py @@ -40,7 +40,7 @@ async def list_with_join( agent_id: str | None = None, agent_name: str | None = None, order_by: str | None = None, - order_direction: Literal["asc", "desc"] = "asc", + order_direction: Literal["asc", "desc"] = "desc", limit: int | None = None, page_number: int | None = None, relationships: list[TaskRelationships] | None = None,