Async streaming tasks fix

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 (
-    <QueryProvider>
+    <>
       <AgentexProvider
         sgpAppURL={sgpAppURL ?? ''}
         agentexAPIBaseURL={agentexAPIBaseURL}
@@ -82,6 +81,6 @@ export function AgentexUIRoot({
         </div>
       </AgentexProvider>
       <ToastContainer />
-    </QueryProvider>
+    </>
   );
 }

agentex-ui/components/agentex/task-sidebar.tsx

@@ -144,12 +144,10 @@ export function TaskSidebar() {
   const [isCollapsed, setIsCollapsed] = useState<boolean>(false);
   const scrollContainerRef = useRef<HTMLDivElement>(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;

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,
           },
         },
       })

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}</>;

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<Agent[]> => {
       return agentexClient.agents.list();
     },
-    staleTime: 5 * 60 * 1000, // 5 minutes - agents don't change often
     refetchOnWindowFocus: false,
   });
 }

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<TaskListResponse> | undefined
+): InfiniteData<TaskListResponse> | 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<string, unknown>;
 };
 
-/**
- * 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<TaskRetrieveResponse>(
+        tasksKeys.individualById(task.id),
+        task
+      );
+      queryClient.setQueryData<InfiniteData<TaskListResponse>>(
+        tasksKeys.all,
+        data => updateTaskInInfiniteQuery(task, variables.agentName, data)
+      );
+      queryClient.setQueryData<InfiniteData<TaskListResponse>>(
+        tasksKeys.byAgentName(variables.agentName),
+        data => updateTaskInInfiniteQuery(task, variables.agentName, data)
+      );
     },
     onError: error => {
       toast.error({

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<Span[], Error>({
-    queryKey: ['spans', traceId],
+    queryKey: spansKeys.byTraceId(traceId),
     queryFn: async ({ signal }) => {
       if (!traceId) {
         return [];

agentex-ui/hooks/use-task-messages.ts

@@ -50,7 +50,6 @@ export function useTaskMessages({
       };
     },
     enabled: !!taskId,
-    staleTime: Infinity,
     refetchOnMount: false,
     refetchOnWindowFocus: false,
   });

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<Task>(tasksKeys.byId(taskId), task);
-          queryClient.setQueryData<Task[]>(tasksKeys.all, old => {
-            if (!old) return [task];
-            return old.map(t => (t.id === task.id ? task : t));
-          });
+          queryClient.setQueryData<TaskRetrieveResponse>(
+            tasksKeys.individualById(taskId),
+            task
+          );
+          queryClient.setQueryData<InfiniteData<TaskListResponse>>(
+            tasksKeys.all,
+            data => updateTaskInInfiniteQuery(task, agentName, data)
+          );
+          queryClient.setQueryData<InfiniteData<TaskListResponse>>(
+            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]);
 }

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<TaskRetrieveResponse> => {
       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,
   });
 }

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,