Approve context

sia/__main__.py

@@ -43,9 +43,20 @@ async def run_web_server():
     await server.start()
 
     try:
-        # Keep running until interrupted
         while True:
-            await asyncio.sleep(1)
+            if server._response_tasks:
+                done, pending = await asyncio.wait(
+                    server._response_tasks,
+                    timeout=0.1,
+                    return_when=asyncio.FIRST_COMPLETED
+                )
+                for task in done:
+                    try:
+                        await task
+                    except Exception as e:
+                        print(f"Task error: {e}")
+                    server._response_tasks.remove(task)
+            await asyncio.sleep(0.1)
     except KeyboardInterrupt:
         print("\nShutting down...")
     finally:

sia/web_agent.py

@@ -54,7 +54,6 @@ class WebAgent(BaseAgent):
         self._command_result: Optional[CommandResult] = None
         self.context = self._compile_context()
 
-        
     @property
     def current_state(self) -> WebAgentState:
         """Get the current state of the agent."""
@@ -93,43 +92,68 @@ class WebAgent(BaseAgent):
     def _notify_response_change(self) -> None:
         """Notify all handlers of response change."""
         for handler in self._response_change_handlers:
-            handler(self._response)
+            handler(self.response)
     
     def _set_response(self, response: str) -> None:
+        """Set response and notify handlers."""
+        self.response = response
+        self._validation_error = self._validator.validate(response)
+        for handler in self._response_change_handlers:
+            try:
+                handler(response)
+            except Exception as e:
+                print(f"Error in response handler: {e}")
+
+    def _set_state(self, new_state: WebAgentState) -> None:
         """
-        Set the current response and notify handlers.
+        Set the current state and notify handlers.
 
         Args:
-            response: New response
+            new_state: New state to set
         """
-        self._response = response
-        self._validation_error = self._validator.validate(self._response)
-        self._notify_response_change()
-        
-    def approve_context(self) -> None:
-        if self._current_state != WebAgentState.CONTEXT_APPROVAL:
-            raise Exception("Not in CONTEXT_APPROVAL state")
-
-        self._set_response("")
-        self._current_state = WebAgentState.INFERENCE
-        self._notify_state_change()
-
-        response_token_iter = self._llm.infer(self._system_prompt, self.context)
-
-        for token in response_token_iter:
-            self._set_response(self._response + token)
-
-        self._current_state = WebAgentState.RESPONSE_APPROVAL
+        self._current_state = new_state
         self._notify_state_change()
             
+    async def approve_context(self) -> None:
+        if self._current_state != WebAgentState.CONTEXT_APPROVAL:
+            error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._current_state})"
+            raise Exception(error_msg)
+                
+        self._set_state(WebAgentState.INFERENCE)
+        self._set_response("")
+        
+        response_token_iter = self._llm.infer(self._system_prompt, self.context)
+        response = ""
+        try:
+            for token in response_token_iter:
+                response += token
+                await self._set_response(response)
+                print(f"{token}")
+        except Exception as e:
+            print(f"Error during inference: {e}")
+            self._set_state(WebAgentState.CONTEXT_APPROVAL)
+            return
+        
+        await self._set_response(response)
+        self._set_state(WebAgentState.RESPONSE_APPROVAL)
+
+    async def _set_response(self, response: str) -> None:
+        """Set response and notify handlers asynchronously."""
+        self.response = response
+        self._validation_error = self._validator.validate(response)
+        for handler in self._response_change_handlers:
+            try:
+                await handler(response)
+            except Exception as e:
+                print(f"Error in response handler: {e}")
+
     def approve_response(self) -> None:
         if self._current_state != WebAgentState.RESPONSE_APPROVAL:
             raise Exception("Not in RESPONSE_APPROVAL state")
 
-        self._current_state = WebAgentState.UPDATE
-        self._notify_state_change()
+        self._set_state(WebAgentState.UPDATE)
 
-        parse_result = self._parser.parse(self._response)
+        parse_result = self._parser.parse(self.response)
         if isinstance(parse_result, Command):
             result = parse_result.execute(self._working_memory)
             self._command_result = result
@@ -137,7 +161,6 @@ class WebAgent(BaseAgent):
             self._working_memory.add_entry(parse_result)
         
         self._working_memory.update()
