cua2-front/src/hooks/useAgentWebSocket.ts

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,
  };
};