HyperCluster/Fara-BrowserUse
5
1import { AgentStep, AgentTrace, AgentTraceMetadata, FinalStep } from '@/types/agent';
2import { create } from 'zustand';
3import { devtools } from 'zustand/middleware';
4
5interface AgentState {
6 // State
7 trace?: AgentTrace;
8 traceId: string | null; // Set by backend heartbeat, persists during connection
9 isAgentProcessing: boolean;
10 isConnectingToE2B: boolean; // New state for E2B connection
11 vncUrl: string;
12 selectedModelId: string;
13 availableModels: string[];
14 isLoadingModels: boolean;
15 isConnected: boolean;
16 error?: string;
17 isDarkMode: boolean;
18 selectedStepIndex: number | null; // null = live mode, number = viewing specific step or 'final'
19 finalStep?: FinalStep; // Special step for success/failure
20
21 // Actions
22 setTrace: (trace: AgentTrace | undefined) => void;
23 setTraceId: (traceId: string | null) => void;
24 updateTraceWithStep: (step: AgentStep, metadata: AgentTraceMetadata) => void;
25 updateStepEvaluation: (stepId: string, evaluation: 'like' | 'dislike' | 'neutral') => void;
26 updateTraceEvaluation: (evaluation: 'success' | 'failed' | 'not_evaluated') => void;
27 completeTrace: (metadata: AgentTraceMetadata, finalState?: 'success' | 'stopped' | 'max_steps_reached' | 'error' | 'sandbox_timeout') => void;
28 setIsAgentProcessing: (processing: boolean) => void;
29 setIsConnectingToE2B: (connecting: boolean) => void;
30 setVncUrl: (url: string) => void;
31 setSelectedModelId: (modelId: string) => void;
32 setAvailableModels: (models: string[]) => void;
33 setIsLoadingModels: (loading: boolean) => void;
34 setIsConnected: (connected: boolean) => void;
35 setError: (error: string | undefined) => void;
36 setSelectedStepIndex: (index: number | null) => void;
37 toggleDarkMode: () => void;
38 resetAgent: () => void;
39}
40
41const initialState = {
42 trace: undefined,
43 traceId: null, // Will be set by backend heartbeat
44 isAgentProcessing: false,
45 isConnectingToE2B: false,
46 vncUrl: '',
47 selectedModelId: 'Qwen/Qwen3-VL-30B-A3B-Instruct',
48 availableModels: [],
49 isLoadingModels: false,
50 isConnected: false,
51 error: undefined,
52 isDarkMode: true, // Default to dark mode
53 selectedStepIndex: null, // null = live mode
54 finalStep: undefined,
55};
56
57export const useAgentStore = create<AgentState>()(
58 devtools(
59 (set) => ({
60 ...initialState,
61
62 // Set the complete trace
63 setTrace: (trace) =>
64 set({ trace }, false, 'setTrace'),
65
66 // Set trace ID (set by backend heartbeat, only cleared on disconnect)
67 setTraceId: (traceId) =>
68 set({ traceId }, false, 'setTraceId'),
69
70 // Update trace with a new step
71 updateTraceWithStep: (step, metadata) =>
72 set(
73 (state) => {
74 if (!state.trace) return state;
75
76 const existingSteps = state.trace.steps || [];
77 const stepExists = existingSteps.some((s) => s.stepId === step.stepId);
78
79 if (stepExists) return state;
80
81 // Preserve existing maxSteps if new metadata has 0
82 const updatedMetadata = {
83 ...metadata,
84 maxSteps: metadata.maxSteps > 0
85 ? metadata.maxSteps
86 : (state.trace.traceMetadata?.maxSteps || 200),
87 };
88
89 return {
90 trace: {
91 ...state.trace,
92 steps: [...existingSteps, step],
93 traceMetadata: updatedMetadata,
94 isRunning: true,
95 },
96 };
97 },
98 false,
99 'updateTraceWithStep'
100 ),
101
102 // Update step evaluation in the store
103 updateStepEvaluation: (stepId, evaluation) =>
104 set(
105 (state) => {
106 if (!state.trace || !state.trace.steps) return state;
107
108 const updatedSteps = state.trace.steps.map((step) =>
109 step.stepId === stepId
110 ? { ...step, step_evaluation: evaluation }
111 : step
112 );
113
114 return {
115 trace: {
116 ...state.trace,
117 steps: updatedSteps,
118 },
119 };
120 },
121 false,
122 'updateStepEvaluation'
123 ),
124
125 // Update trace evaluation in the store
126 updateTraceEvaluation: (evaluation) =>
127 set(
128 (state) => {
129 if (!state.trace || !state.trace.traceMetadata) return state;
130
131 const updatedMetadata = {
132 ...state.trace.traceMetadata,
133 user_evaluation: evaluation,
134 };
135
136 return {
137 trace: {
138 ...state.trace,
139 traceMetadata: updatedMetadata,
140 },
141 // Also update finalStep metadata if it exists
142 finalStep: state.finalStep ? {
143 ...state.finalStep,
144 metadata: {
145 ...state.finalStep.metadata,
146 user_evaluation: evaluation,
147 },
148 } : state.finalStep,
149 };
150 },
151 false,
152 'updateTraceEvaluation'
153 ),
154
155 // Complete the trace
156 completeTrace: (metadata, finalState?: 'success' | 'stopped' | 'max_steps_reached' | 'error' | 'sandbox_timeout') =>
157 set(
158 (state) => {
159 if (!state.trace) return state;
160
161 // Preserve existing maxSteps if new metadata has 0
162 const updatedMetadata = {
163 ...metadata,
164 maxSteps: metadata.maxSteps > 0
165 ? metadata.maxSteps
166 : (state.trace.traceMetadata?.maxSteps || 200),
167 completed: true,
168 };
169
170 // Determine the final step type based on final_state from backend
171 let stepType: 'success' | 'failure' | 'stopped' | 'max_steps_reached' | 'sandbox_timeout';
172 let stepMessage: string | undefined;
173
174 if (finalState === 'stopped') {
175 stepType = 'stopped';
176 stepMessage = 'Task stopped by user';
177 } else if (finalState === 'max_steps_reached') {
178 stepType = 'max_steps_reached';
179 stepMessage = 'Maximum steps reached';
180 } else if (finalState === 'sandbox_timeout') {
181 stepType = 'sandbox_timeout';
182 stepMessage = 'Sandbox timeout';
183 } else if (finalState === 'error' || state.error) {
184 stepType = 'failure';
185 stepMessage = state.error || 'Task failed';
186 } else {
187 stepType = 'success';
188 stepMessage = undefined;
189 }
190
191 const finalStep: FinalStep = {
192 type: stepType,
193 message: stepMessage,
194 metadata: updatedMetadata,
195 };
196
197 return {
198 trace: {
199 ...state.trace,
200 isRunning: false,
201 traceMetadata: updatedMetadata,
202 },
203 finalStep,
204 // Keep error in state for display
205 selectedStepIndex: null, // Reset to live mode on completion
206 };
207 },
208 false,
209 'completeTrace'
210 ),
211
212 // Set processing state
213 setIsAgentProcessing: (isAgentProcessing) =>
214 set({ isAgentProcessing }, false, 'setIsAgentProcessing'),
215
216 // Set E2B connection state
217 setIsConnectingToE2B: (isConnectingToE2B) =>
218 set({ isConnectingToE2B }, false, 'setIsConnectingToE2B'),
219
220 // Set VNC URL
221 setVncUrl: (vncUrl) =>
222 set({ vncUrl }, false, 'setVncUrl'),
223
224 // Set selected model ID
225 setSelectedModelId: (selectedModelId) =>
226 set({ selectedModelId }, false, 'setSelectedModelId'),
227
228 // Set available models
229 setAvailableModels: (availableModels) =>
230 set({ availableModels }, false, 'setAvailableModels'),
231
232 // Set loading models state
233 setIsLoadingModels: (isLoadingModels) =>
234 set({ isLoadingModels }, false, 'setIsLoadingModels'),
235
236 // Set connection status
237 setIsConnected: (isConnected) =>
238 set({ isConnected }, false, 'setIsConnected'),
239
240 // Set error
241 setError: (error) =>
242 set(
243 (state) => {
244 // If there's an error and a trace, mark it as failed
245 if (error && state.trace) {
246 const metadata = state.trace.traceMetadata || {
247 traceId: state.trace.id,
248 inputTokensUsed: 0,
249 outputTokensUsed: 0,
250 duration: 0,
251 numberOfSteps: state.trace.steps?.length || 0,
252 maxSteps: 200,
253 completed: false,
254 final_state: null,
255 user_evaluation: 'not_evaluated' as const,
256 };
257
258 // Ensure maxSteps is not 0
259 const finalMetadata: AgentTraceMetadata = {
260 ...metadata,
261 maxSteps: metadata.maxSteps > 0 ? metadata.maxSteps : 200,
262 final_state: metadata.final_state || null,
263 user_evaluation: metadata.user_evaluation || 'not_evaluated',
264 };
265
266 const finalStep: FinalStep = {
267 type: 'failure',
268 message: error,
269 metadata: finalMetadata,
270 };
271
272 return {
273 error,
274 finalStep,
275 trace: {
276 ...state.trace,
277 isRunning: false,
278 },
279 selectedStepIndex: null, // Reset to live mode on error
280 };
281 }
282 return { error };
283 },
284 false,
285 'setError'
286 ),
287
288 // Set selected step index for time travel
289 setSelectedStepIndex: (selectedStepIndex) =>
290 set({ selectedStepIndex }, false, 'setSelectedStepIndex'),
291
292 // Toggle dark mode
293 toggleDarkMode: () =>
294 set((state) => ({ isDarkMode: !state.isDarkMode }), false, 'toggleDarkMode'),
295
296 // Reset agent state (but preserve traceId from backend during connection)
297 resetAgent: () =>
298 set((state) => ({
299 ...initialState,
300 traceId: state.traceId, // IMPORTANT: Keep traceId from backend
301 isDarkMode: state.isDarkMode, // Keep dark mode preference
302 isConnected: state.isConnected, // Keep connection status
303 selectedModelId: state.selectedModelId, // Keep selected model
304 availableModels: state.availableModels, // Keep available models
305 isLoadingModels: state.isLoadingModels // Keep loading state
306 }), false, 'resetAgent'),
307 }),
308 { name: 'AgentStore' }
309 )
310);
311
312// Selectors for better performance
313export const selectTrace = (state: AgentState) => state.trace;
314export const selectTraceId = (state: AgentState) => state.traceId;
315export const selectIsAgentProcessing = (state: AgentState) => state.isAgentProcessing;
316export const selectIsConnectingToE2B = (state: AgentState) => state.isConnectingToE2B;
317export const selectVncUrl = (state: AgentState) => state.vncUrl;
318export const selectSelectedModelId = (state: AgentState) => state.selectedModelId;
319export const selectAvailableModels = (state: AgentState) => state.availableModels;
320export const selectIsLoadingModels = (state: AgentState) => state.isLoadingModels;
321export const selectIsConnected = (state: AgentState) => state.isConnected;
322export const selectSteps = (state: AgentState) => state.trace?.steps;
323export const selectMetadata = (state: AgentState) => state.trace?.traceMetadata;
324export const selectError = (state: AgentState) => state.error;
325export const selectIsDarkMode = (state: AgentState) => state.isDarkMode;
326export const selectSelectedStepIndex = (state: AgentState) => state.selectedStepIndex;
327export const selectFinalStep = (state: AgentState) => state.finalStep;
328
329// Composite selector for selected step (avoids infinite loops)
330export const selectSelectedStep = (state: AgentState) => {
331 const steps = state.trace?.steps;
332 const selectedIndex = state.selectedStepIndex;
333
334 if (selectedIndex === null || !steps || selectedIndex >= steps.length) {
335 return null;
336 }
337
338 return steps[selectedIndex];
339};
340