Spaces:
Running
Running
File size: 5,205 Bytes
67b042f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
import { useCallback, useEffect } from 'react';
import { useWebSocket } from './useWebSocket';
import { useAgentStore } from '@/stores/agentStore';
import { WebSocketEvent, AgentTrace, AgentStep } from '@/types/agent';
import { ulid } from 'ulid';
interface UseAgentWebSocketOptions {
url: string;
}
export const useAgentWebSocket = ({ url }: UseAgentWebSocketOptions) => {
const {
setTrace,
updateTraceWithStep,
completeTrace,
setIsAgentProcessing,
setIsConnectingToE2B,
setVncUrl,
setError,
setIsConnected,
selectedModelId,
resetAgent,
} = useAgentStore();
// Handle incoming WebSocket messages
const handleWebSocketMessage = useCallback(
(event: WebSocketEvent) => {
console.log('WebSocket event received:', event);
switch (event.type) {
case 'agent_start': {
// Clear previous state (especially finalStep)
resetAgent();
setIsAgentProcessing(true);
setIsConnectingToE2B(true); // Start connecting to E2B
setError(undefined); // Clear any previous error
// Ensure trace has proper metadata with default maxSteps if not provided
const traceWithMetadata = {
...event.agentTrace,
traceMetadata: event.agentTrace.traceMetadata ? {
...event.agentTrace.traceMetadata,
maxSteps: event.agentTrace.traceMetadata.maxSteps > 0
? event.agentTrace.traceMetadata.maxSteps
: 200, // Default if backend sends 0
} : {
traceId: event.agentTrace.id,
inputTokensUsed: 0,
outputTokensUsed: 0,
duration: 0,
numberOfSteps: 0,
maxSteps: 200,
completed: false,
},
};
setTrace(traceWithMetadata);
console.log('Agent start received:', traceWithMetadata);
break;
}
case 'agent_progress':
// Add new step from agent trace run with image, generated text, actions, tokens and timestamp
setIsConnectingToE2B(false); // Connected! First step received
updateTraceWithStep(event.agentStep, event.traceMetadata);
console.log('Agent progress received:', event.agentStep);
break;
case 'agent_complete':
setIsAgentProcessing(false);
setIsConnectingToE2B(false);
completeTrace(event.traceMetadata);
console.log('Agent complete received:', event.traceMetadata);
break;
case 'agent_error':
setIsAgentProcessing(false);
setIsConnectingToE2B(false);
setError(event.error);
console.error('Agent error received:', event.error);
break;
case 'vnc_url_set':
setIsConnectingToE2B(false); // Connected! VNC URL received
setVncUrl(event.vncUrl);
console.log('VNC URL set received:', event.vncUrl);
break;
case 'vnc_url_unset':
setVncUrl('');
console.log('VNC URL unset received');
break;
case 'heartbeat':
console.log('Heartbeat received:', event);
break;
}
},
[setTrace, updateTraceWithStep, completeTrace, setIsAgentProcessing, setIsConnectingToE2B, setVncUrl, setError, resetAgent]
);
// Handle WebSocket errors
const handleWebSocketError = useCallback(() => {
// WebSocket Frontend Error handling
console.error('WebSocket connection error');
}, []);
// Initialize WebSocket connection
const { isConnected, connectionState, sendMessage, manualReconnect } = useWebSocket({
url,
onMessage: handleWebSocketMessage,
onError: handleWebSocketError,
});
// Sync connection state to store
useEffect(() => {
setIsConnected(isConnected);
}, [isConnected, setIsConnected]);
// Create a global sendNewTask function that can be called from anywhere
useEffect(() => {
// Store sendNewTask in window for global access
(window as Window & { __sendNewTask?: (instruction: string, modelId: string) => void }).__sendNewTask = (instruction: string, modelId: string) => {
// Reset agent state before starting a new task
resetAgent();
const traceId = ulid();
const trace: AgentTrace = {
id: traceId,
instruction,
modelId: modelId,
timestamp: new Date(),
isRunning: true,
traceMetadata: {
traceId: traceId,
inputTokensUsed: 0,
outputTokensUsed: 0,
duration: 0,
numberOfSteps: 0,
maxSteps: 200, // Default max steps, will be updated by backend
completed: false,
},
};
setTrace(trace);
setIsAgentProcessing(true);
setIsConnectingToE2B(true); // Start connecting when task is sent
// Send message to Python backend via WebSocket
sendMessage({
type: 'user_task',
trace: trace,
});
console.log('Task sent:', trace);
};
}, [setTrace, setIsAgentProcessing, setIsConnectingToE2B, sendMessage, resetAgent]);
return {
isConnected,
connectionState,
manualReconnect,
};
};
|