fix: Heartbeat UI refresh - smart polling implementation

aggregator-web/src/hooks/useHeartbeat.ts

@@ -1,31 +1,56 @@
 import { useQuery, useQueryClient } from '@tanstack/react-query';
 import { agentApi } from '@/lib/api';
 import type { UseQueryResult } from '@tanstack/react-query';
+import { useState } from 'react';
 
 export interface HeartbeatStatus {
   enabled: boolean;
   until: string | null;
   active: boolean;
   duration_minutes: number;
+  source?: string | null;
 }
 
-export const useHeartbeatStatus = (agentId: string, enabled: boolean = true): UseQueryResult<HeartbeatStatus, Error> => {
+interface UseHeartbeatResult extends UseQueryResult<HeartbeatStatus, Error> {
-  const queryClient = useQueryClient();
+  recentlyTriggered: boolean;
+  setRecentlyTriggered: (value: boolean) => void;
+}
 
-  return useQuery({
+export const useHeartbeatStatus = (agentId: string, enabled: boolean = true): UseHeartbeatResult => {
+  const queryClient = useQueryClient();
+  const [recentlyTriggered, setRecentlyTriggered] = useState(false);
+
+  const query = useQuery({
     queryKey: ['heartbeat', agentId],
     queryFn: () => agentApi.getHeartbeatStatus(agentId),
     enabled: enabled && !!agentId,
-    staleTime: 1000, // Data is fresh for 1 second
     refetchInterval: (data) => {
-      // Smart polling: fast during active heartbeat, slow when idle
+      // Fast polling after button click (wait for agent to report)
-      if (data?.enabled && data?.active) {
+      if (recentlyTriggered) {
-        return 10000; // 10 seconds when heartbeat is active (catch transitions)
+        return 5000; // 5 seconds
       }
-      return 120000; // 2 minutes when idle (save resources)
+
+      // Fast polling during active operations
+      if (data?.enabled && data?.active) {
+        return 10000; // 10 seconds
+      }
+
+      // Slow polling when idle
+      return 120000; // 2 minutes
     },
-    refetchOnWindowFocus: true, // Refresh when you return to the tab
+    refetchOnWindowFocus: true,
   });
+
+  // Clear flag when agent reports active
+  if (recentlyTriggered && query.data?.active) {
+    setRecentlyTriggered(false);
+  }
+
+  return {
+    ...query,
+    recentlyTriggered,
+    setRecentlyTriggered,
+  };
 };
 
 // Hook to manually invalidate heartbeat cache (used after commands)

aggregator-web/src/pages/Agents.tsx

@@ -242,11 +242,11 @@ const Agents: React.FC = () => {
   const selectedAgent = selectedAgentData || agents.find(a => a.id === id);
 
   // Get heartbeat status for selected agent (smart polling - only when active)
-  const { data: heartbeatStatus } = useHeartbeatStatus(selectedAgent?.id || '', !!selectedAgent);
+  const { data: heartbeatStatus, recentlyTriggered, setRecentlyTriggered } = useHeartbeatStatus(selectedAgent?.id || '', !!selectedAgent);
   const invalidateHeartbeat = useInvalidateHeartbeat();
   const syncAgentData = useHeartbeatAgentSync(selectedAgent?.id || '', heartbeatStatus);
 
-  
+
   // Simple completion handling - clear loading state quickly
   useEffect(() => {
     if (!heartbeatCommandId) return;
@@ -262,6 +262,14 @@ const Agents: React.FC = () => {
     };
   }, [heartbeatCommandId]);
 
+  // Refresh heartbeat status when switching to overview tab
+  useEffect(() => {
+    if (activeTab === 'overview' && selectedAgent?.id) {
+      // Invalidate heartbeat cache to force fresh data on tab switch
+      invalidateHeartbeat(selectedAgent.id);
+    }
+  }, [activeTab, selectedAgent?.id]);
+
   // Filter agents based on OS
   const filteredAgents = agents.filter(agent => {
     if (osFilter === 'all') return true;
@@ -358,8 +366,9 @@ const Agents: React.FC = () => {
       const duration = durationMinutes || heartbeatDuration;
       const result = await agentApi.toggleHeartbeat(agentId, enabled, duration);
 
-      // Immediately invalidate cache to force fresh data
+      // Trigger fast polling for 15 seconds to wait for agent response
-      invalidateHeartbeat(agentId);
+      setRecentlyTriggered(true);
+      setTimeout(() => setRecentlyTriggered(false), 15000);
 
       // Store the command ID for minimal tracking
       if (result.command_id) {