feat(agent-ui): optimize session status action design (#1766)

Author: ulivz

multimodal/tarko/agent-ui/src/common/hooks/useSession.ts

@@ -21,7 +21,6 @@ import {
   deleteSessionAction,
   sendMessageAction,
   abortQueryAction,
-  checkSessionStatusAction,
 } from '../state/actions/sessionActions';
 import {
   initConnectionMonitoringAction,
@@ -64,29 +63,6 @@ export function useSession() {
   const abortQuery = useSetAtom(abortQueryAction);
   const initConnectionMonitoring = useSetAtom(initConnectionMonitoringAction);
   const checkServerStatus = useSetAtom(checkConnectionStatusAction);
-  const checkSessionStatus = useSetAtom(checkSessionStatusAction);
-
-  const statusCheckTimeoutRef = useRef<NodeJS.Timeout | null>(null);
-
-  useEffect(() => {
-    if (!activeSessionId || !connectionStatus.connected || isReplayMode) return;
-
-    if (statusCheckTimeoutRef.current) {
-      clearTimeout(statusCheckTimeoutRef.current);
-    }
-
-    statusCheckTimeoutRef.current = setTimeout(() => {
-      if (activeSessionId && connectionStatus.connected && !isReplayMode) {
-        checkSessionStatus(activeSessionId);
-      }
-    }, 200);
-
-    return () => {
-      if (statusCheckTimeoutRef.current) {
-        clearTimeout(statusCheckTimeoutRef.current);
-      }
-    };
-  }, [activeSessionId, connectionStatus.connected, checkSessionStatus, isReplayMode]);
 
   const sessionState = useMemo(
     () => ({
@@ -122,7 +98,7 @@ export function useSession() {
       initConnectionMonitoring,
       checkServerStatus,
 
-      checkSessionStatus,
+
     }),
     [
       sessions,
@@ -151,7 +127,6 @@ export function useSession() {
       setWorkspaceDisplayState,
       initConnectionMonitoring,
       checkServerStatus,
-      checkSessionStatus,
     ],
   );
 

multimodal/tarko/agent-ui/src/common/state/actions/eventProcessors/handlers/AgentRunHandler.ts

@@ -1,7 +1,6 @@
-import { isProcessingAtom } from '@/common/state/atoms/ui';
+import { sessionProcessingStatesAtom } from '@/common/state/atoms/ui';
 import { AgentEventStream } from '@/common/types';
 import { EventHandler, EventHandlerContext } from '../types';
-import { shouldUpdateProcessingState } from '../utils/panelContentUpdater';
 
 export class AgentRunStartHandler implements EventHandler<AgentEventStream.AgentRunStartEvent> {
   canHandle(event: AgentEventStream.Event): event is AgentEventStream.AgentRunStartEvent {
@@ -15,10 +14,11 @@ export class AgentRunStartHandler implements EventHandler<AgentEventStream.Agent
   ): void {
     const { set } = context;
 
-    // Update processing state
-    if (shouldUpdateProcessingState(sessionId)) {
-      set(isProcessingAtom, true);
-    }
+    // Update session-isolated processing state
+    set(sessionProcessingStatesAtom, (prev) => ({
+      ...prev,
+      [sessionId]: true,
+    }));
   }
 }
 
@@ -30,9 +30,10 @@ export class AgentRunEndHandler implements EventHandler<AgentEventStream.Event>
   handle(context: EventHandlerContext, sessionId: string, event: AgentEventStream.Event): void {
     const { set } = context;
 
-    // Update processing state
-    if (shouldUpdateProcessingState(sessionId)) {
-      set(isProcessingAtom, false);
-    }
+    // Update session-isolated processing state
+    set(sessionProcessingStatesAtom, (prev) => ({
+      ...prev,
+      [sessionId]: false,
+    }));
   }
 }

multimodal/tarko/agent-ui/src/common/state/actions/eventProcessors/utils/panelContentUpdater.ts

@@ -17,12 +17,3 @@ export function shouldUpdatePanelContent(get: Getter, sessionId: string): boolea
 
   return true;
 }
-
-/**
- * Helper function to determine if processing state should be updated for a session
- * Always allow processing state updates for the session that owns the event
- */
-export function shouldUpdateProcessingState(sessionId: string): boolean {
-  // Processing state is now session-isolated, so we always update for the event's session
-  return Boolean(sessionId);
-}

multimodal/tarko/agent-ui/src/common/state/actions/sessionActions.ts

@@ -4,7 +4,7 @@ import { apiService } from '../../services/apiService';
 import { sessionsAtom, activeSessionIdAtom } from '../atoms/session';
 import { messagesAtom } from '../atoms/message';
 import { toolResultsAtom, toolCallResultMap } from '../atoms/tool';
-import { sessionPanelContentAtom, isProcessingAtom } from '../atoms/ui';
+import { sessionPanelContentAtom, sessionProcessingStatesAtom, isProcessingAtom } from '../atoms/ui';
 import { processEventAction } from './eventProcessors';
 import { Message, SessionInfo } from '@/common/types';
 import { connectionStatusAtom } from '../atoms/ui';
@@ -150,7 +150,7 @@ export const setActiveSessionAction = atom(null, async (get, set, sessionId: str
       });
     }
 
-    // Processing state will be managed by SSE events
+
 
     toolCallResultMap.clear();
 
@@ -169,6 +169,27 @@ export const setActiveSessionAction = atom(null, async (get, set, sessionId: str
       }
     }
 
+    // Status recovery with higher priority - always called after events processing
+    // This ensures accurate processing state when SSE connection is established
+    // and overrides any incorrect state from events processing
+    if (!replayState.isActive) {
+      try {
+        const status = await apiService.getSessionStatus(sessionId);
+        set(sessionProcessingStatesAtom, (prev) => ({
+          ...prev,
+          [sessionId]: status.isProcessing,
+        }));
+        console.log(`Recovered processing state for session ${sessionId}: ${status.isProcessing}`);
+      } catch (error) {
+        console.warn(`Failed to recover session status for ${sessionId}:`, error);
+        // Default to false on error to avoid stuck processing state
+        set(sessionProcessingStatesAtom, (prev) => ({
+          ...prev,
+          [sessionId]: false,
+        }));
+      }
+    }
+
     // Always ensure we have the latest session metadata (including modelConfig)
     // This is lightweight since server always provides it
     try {
@@ -410,46 +431,3 @@ export const abortQueryAction = atom(null, async (get, set) => {
     return false;
   }
 });
-
-// Cache to prevent frequent status checks for the same session
-const statusCheckCache = new Map<string, { timestamp: number; promise?: Promise<any> }>();
-const STATUS_CACHE_TTL = 2000; // 2 seconds cache
-
-export const checkSessionStatusAction = atom(null, async (get, set, sessionId: string) => {
-  if (!sessionId) return;
-
-  const now = Date.now();
-  const cached = statusCheckCache.get(sessionId);
-
-  // If we have a recent check or an ongoing request, skip
-  if (cached) {
-    if (cached.promise) {
-      // There's already an ongoing request for this session
-      return cached.promise;
-    }
-    if (now - cached.timestamp < STATUS_CACHE_TTL) {
-      // Recent check, skip
-      return;
-    }
-  }
-
-  try {
-    // Mark that we're making a request
-    const promise = apiService.getSessionStatus(sessionId);
-    statusCheckCache.set(sessionId, { timestamp: now, promise });
-
-    const status = await promise;
-
-    // Update simple processing state
-    set(isProcessingAtom, status.isProcessing);
-
-    // Clear the promise and update timestamp
-    statusCheckCache.set(sessionId, { timestamp: now });
-
-    return status;
-  } catch (error) {
-    console.error('Failed to check session status:', error);
-    // Clear the failed request
-    statusCheckCache.delete(sessionId);
-  }
-});

multimodal/tarko/agent-ui/src/common/state/atoms/ui.ts

@@ -57,9 +57,31 @@ export const sidebarCollapsedAtom = atom<boolean>(true);
 export const workspacePanelCollapsedAtom = atom<boolean>(false);
 
 /**
- * Simple processing state atom based on SSE events
+ * Session-isolated processing state atom based on SSE events
+ * Maps session IDs to their processing states
  */
-export const isProcessingAtom = atom<boolean>(false);
+export const sessionProcessingStatesAtom = atom<Record<string, boolean>>({});
+
+/**
+ * Derived atom for the current active session's processing state
+ * Automatically isolates state by active session
+ */
+export const isProcessingAtom = atom(
+  (get) => {
+    const activeSessionId = get(activeSessionIdAtom);
+    const sessionProcessingStates = get(sessionProcessingStatesAtom);
+    return activeSessionId ? sessionProcessingStates[activeSessionId] ?? false : false;
+  },
+  (get, set, update: boolean) => {
+    const activeSessionId = get(activeSessionIdAtom);
+    if (activeSessionId) {
+      set(sessionProcessingStatesAtom, (prev) => ({
+        ...prev,
+        [activeSessionId]: update,
+      }));
+    }
+  },
+);
 
 /**
  * Workspace display mode - determines what content is shown in the workspace