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 (