diff --git a/aggregator-web/src/hooks/useHeartbeat.ts b/aggregator-web/src/hooks/useHeartbeat.ts index 048bfa4..472991c 100644 --- a/aggregator-web/src/hooks/useHeartbeat.ts +++ b/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 => { - const queryClient = useQueryClient(); +interface UseHeartbeatResult extends UseQueryResult { + 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 - if (data?.enabled && data?.active) { - return 10000; // 10 seconds when heartbeat is active (catch transitions) + // Fast polling after button click (wait for agent to report) + if (recentlyTriggered) { + 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) diff --git a/aggregator-web/src/pages/Agents.tsx b/aggregator-web/src/pages/Agents.tsx index 138c2e9..df10170 100644 --- a/aggregator-web/src/pages/Agents.tsx +++ b/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 - invalidateHeartbeat(agentId); + // Trigger fast polling for 15 seconds to wait for agent response + setRecentlyTriggered(true); + setTimeout(() => setRecentlyTriggered(false), 15000); // Store the command ID for minimal tracking if (result.command_id) {