-        self._context = self._compile_context()
+        self.context = self._compile_context()
 
-        self._current_state = WebAgentState.CONTEXT_APPROVAL
-        self._notify_state_change()
+        self._set_state(WebAgentState.CONTEXT_APPROVAL)

sia/web_io_buffer.py

@@ -1,38 +1,38 @@
+from threading import Lock
+
 from .io_buffer import IOBuffer
 
 class WebIOBuffer(IOBuffer):
     """
-    IOBuffer implementation for web interface communication.
-    
-    This class manages input/output through memory buffers instead of system IO,
-    allowing for asynchronous communication through a web interface. It maintains
-    separate buffers for stdin and stdout to simulate bidirectional communication.
+    Thread-safe WebIOBuffer that maintains the synchronous IOBuffer interface.
+    Uses threading primitives instead of asyncio for synchronization.
     """
-    
     def __init__(self):
-        """Initialize empty input and output buffers."""
         self._stdin_buffer = ""
         self._stdout_buffer = ""
+        self._lock = Lock()
         
+    def append_stdin(self, content: str) -> None:
+        """Thread-safe append to stdin buffer."""
+        if not content:
+            return
+            
+        with self._lock:
+            self._stdin_buffer += content
+    
     def read(self) -> str:
-        """
-        Read and clear all available input from the stdin buffer.
-        
-        Returns:
-            str: Content from the stdin buffer, or empty string if buffer is empty
-        """
-        content = self._stdin_buffer
-        self._stdin_buffer = ""
-        return content
-        
+        """Thread-safe read from stdin buffer."""
+        with self._lock:
+            content = self._stdin_buffer
+            self._stdin_buffer = ""
+            return content
+            
     def write(self, content: str) -> None:
-        """
-        Append content to the stdout buffer.
-        
-        Args:
-            content: String content to append to the buffer
-        """
-        if content:
+        """Thread-safe write to stdout buffer."""
+        if not content:
+            return
+            
+        with self._lock:
             self._stdout_buffer += content
             
     def buffer_length(self) -> int:
@@ -55,14 +55,16 @@ class WebIOBuffer(IOBuffer):
             self._stdin_buffer += content
             
     def get_stdout(self) -> str:
-        """
-        Get the current content of the stdout buffer.
-        
-        Returns:
-            str: Current content of the stdout buffer
-        """
-        return self._stdout_buffer
-        
+        """Thread-safe get stdout content."""
+        with self._lock:
+            return self._stdout_buffer
+            
     def clear_stdout(self) -> None:
-        """Clear the stdout buffer."""
-        self._stdout_buffer = ""
+        """Thread-safe clear stdout buffer."""
+        with self._lock:
+            self._stdout_buffer = ""
+            
+    def buffer_length(self) -> int:
+        """Thread-safe get stdin buffer length."""
+        with self._lock:
+            return len(self._stdin_buffer)

sia/web_server.py

@@ -5,6 +5,7 @@ import asyncio
 import json
 from dataclasses import dataclass
 import mimetypes
+from functools import partial
 
 from .web_agent import WebAgent, WebAgentState
 from .web_io_buffer import WebIOBuffer
