Basic web loop works

web/src/components/App.jsx

@@ -24,35 +24,41 @@ const App = () => {
   const [showSidebar, setShowSidebar] = useState(false);
   const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
 
-  const { wsState, sendMessage, setMessageHandler } = useWebSocket('ws://localhost:8080/ws');
+  const { wsState, sendMessage, addMessageHandler } = useWebSocket('ws://localhost:8080/ws');
 
-  // Set up WebSocket message handlers
   useEffect(() => {
-    setMessageHandler(ServerMessageType.STATE_CHANGE, (message) => {
-      setAgentState(message.state);
+    addMessageHandler((message) => {
+      switch(message.type) {
+        case ServerMessageType.STATE_CHANGE:
+          setAgentState(message.state);
+          break;
+        case ServerMessageType.CONTEXT_UPDATE:
+          setOriginalContext(message.context);
+          setModifiedContext(message.context);
+          break;
+        case ServerMessageType.RESPONSE_UPDATE:
+          setOriginalResponse(message.response);
+          setModifiedResponse(message.response);
+          setValidationError(message.validation_error);
+          break;
+        case ServerMessageType.OUTPUT_UPDATE:
+          setOutput(prev => prev + message.data);
+          break;
+      }
     });
-
-    setMessageHandler(ServerMessageType.CONTEXT_UPDATE, (message) => {
-      setOriginalContext(message.context);
-      setModifiedContext(message.context);
-    });
-
-    setMessageHandler(ServerMessageType.RESPONSE_UPDATE, (message) => {
-      setOriginalResponse(message.response);
-      setModifiedResponse(message.response);
-      setValidationError(message.validation_error);
-    });
-
-    setMessageHandler(ServerMessageType.OUTPUT_UPDATE, (message) => {
-      setOutput(prev => prev + message.data);
-    });
-  }, [setMessageHandler]);
+  }, [addMessageHandler]);
 
   const handleApprove = () => {
     if (agentState === AgentState.CONTEXT_APPROVAL) {
-      sendMessage(ClientMessageType.APPROVE_CONTEXT);
+      sendMessage({
+        "type": ClientMessageType.APPROVE_CONTEXT,
+        "context": modifiedContext
+      });
     } else if (agentState === AgentState.RESPONSE_APPROVAL) {
-      sendMessage(ClientMessageType.APPROVE_RESPONSE, modifiedResponse);
+      sendMessage({
+        "type": ClientMessageType.APPROVE_RESPONSE,
+        "response": modifiedResponse
+      });
     }
   };
 
@@ -63,15 +69,28 @@ const App = () => {
     }
   };
 
+  useEffect(() => {
+    if (agentState === AgentState.APPROVE_RESPONSE) {
+      sendMessage({
+        "type": ClientMessageType.MODIFY_RESPONSE,
+        "response": modifiedResponse
+      });
+    }
+  }, [modifiedResponse]);
+
   useEffect(() => {
     switch (agentState) {
+      case AgentState.UPDATE: // fall through
       case AgentState.CONTEXT_APPROVAL:
         setActiveTab(Tabs.CONTEXT);
         break;
+      case AgentState.INFERENCE: // fall through
       case AgentState.RESPONSE_APPROVAL:
         setActiveTab(Tabs.RESPONSE);
         break;
     }
+    setShowDiff(false);
+    setShowSidebar(false);
   }, [agentState]);
 
   return (

web/src/constants.jsx

@@ -19,7 +19,6 @@ export const Tabs = {
 };
 
 export const ClientMessageType = {
-  MODIFY_CONTEXT: 'MODIFY_CONTEXT',
   APPROVE_CONTEXT: 'APPROVE_CONTEXT',
   MODIFY_RESPONSE: 'MODIFY_RESPONSE',
   APPROVE_RESPONSE: 'APPROVE_RESPONSE',

web/src/hooks/useWebSocket.jsx

@@ -4,7 +4,7 @@ import { WebSocketState, MessageType } from '../constants';
 export const useWebSocket = (url) => {
   const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
   const wsRef = useRef(null);
-  const messageHandlersRef = useRef({});
+  const handlersRef = useRef([]);
 
   const connect = useCallback(() => {
     setWsState(WebSocketState.CONNECTING);
@@ -27,12 +27,7 @@ export const useWebSocket = (url) => {
     ws.onmessage = (event) => {
       try {
         const message = JSON.parse(event.data);
-        const handler = messageHandlersRef.current[message.type];
-        if (handler) {
-          handler(message);
-        } else {
-          console.warn('No handler for message type:', message.type);
-        }
+        handlersRef.current.forEach(handler => handler(message));
       } catch (error) {
         console.error('Error handling WebSocket message:', error);
       }
@@ -48,25 +43,21 @@ export const useWebSocket = (url) => {
     };
   }, [connect]);
 
-  const sendMessage = useCallback((type, data = null) => {
+  const sendMessage = useCallback((data) => {
     if (wsRef.current?.readyState === WebSocket.OPEN) {
-      const message = { 
-        type,
-        data: data || undefined
-      };
-      wsRef.current.send(JSON.stringify(message));
+      wsRef.current.send(JSON.stringify(data));
     } else {
       console.error('WebSocket not connected, state:', wsRef.current?.readyState);
     }
   }, []);
 
-  const setMessageHandler = useCallback((type, handler) => {
-    messageHandlersRef.current[type] = handler;
+  const addMessageHandler = useCallback((handler) => {
+    handlersRef.current.push(handler);
   }, []);
 
   return {
     wsState,
     sendMessage,
-    setMessageHandler
+    addMessageHandler
   };
 };