HyperCluster/Fara-BrowserUse
5
1import { useAgentStore } from '@/stores/agentStore';
2import { uploadTraceToModal } from '@/services/api';
3import { getTraceExportData } from '@/services/jsonExporter';
4import { AgentTrace, AgentTraceMetadata, WebSocketEvent } from '@/types/agent';
5import { useCallback, useEffect } from 'react';
6import { useWebSocket } from './useWebSocket';
7
8interface UseAgentWebSocketOptions {
9 url: string;
10}
11
12export const useAgentWebSocket = ({ url }: UseAgentWebSocketOptions) => {
13 const {
14 setTrace,
15 traceId,
16 setTraceId,
17 updateTraceWithStep,
18 completeTrace,
19 setIsAgentProcessing,
20 setIsConnectingToE2B,
21 setVncUrl,
22 setError,
23 setIsConnected,
24 selectedModelId,
25 resetAgent,
26 } = useAgentStore();
27
28 // Handle incoming WebSocket messages
29 const handleWebSocketMessage = useCallback(
30 (event: WebSocketEvent) => {
31 console.log('WebSocket event received:', event);
32
33 switch (event.type) {
34 case 'agent_start': {
35 // Clear previous state (especially finalStep)
36 resetAgent();
37
38 setIsAgentProcessing(true);
39 setIsConnectingToE2B(true); // Start connecting to E2B
40 setError(undefined); // Clear any previous error
41
42 // Ensure trace has proper metadata with default maxSteps if not provided
43 const traceWithMetadata = {
44 ...event.agentTrace,
45 traceMetadata: event.agentTrace.traceMetadata ? {
46 ...event.agentTrace.traceMetadata,
47 maxSteps: event.agentTrace.traceMetadata.maxSteps > 0
48 ? event.agentTrace.traceMetadata.maxSteps
49 : 200, // Default if backend sends 0
50 } : {
51 traceId: event.agentTrace.id,
52 inputTokensUsed: 0,
53 outputTokensUsed: 0,
54 duration: 0,
55 numberOfSteps: 0,
56 maxSteps: 200,
57 completed: false,
58 final_state: null,
59 },
60 };
61
62 setTrace(traceWithMetadata);
63 console.log('Agent start received:', traceWithMetadata);
64 break;
65 }
66
67 case 'agent_progress':
68 // Add new step from agent trace run with image, generated text, actions, tokens and timestamp
69 setIsConnectingToE2B(false); // Connected! First step received
70 updateTraceWithStep(event.agentStep, event.traceMetadata);
71 console.log('Agent progress received:', event.agentStep);
72 break;
73
74 case 'agent_complete':
75 setIsAgentProcessing(false);
76 setIsConnectingToE2B(false);
77 completeTrace(event.traceMetadata, event.final_state);
78 console.log('Agent complete received:', event.traceMetadata, 'Final state:', event.final_state);
79
80 // Auto-upload trace when task completes
81 // Use setTimeout to ensure store is updated first
82 setTimeout(() => {
83 const state = useAgentStore.getState();
84 const trace = state.trace;
85 const steps = state.trace?.steps || [];
86 const finalStep = state.finalStep;
87 const metadata = state.trace?.traceMetadata;
88
89 if (trace && !trace.isRunning) {
90 console.log('🚀 Auto-uploading trace after task completion...');
91 const traceData = getTraceExportData(trace, steps, metadata, finalStep);
92 uploadTraceToModal(traceData)
93 .then((result) => {
94 if (result.success) {
95 console.log('✅ Trace auto-uploaded successfully:', result);
96 } else {
97 console.warn('⚠️ Trace auto-upload failed:', result.error);
98 }
99 })
100 .catch((error) => {
101 console.error('❌ Error auto-uploading trace:', error);
102 });
103 }
104 }, 100);
105 break;
106
107 case 'agent_error':
108 setIsAgentProcessing(false);
109 setIsConnectingToE2B(false);
110 setError(event.error);
111 console.error('Agent error received:', event.error);
112 break;
113
114 case 'vnc_url_set':
115 setIsConnectingToE2B(false); // Connected! VNC URL received
116 setVncUrl(event.vncUrl);
117 console.log('VNC URL set received:', event.vncUrl);
118 break;
119
120 case 'vnc_url_unset':
121 setVncUrl('');
122 console.log('VNC URL unset received');
123 break;
124
125 case 'heartbeat':
126 console.log('Heartbeat received:', event);
127 setTraceId(event.uuid);
128 console.log('TraceId set from backend:', event.uuid);
129 break;
130
131 }
132 },
133 [setTrace, updateTraceWithStep, completeTrace, setIsAgentProcessing, setIsConnectingToE2B, setVncUrl, setError, resetAgent, setTraceId, traceId]
134 );
135
136 // Handle WebSocket errors
137 const handleWebSocketError = useCallback(() => {
138 // WebSocket Frontend Error handling
139 console.error('WebSocket connection error');
140 }, []);
141
142 // Initialize WebSocket connection
143 const { isConnected, connectionState, sendMessage, manualReconnect } = useWebSocket({
144 url,
145 onMessage: handleWebSocketMessage,
146 onError: handleWebSocketError,
147 });
148
149 // Sync connection state to store and clear traceId on disconnect
150 useEffect(() => {
151 setIsConnected(isConnected);
152
153 // Clear traceId when websocket disconnects
154 if (!isConnected) {
155 setTraceId(null);
156 console.log('WebSocket disconnected - traceId cleared');
157 }
158 }, [isConnected, setIsConnected, setTraceId]);
159
160 // Create a global sendNewTask function that can be called from anywhere
161 useEffect(() => {
162 // Store sendNewTask in window for global access
163 (window as Window & { __sendNewTask?: (instruction: string, modelId: string) => void }).__sendNewTask = (instruction: string, modelId: string) => {
164 // Reset agent state before starting a new task
165 resetAgent();
166
167 // Ensure traceId is set before creating trace
168 if (!traceId) {
169 console.error('Internal error: Cannot send task. TraceId not set. Refreshing page...');
170 window.location.reload();
171 return;
172 }
173
174 const trace: AgentTrace = {
175 id: traceId,
176 instruction,
177 modelId: modelId,
178 timestamp: new Date(),
179 isRunning: true,
180 traceMetadata: {
181 traceId: traceId,
182 inputTokensUsed: 0,
183 outputTokensUsed: 0,
184 duration: 0,
185 numberOfSteps: 0,
186 maxSteps: 200, // Default max steps, will be updated by backend
187 completed: false,
188 final_state: null,
189 } as AgentTraceMetadata,
190 };
191
192 setTrace(trace);
193 setIsAgentProcessing(true);
194 setIsConnectingToE2B(true); // Start connecting when task is sent
195
196 // Send message to Python backend via WebSocket
197 sendMessage({
198 type: 'user_task',
199 trace: trace,
200 });
201
202 console.log('Task sent:', trace);
203 };
204 }, [setTrace, setIsAgentProcessing, setIsConnectingToE2B, sendMessage, resetAgent, traceId]);
205
206 // Function to stop the current task
207 const stopCurrentTask = useCallback(() => {
208 const trace = useAgentStore.getState().trace;
209 if (trace?.id && trace.isRunning) {
210 sendMessage({
211 type: 'stop_task',
212 trace_id: trace.id,
213 });
214 console.log('Stop task sent for trace:', trace.id);
215
216 // Don't update UI state here - wait for backend to send agent_complete with final_state='stopped'
217 }
218 }, [sendMessage]);
219
220 return {
221 isConnected,
222 connectionState,
223 manualReconnect,
224 stopCurrentTask,
225 };
226};
227