Fix issue with new async message sse being closed

Author: declan-scale

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

@@ -1,18 +1,14 @@
 'use client';
 
-import {
-  InfiniteData,
-  useMutation,
-  useQueryClient,
-} from '@tanstack/react-query';
+import { useMutation, useQueryClient } from '@tanstack/react-query';
 import { agentRPCNonStreaming } from 'agentex/lib';
 
 import { toast } from '@/components/agentex/toast';
 
 import { tasksKeys } from './use-tasks';
 
 import type AgentexSDK from 'agentex';
-import type { Task, TaskListResponse } from 'agentex/resources';
+import type { Task } from 'agentex/resources';
 
 type CreateTaskParams = {
   agentName: string;
@@ -49,40 +45,11 @@ export function useCreateTask({
 
       return response.result;
     },
-    onSuccess: (newTask, variables) => {
-      // Helper function to update infinite query cache
-      const updateInfiniteCache = (queryKey: readonly unknown[]) => {
-        queryClient.setQueryData<InfiniteData<TaskListResponse>>(
-          queryKey,
-          old => {
-            if (!old || !old.pages || old.pages.length === 0) {
-              // If no cache exists, create initial structure
-              return {
-                pages: [[newTask]],
-                pageParams: [1],
-              };
-            }
-
-            // Add new task to the first page (prepend to show at top)
-            const firstPage = old.pages[0] ?? [];
-            if (firstPage.some(t => t.id === newTask.id)) {
-              return old; // Avoid duplicates
-            }
-
-            return {
-              ...old,
-              pages: [[newTask, ...firstPage], ...old.pages.slice(1)],
-            };
-          }
-        );
-      };
-
-      // Update both the agent-specific cache and the generic "all tasks" cache
-      updateInfiniteCache(tasksKeys.byAgentName(variables.agentName));
-      updateInfiniteCache(tasksKeys.all);
-
-      // Invalidate all task queries to ensure consistency
+    onSuccess: (_, variables) => {
       queryClient.invalidateQueries({ queryKey: tasksKeys.all });
+      queryClient.invalidateQueries({
+        queryKey: tasksKeys.byAgentName(variables.agentName),
+      });
     },
     onError: error => {
       toast.error({

agentex-ui/hooks/use-tasks.ts

@@ -62,7 +62,7 @@ export function useInfiniteTasks(
       return agentexClient.tasks.list(params);
     },
     getNextPageParam: (lastPage, allPages) => {
-      if (lastPage.length < limit) {
+      if (!lastPage || lastPage.length < limit) {
         return undefined;
       }
       return allPages.length + 1;