@@ -37,7 +38,7 @@ class WebServer:
     integrates with the WebAgent. Serves static files from a web directory.
     """
     def __init__(self, agent: WebAgent, io_buffer: WebIOBuffer, 
-                 host: str = "0.0.0.0", port: int = 8080, 
+                 host: str = "0.0.0.0", port: int = 8080,
                  static_dir: Optional[Path] = None):
         """Initialize the web server."""
         self.agent = agent
@@ -49,11 +50,22 @@ class WebServer:
         self._init_routes()
         self._clients: Set[web.WebSocketResponse] = set()
         self._lock = asyncio.Lock()
+        self._response_tasks: List[asyncio.Task] = []  # Initialize tasks list
         self.runner: Optional[web.AppRunner] = None
         
-        # Register agent handlers
-        self.agent.add_state_change_handler(self._handle_state_change)
-        self.agent.add_response_change_handler(self._handle_response_change)
+        # Register wrapped agent handlers
+        self.agent.add_state_change_handler(self._wrap_async_handler(self._handle_state_change))
+        self.agent.add_response_change_handler(self._wrap_async_handler(self._handle_response_change))
+
+    def _wrap_async_handler(self, coro_func):
+        """Wrap an async handler function to be called synchronously."""
+        def wrapper(*args, **kwargs):
+            loop = asyncio.get_event_loop()
+            task = loop.create_task(coro_func(*args, **kwargs))
+            self._response_tasks.append(task)
+            self._response_tasks = [t for t in self._response_tasks if not t.done()]
+            return task
+        return wrapper
 
     def _init_routes(self):
         """Initialize application routes."""
@@ -61,16 +73,44 @@ class WebServer:
         self.app.router.add_get("/ws", self._handle_websocket)
 
         if self.static_dir.exists():
-                self.app.router.add_get("/", self._serve_index)
-                self.app.router.add_static("/static/", self.static_dir, show_index=False)
-                self.app.router.add_static("/assets/", self.static_dir / "assets", show_index=False)
-                self.app.router.add_get("/{path:.*}", self._serve_index)
+            self.app.router.add_get("/", self._serve_index)
+            self.app.router.add_static("/static/", self.static_dir, show_index=False)
+            self.app.router.add_static("/assets/", self.static_dir / "assets", show_index=False)
+            self.app.router.add_get("/{path:.*}", self._serve_index)
+
+    async def _handle_client_message(self, data: dict):
+        """Handle incoming client message."""
+        message_type = data.get("type")
+        
+        if message_type == ClientMessage.APPROVE_CONTEXT:
+            if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
+                self.agent.approve_context()
+                
+        elif message_type == ClientMessage.APPROVE_RESPONSE:
+            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
+                self.agent.approve_response()
+                
+        elif message_type == ClientMessage.MODIFY_RESPONSE:
+            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
+                new_response = data.get("data")
+                if new_response:
+                    self.agent._set_response(new_response)
+                    
+        elif message_type == ClientMessage.SEND_INPUT:
+            input_text = data.get("data")
+            if input_text:
+                if not input_text.endswith('\n'):
+                    input_text += '\n'
+                
+                self.io_buffer.append_stdin(input_text)
+                
+                current_length = self.io_buffer.buffer_length()
+                return
+        else:
+            print(f"Agent in wrong state: {self.agent.current_state}")
 
     async def _serve_index(self, request: web.Request) -> web.Response:
-        """
-        Serve the React application HTML for any unmatched routes to support
-        client-side routing.
-        """
+        """Serve the React application HTML for any unmatched routes."""
         index_path = self.static_dir / "index.html"
         if not index_path.exists():
             raise web.HTTPNotFound()
@@ -129,7 +169,8 @@ class WebServer:
             try:
                 await ws.send_json(message)
                 return True
-            except Exception:
+            except Exception as e:
+                print(f"Error sending to client: {e}")
                 return False
         return False
             
@@ -216,27 +257,28 @@ class WebServer:
         return ws
 
     async def _handle_client_message(self, data: dict):
-        """Handle incoming client message."""
         message_type = data.get("type")
-        
+    
         if message_type == ClientMessage.APPROVE_CONTEXT:
             if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
-                await self.agent.approve_context()
-                
+                await self.agent.approve_context()                
         elif message_type == ClientMessage.APPROVE_RESPONSE:
             if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
-                await self.agent.approve_response()
+                self.agent.approve_response()  # Now synchronous
                 
         elif message_type == ClientMessage.MODIFY_RESPONSE:
             if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
                 new_response = data.get("data")
                 if new_response:
-                    await self.agent._set_response(new_response)
+                    self.agent._set_response(new_response)
                     
         elif message_type == ClientMessage.SEND_INPUT:
             input_text = data.get("data")
             if input_text:
+                self.agent.set_input(input_text)
                 self.io_buffer.append_stdin(input_text + "\n")
+        else:
+            print(f"Agent in wrong state: {self.agent.current_state}")
         
     async def _handle_state_change(self, new_state: WebAgentState):
         """Handle state changes from the WebAgent."""
@@ -249,3 +291,27 @@ class WebServer:
             response,
             validation_error=self.agent._validation_error
         )
+
+    async def broadcast(self, message_type: str, data: str, **kwargs):
+        """Broadcast a message to all connected clients."""
+        async with self._lock:
+            message = {
+                "type": message_type,
+                "data": data,
+                **kwargs
+            }
+            
+            tasks = []
+            disconnected = set()
+            
+            for ws in self._clients:
+                if not ws.closed:
+                    try:
+                        await self._send_to_client(ws, message)
+                    except Exception as e:
+                        print(f"Error sending to client: {e}")
+                        disconnected.add(ws)
+                else:
+                    disconnected.add(ws)
+                    
+            self._clients -= disconnected

test/mock_web_agent.py

@@ -1,4 +1,3 @@
-import asyncio
 from typing import Optional, List, Callable
 from sia.web_agent import WebAgentState
 
@@ -21,35 +20,34 @@ class MockWebAgent:
         
     def add_response_change_handler(self, handler):
         self._response_handlers.append(handler)
+
+    def set_input(self, input_text):
+        pass
         
-    async def _notify_state_handlers(self, new_state: WebAgentState):
+    def _notify_state_handlers(self, new_state: WebAgentState):
         """Notify all state handlers of a state change."""
         for handler in self._state_handlers:
-            await handler(new_state)
+            handler(new_state)
             
-    async def _notify_response_handlers(self, response: str):
+    def _notify_response_handlers(self, response: str):
         """Notify all response handlers of a response change."""
         for handler in self._response_handlers:
-            await handler(response)
+            handler(response)
             
-    async def _set_response(self, response: str):
+    def _set_response(self, response: str):
         """Set response and notify handlers."""
         self.response = response
-        await self._notify_response_handlers(response)
+        self._notify_response_handlers(response)
             
-    async def approve_context(self):
+    def approve_context(self):
         """Handle context approval with immediate state updates."""
         # Change to INFERENCE state and notify
         self._current_state = WebAgentState.INFERENCE
-        for handler in self._state_handlers:
-            await handler(WebAgentState.INFERENCE)
+        self._notify_state_handlers(WebAgentState.INFERENCE)
         
         # Set response and notify
-        self.response = "<reasoning>test</reasoning>"
-        for handler in self._response_handlers:
-            await handler(self.response)
+        self._set_response("<reasoning>test</reasoning>")
         
         # Change to RESPONSE_APPROVAL state and notify
         self._current_state = WebAgentState.RESPONSE_APPROVAL
-        for handler in self._state_handlers:
-            await handler(WebAgentState.RESPONSE_APPROVAL)
+        self._notify_state_handlers(WebAgentState.RESPONSE_APPROVAL)

test/standard_io_buffer_test.py

@@ -159,11 +159,6 @@ if __name__ == '__main__':
         stdout, stderr, code = self.run_script('read', special_chars)
         self.assertEqual(code, 0, f"Script failed with stderr: {stderr}")
         expected = f'Read: {special_chars}'
-        # Print actual bytes for debugging
-        if stdout != expected:
-            print("\nDebug output:")
-            print("Expected bytes:", expected.encode('utf-8'))
-            print("Actual bytes:", stdout.encode('utf-8'))
         self.assertEqual(stdout, expected)
 
     def test_empty_input(self):

test/web_agent_test.py

@@ -1,3 +1,4 @@
+import asyncio
 import unittest
 from datetime import datetime
 from unittest.mock import Mock, MagicMock, patch, call
@@ -123,7 +124,7 @@ class WebAgentTest(unittest.TestCase):
         self.agent._current_state = WebAgentState.UPDATE
         
         with self.assertRaises(Exception) as context:
-            self.agent.approve_context()
+            asyncio.run(self.agent.approve_context())
         self.assertIn("Not in CONTEXT_APPROVAL state", str(context.exception))
         
     def test_approve_response_state_error(self):
@@ -132,7 +133,7 @@ class WebAgentTest(unittest.TestCase):
             self.agent.approve_response()
         self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
         
-    def test_context_approval_flow(self):
+    async def test_context_approval_flow(self):
         """Test complete context approval flow with state transitions."""
         # Register handlers
         self.agent.add_state_change_handler(self.state_change_handler)
@@ -140,6 +141,7 @@ class WebAgentTest(unittest.TestCase):
         
         # Setup LLM response
         def mock_infer(prompt: str, context: str) -> Iterator[str]:
+            # Return complete response at once for testing
             yield "<reasoning>test reasoning</reasoning>"
         self.mock_llm.infer.side_effect = mock_infer
         
@@ -149,6 +151,9 @@ class WebAgentTest(unittest.TestCase):
         # Approve context
         self.agent.approve_context()
         
+        # Wait for response tasks to complete
+        await asyncio.gather(*self.web_server._response_tasks)
+        
         # Verify state transitions
         expected_states = [
             WebAgentState.INFERENCE,
@@ -156,10 +161,11 @@ class WebAgentTest(unittest.TestCase):
         ]
         self.assertEqual(self.state_changes, expected_states)
         
-        # Verify response updates
-        self.assertEqual(len(self.response_changes), 2)  # Empty + final response
-        self.assertEqual(self.response_changes[0], "")
-        self.assertEqual(self.response_changes[1], "<reasoning>test reasoning</reasoning>")
+        # Updated test: Should now expect streaming updates
+        self.assertEqual(len(self.response_changes), len(self.received_messages))
+        for i, msg in enumerate(self.received_messages):
+            self.assertEqual(msg["type"], "RESPONSE_UPDATE")
+            self.assertEqual(msg["data"], self.response_changes[i])
         
     def test_response_approval_flow_command(self):
         """Test response approval flow with command execution."""
@@ -200,7 +206,7 @@ class WebAgentTest(unittest.TestCase):
         
         # Set initial state and response
         self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
-        self.agent._response = "<reasoning>test reasoning</reasoning>"
+        self.agent.response = "<reasoning>test reasoning</reasoning>"
         
         # Approve response
         self.agent.approve_response()
@@ -218,7 +224,7 @@ class WebAgentTest(unittest.TestCase):
         ]
         self.assertEqual(self.state_changes, expected_states)
         
-    def test_response_validation(self):
+    async def test_response_validation(self):
         """Test response validation during response setting."""
         # Register handlers
         self.agent.add_response_change_handler(self.response_change_handler)
@@ -227,8 +233,8 @@ class WebAgentTest(unittest.TestCase):
         error_message = "Invalid XML"
         self.mock_validator.validate.return_value = error_message
         
-        # Set response
-        self.agent._set_response("<invalid>")
+        # Set response using async method
+        await self.agent._set_response("<invalid>")
         
         # Verify validation error stored
         self.assertEqual(self.agent._validation_error, error_message)

test/web_server_test.py

@@ -6,6 +6,7 @@ from sia.web_agent import WebAgentState
 from unittest.mock import patch
 import asyncio
 import asyncio
+import time
 import json
 import json
 import shutil
@@ -29,6 +30,7 @@ class WebServerTest(AioHTTPTestCase):
         self.static_dir.mkdir(exist_ok=True)
         
         # Create web server with temp static directory
+        self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
         self.web_server = WebServer(
             self.agent,
             self.io_buffer,
@@ -54,7 +56,6 @@ class WebServerTest(AioHTTPTestCase):
                 elif msg.type == WSMsgType.BINARY:
                     messages.append(json.loads(msg.data.decode()))
                 elif msg.type == WSMsgType.CLOSED:
-                    print("WebSocket closed normally")
                     break
                 elif msg.type == WSMsgType.ERROR:
                     self.fail(f"WebSocket error while waiting for messages: {msg.data}")
@@ -108,6 +109,7 @@ class WebServerTest(AioHTTPTestCase):
             await ws.send_json({
                 "type": ClientMessage.APPROVE_CONTEXT,
             })
+            await asyncio.sleep(0.1)  # Let state changes propagate
             
             # Should receive 3 messages: INFERENCE state, response, RESPONSE_APPROVAL state
             messages = await self.wait_for_messages(ws, count=3)
@@ -124,7 +126,7 @@ class WebServerTest(AioHTTPTestCase):
             self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE)
             self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name)
 
-    @unittest_run_loop        
+    @unittest_run_loop
     async def test_validation_error(self):
         """Test handling of validation errors."""
         async with self.client.ws_connect("/ws") as ws:
@@ -134,7 +136,7 @@ class WebServerTest(AioHTTPTestCase):
             
             # Set validation error and response
             self.agent._validation_error = "Invalid XML"
-            await self.agent._set_response("<invalid>")
+            self.agent._set_response("<invalid>")  # Now synchronous
             await asyncio.sleep(0.1)  # Let changes propagate
             
             # Should receive response update with error
@@ -204,24 +206,34 @@ class WebServerTest(AioHTTPTestCase):
             await ws1.close()
             await ws2.close()
             
-    @unittest_run_loop
+    async def wait_for_buffer_content(self, timeout=1.0):
+        start = time.time()
+        while time.time() - start < timeout:
+            if self.io_buffer.buffer_length() > 0:
+                return True
+            await asyncio.sleep(0.01)
+        return False
+
     async def test_send_input(self):
-        """Test sending input to agent."""
         async with self.client.ws_connect("/ws") as ws:
-            # Drain initial messages
             await self.drain_messages(ws)
-            await asyncio.sleep(0.1)  # Let connection settle
-            
-            # Send input
+            await asyncio.sleep(0.1)
+
             test_input = "test input"
             await ws.send_json({
                 "type": ClientMessage.SEND_INPUT,
                 "data": test_input
             })
-            await asyncio.sleep(0.1)  # Let input process
+
+            # Wait for buffer to have content
+            self.assertTrue(await self.wait_for_buffer_content())
             
-            # Verify input added to buffer with newline
-            self.assertEqual(self.io_buffer.read(), test_input + "\n")
+            content = self.io_buffer.read()
+            self.assertEqual(
+                content, 
+                test_input + "\n",
+                f"Expected '{test_input}\\n' but got '{content}'"
+            )
 
     @unittest_run_loop
     async def test_invalid_json(self):

web/src/components/App.jsx

@@ -1,12 +1,10 @@
 import React, { useState, useEffect, useRef } from 'react';
 import Editor from '@monaco-editor/react';
-import { Card, CardContent, CardHeader, CardTitle } from '../components/ui/card.jsx';
-import { Button } from '../components/ui/button.jsx';
-import { ScrollArea } from '../components/ui/scroll-area.jsx';
-import { Alert, AlertDescription } from '../components/ui/alert.jsx';
-import { Tabs, TabsContent, TabsList, TabsTrigger } from '../components/ui/tabs.jsx';
-import { Textarea } from '../components/ui/textarea.jsx';
-import { Input } from '../components/ui/input.jsx';
+import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
+import { Button } from '@/components/ui/button';
+import { ScrollArea } from '@/components/ui/scroll-area';
+import { Alert, AlertDescription } from '@/components/ui/alert';
+import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
 
 const WebSocketState = {
   CONNECTING: 'CONNECTING',
@@ -29,6 +27,13 @@ const MessageType = {
   VALIDATION_ERROR: 'VALIDATION_ERROR',
 };
 
+const ClientMessageType = {
+  APPROVE_CONTEXT: 'APPROVE_CONTEXT',
+  APPROVE_RESPONSE: 'APPROVE_RESPONSE',
+  MODIFY_RESPONSE: 'MODIFY_RESPONSE',
+  SEND_INPUT: 'SEND_INPUT',
+};
+
 const SIAInterface = () => {
   // Editor content state
   const [context, setContext] = useState('');
@@ -40,7 +45,7 @@ const SIAInterface = () => {
 
   // UI state
   const [activeTab, setActiveTab] = useState('input');
-  const [agentState, setAgentState] = useState(AgentState.UPDATE);
+  const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
   const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
 
   // WebSocket ref
@@ -58,10 +63,7 @@ const SIAInterface = () => {
   };
 
   useEffect(() => {
-    // Connect to WebSocket
     connectWebSocket();
-
-    // Cleanup on unmount
     return () => {
       if (wsRef.current) {
         wsRef.current.close();
@@ -80,10 +82,13 @@ const SIAInterface = () => {
 
     ws.onclose = () => {
       setWsState(WebSocketState.DISCONNECTED);
-      // Attempt to reconnect after delay
       setTimeout(connectWebSocket, 2000);
     };
 
+    ws.onerror = (error) => {
+      console.error('WebSocket error:', error);
+    };
+
     ws.onmessage = (event) => {
       try {
         const message = JSON.parse(event.data);
@@ -113,21 +118,26 @@ const SIAInterface = () => {
     }
   };
 
-  const sendWebSocketMessage = (type, data) => {
+  const sendWebSocketMessage = (type, data = null) => {
     if (wsRef.current?.readyState === WebSocket.OPEN) {
-      wsRef.current.send(JSON.stringify({ type, data }));
+      const message = { 
+        type: type,  // Ensure exact string match
+        data: data || undefined // Only include if not null
+      };
+      console.log('Sending message:', JSON.stringify(message)); // Log exact message being sent
+      wsRef.current.send(JSON.stringify(message));
+    } else {
+      console.error('WebSocket not connected, state:', wsRef.current?.readyState);
     }
   };
 
   const handleSendInput = (withStdin = false) => {
     if (!input.trim()) return;
 
-    // Send input to backend
-    sendWebSocketMessage('SEND_INPUT', input);
+    sendWebSocketMessage(ClientMessageType.SEND_INPUT, input);
     setInput('');
 
     if (withStdin) {
-      // Prepare read_stdin response
       const readStdinXml = `<read_stdin>
         <!-- Content will be populated with stdin -->
       </read_stdin>`;
@@ -139,12 +149,26 @@ const SIAInterface = () => {
   };
 
   const handleApproveContext = () => {
-    sendWebSocketMessage('APPROVE_CONTEXT');
+    console.log('Current state:', agentState);
+    console.log('WebSocket state:', wsRef.current?.readyState);
+    if (agentState === AgentState.CONTEXT_APPROVAL && 
+        wsRef.current?.readyState === WebSocket.OPEN) {
+      console.log('Sending approve context message');
+      sendWebSocketMessage(ClientMessageType.APPROVE_CONTEXT);
+    } else {
+      console.warn(
+        'Cannot approve context.',
+        'Agent state:', agentState,
+        'WebSocket state:', wsRef.current?.readyState
+      );
+    }
   };
 
   const handleApproveResponse = (modified = false) => {
-    const response = modified ? modifiedResponse : originalResponse;
-    sendWebSocketMessage('APPROVE_RESPONSE', response);
+    if (agentState === AgentState.RESPONSE_APPROVAL) {
+      const response = modified ? modifiedResponse : originalResponse;
+      sendWebSocketMessage(ClientMessageType.APPROVE_RESPONSE, response);
+    }
   };
 
   const handleResetModified = () => {
@@ -190,7 +214,7 @@ const SIAInterface = () => {
             <div className="mt-4">
               <Button 
                 onClick={handleApproveContext} 
-                disabled={agentState !== AgentState.CONTEXT_APPROVAL}
+                disabled={agentState !== AgentState.CONTEXT_APPROVAL || wsState !== WebSocketState.CONNECTED}
                 className="w-full"
               >
                 Approve Context
@@ -223,6 +247,7 @@ const SIAInterface = () => {
                   <div className="flex space-x-2">
                     <Button 
                       onClick={() => handleSendInput(false)} 
+                      disabled={wsState !== WebSocketState.CONNECTED}
                       className="flex-1"
                     >
                       Send Input
@@ -230,6 +255,7 @@ const SIAInterface = () => {
                     <Button 
                       onClick={() => handleSendInput(true)} 
                       variant="secondary"
+                      disabled={wsState !== WebSocketState.CONNECTED}
                       className="flex-1"
                     >
                       Send & Read Stdin
@@ -259,7 +285,11 @@ const SIAInterface = () => {
                   )}
                   <Button 
                     onClick={() => handleApproveResponse(false)}
-                    disabled={agentState !== AgentState.RESPONSE_APPROVAL || validationError}
+                    disabled={
+                      agentState !== AgentState.RESPONSE_APPROVAL || 
+                      validationError || 
+                      wsState !== WebSocketState.CONNECTED
+                    }
                     className="w-full"
                   >
                     Approve Original Response
@@ -297,7 +327,11 @@ const SIAInterface = () => {
                     </Button>
                     <Button 
                       onClick={() => handleApproveResponse(true)}
-                      disabled={agentState !== AgentState.RESPONSE_APPROVAL || validationError}
+                      disabled={
+                        agentState !== AgentState.RESPONSE_APPROVAL || 
+                        validationError || 
+                        wsState !== WebSocketState.CONNECTED
+                      }
                       className="flex-1"
                     >
                       Approve Modified Response
@@ -334,4 +368,4 @@ const SIAInterface = () => {
   );
 };
 
-export default SIAInterface;
+export default SIAInterface;