Separate and tested web_socket_manager

Dockerfile

@@ -11,7 +11,7 @@ FROM requirements AS sia-test
 COPY ./ /root/sia/
 WORKDIR /root/sia/
 RUN mkdir -p /root/model
-CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py", "-v"]
+CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py"]
 
 FROM node:20-alpine AS web-test
 WORKDIR /app

readme.md

@@ -414,34 +414,32 @@ classDiagram
 
     class WebAgent {
         +context: str
-        +response: str
+        +response: str readonly
         +current_state WebAgentState readonly
         +command_result Optional[CommandResult] readonly
+        +validation_error Optional[str] readonly
 
         +add_state_change_handler(handler Callable) void
         +add_response_change_handler(handler Callable) void
         +approve_context() void
+        +set__response(response str) void
         +approve_response() void
     }
 
     class WebAgentState {
+
         <<enumeration>>
         UPDATE
         CONTEXT_APPROVAL
         INFERENCE
         RESPONSE_APPROVAL
+        STOPPED
     }
 
-    class WebServer {
+    class WebSocketManager {
-        -agent: WebAgent
+        -web_sockets: Set~WebSocket~
-        -app: Application
-        -clients: Set~WebSocketResponse~
-        +clients Set~WebSocketResponse~ readonly
 
-        +WebServer(agent WebAgent, io_buffer WebIOBuffer, host str, port int)
+        +WebServer(agent WebAgent, io_buffer WebIOBuffer, static_files path, host str, port int)
-        +start() void
-        +stop() void
-        +broadcast(message_type str, data str, **kwargs) void
     }
 
     class ClientMessage {

sia/__main__.py

@@ -1,65 +1,105 @@
-import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from aiohttp import web
 from pathlib import Path
+import asyncio
+import mimetypes
+import time
 
 from .llm_engine import LlmEngine
 from .system_metrics import SystemMetrics
 from .web_agent import WebAgent
 from .web_io_buffer import WebIOBuffer
-from .web_server import WebServer
+from .web_socket_manager import WebSocketManager
 from .working_memory import WorkingMemory
 from .xml_validator import XMLValidator
 from .response_parser import ResponseParser
+from .web_agent import WebAgent
+from .web_io_buffer import WebIOBuffer
 
-async def run_web_server():
+mimetypes.add_type("application/javascript", ".js")
-    """Run the web server with default configuration."""
+mimetypes.add_type("application/javascript", ".jsx")
-    base_dir = Path(__file__).parent.parent
+mimetypes.add_type("text/javascript", ".js")
+mimetypes.add_type("text/javascript", ".jsx")
 
-    # Load system prompt and action schema
+class TestLLM:
-    system_prompt = (base_dir / "system_prompt.md").read_text()
+    def infer(self, prompt: str, context: str):
-    action_schema = (base_dir / "action_schema.xsd").read_text()
+        yield "<reasoning>"
+        time.sleep(2)
+        yield "test reasoning"
+        time.sleep(2)
+        yield "</reasoning>"
 
-    # Initialize core components
+class Main:
-    llm = LlmEngine("/root/model")
+    def __init__(self):
-    io_buffer = WebIOBuffer()
+        self._base_dir = Path(__file__).parent.parent
-    agent = WebAgent(
+        self._system_prompt = (self._base_dir / "system_prompt.md").read_text()
-        system_prompt=system_prompt,
+        self._action_schema = (self._base_dir / "action_schema.xsd").read_text()
-        action_schema=action_schema,
+        self._static_dir = self._base_dir / "static"
-        working_memory=WorkingMemory(),
-        system_metrics=SystemMetrics(),
-        llm=llm,
-        validator=XMLValidator(action_schema),
-        parser=ResponseParser(io_buffer)
-    )
 
-    # Start the web server
+        self._llm = LlmEngine("/root/model")
-    server = WebServer(agent, io_buffer, "0.0.0.0", 8080)
+        #self._llm = TestLLM()
-    await server.start()
+        self._io_buffer = WebIOBuffer()
+        self._agent = WebAgent(
+            system_prompt=self._system_prompt,
+            action_schema=self._action_schema,
+            working_memory=WorkingMemory(),
+            system_metrics=SystemMetrics(),
+            llm=self._llm,
+            validator=XMLValidator(self._action_schema),
+            parser=ResponseParser(self._io_buffer)
+        )
+        self._ws_manager = WebSocketManager(
+            self._agent,
+            self._io_buffer
+        )
 
-    try:
+        self._app = web.Application()
-        while True:
+        self._init_routes()
-            if server._response_tasks:
-                done, _ = await asyncio.wait(
-                    server._response_tasks,
-                    timeout=0.1,
-                    return_when=asyncio.FIRST_COMPLETED
-                )
-                for task in done:
-                    try:
-                        await task
-                    except asyncio.CancelledError:
-                        print("Task was cancelled")
-                    except Exception as e:
-                        print(f"Task error: {e}")
-                    server._response_tasks.discard(task)  # Use discard to prevent KeyError
-            await asyncio.sleep(0.1)
-    except KeyboardInterrupt:
-        print("\nShutting down...")
-    finally:
-        await server.stop()
 
-def main():
+    async def app(self):
-    """Main entry point for the SIA application."""
+        return self._app
-    asyncio.run(run_web_server())
+
+    def _init_routes(self):
+        """Initialize application routes."""
+        self._app.middlewares.append(self._cors_middleware)
+        self._app.router.add_get("/ws", self._ws_manager.handle_websocket)
+        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 _serve_index(self, request: web.Request) -> web.Response:
+        """Serve the React application HTML for any unmatched routes."""
+        index_path = self._static_dir / "index.html"
+        if not index_path.exists():
+            raise web.HTTPNotFound()
+
+        with open(index_path, "r") as f:
+            html_content = f.read()
+
+        return web.Response(
+            text=html_content,
+            content_type="text/html"
+        )
+
+    @web.middleware
+    async def _cors_middleware(self, request: web.Request, handler):
+        """Handle CORS headers."""
+        if request.method == "OPTIONS":
+            response = web.Response()
+        else:
+            response = await handler(request)
+
+        response.headers.update({
+            'Access-Control-Allow-Origin': '*',
+            'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
+            'Access-Control-Allow-Headers': 'Content-Type',
+            'Access-Control-Max-Age': '86400',
+        })
+        return response
 
 if __name__ == "__main__":
-    main()
+    main = Main()
+    app = asyncio.run(main.app())
+    print("Web server started at http://localhost:8080")
+    web.run_app(app, port=8080)

sia/web_agent.py

@@ -24,6 +24,7 @@ class WebAgentState(Enum):
     CONTEXT_APPROVAL = auto()
     INFERENCE = auto()
     RESPONSE_APPROVAL = auto()
+    STOPPED = auto()
 
 class WebAgent(BaseAgent):
     """
@@ -45,12 +46,12 @@ class WebAgent(BaseAgent):
         Initialize web agent with required components.
         """
         super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser)
-        self.response = ""
+        self._response = ""
         self._system_prompt = system_prompt
         self._current_state = WebAgentState.CONTEXT_APPROVAL
         self._validation_error: Optional[str] = None
         self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
-        self._response_change_handlers: List[Callable[[str], None]] = []
+        self._response_change_handlers: List[Callable[[str, str], None]] = []
         self._command_result: Optional[CommandResult] = None
         self.context = self._compile_context()
 
@@ -63,6 +64,16 @@ class WebAgent(BaseAgent):
     def command_result(self) -> Optional[CommandResult]:
         """Get the result of the last command execution."""
         return self._command_result
+
+    @property
+    def validation_error(self) -> Optional[str]:
+        """Get the validation error, if any."""
+        return self._validation_error
+    
+    @property
+    def response(self) -> str:
+        """Get the current response."""
+        return self._response
         
     def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
         """
@@ -74,7 +85,7 @@ class WebAgent(BaseAgent):
         if handler not in self._state_change_handlers:
             self._state_change_handlers.append(handler)
     
-    def add_response_change_handler(self, handler: Callable[[str], None]) -> None:
+    def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None:
         """
         Add a callback for response changes.
 
@@ -89,20 +100,13 @@ class WebAgent(BaseAgent):
         for handler in self._state_change_handlers:
             handler(self._current_state)
     
-    def _notify_response_change(self) -> None:
+    def set_response(self, response: str) -> None:
-        """Notify all handlers of response change."""
-        for handler in self._response_change_handlers:
-            handler(self.response)
-    
-    def _set_response(self, response: str) -> None:
         """Set response and notify handlers."""
-        self.response = response
+        self._response = response
-        self._validation_error = self._validator.validate(response)
+        validation_error = self._validator.validate(response)
+        self._validation_error = validation_error
         for handler in self._response_change_handlers:
-            try:
+            handler(response, validation_error)
-                handler(response)
-            except Exception as e:
-                print(f"Error in response handler: {e}")
 
     def _set_state(self, new_state: WebAgentState) -> None:
         """
@@ -114,39 +118,22 @@ class WebAgent(BaseAgent):
         self._current_state = new_state
         self._notify_state_change()
             
-    async def approve_context(self) -> None:
+    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("")
+        self.set_response("")
         
         response_token_iter = self._llm.infer(self._system_prompt, self.context)
         response = ""
-        try:
+        for token in response_token_iter:
-            for token in response_token_iter:
+            response += token
-                response += token
+            self.set_response(response)
-                await self._set_response(response)
+            print(f"{token}")
-                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")
@@ -157,10 +144,12 @@ class WebAgent(BaseAgent):
         if isinstance(parse_result, Command):
             result = parse_result.execute(self._working_memory)
             self._command_result = result
+            if result.should_stop:
+                self._set_state(WebAgentState.STOPPED)
+                return
         else:
             self._working_memory.add_entry(parse_result)
         
         self._working_memory.update()
         self.context = self._compile_context()
-
         self._set_state(WebAgentState.CONTEXT_APPROVAL)

sia/web_io_buffer.py

@@ -10,30 +10,16 @@ class WebIOBuffer(IOBuffer):
     def __init__(self):
         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:
         """Thread-safe read from stdin buffer."""
-        with self._lock:
+        content = self._stdin_buffer
-            content = self._stdin_buffer
+        self._stdin_buffer = ""
-            self._stdin_buffer = ""
+        return content
-            return content
             
     def write(self, content: str) -> None:
         """Thread-safe write to stdout buffer."""
-        if not content:
+        self._stdout_buffer += content
-            return
-            
-        with self._lock:
-            self._stdout_buffer += content
             
     def buffer_length(self) -> int:
         """
@@ -51,20 +37,12 @@ class WebIOBuffer(IOBuffer):
         Args:
             content: String content to append to the stdin buffer
         """
-        if content:
+        self._stdin_buffer += content
-            self._stdin_buffer += content
             
     def get_stdout(self) -> str:
         """Thread-safe get stdout content."""
-        with self._lock:
+        return self._stdout_buffer
-            return self._stdout_buffer
             
     def clear_stdout(self) -> None:
         """Thread-safe clear stdout buffer."""
-        with self._lock:
+        self._stdout_buffer = ""
-            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

@@ -1,317 +0,0 @@
-from aiohttp import web, WSMsgType
-from pathlib import Path
-from typing import Set, Optional
-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
-
-mimetypes.add_type("application/javascript", ".js")
-mimetypes.add_type("application/javascript", ".jsx")
-mimetypes.add_type("text/javascript", ".js")
-mimetypes.add_type("text/javascript", ".jsx")
-
-@dataclass
-class ClientMessage:
-    """Messages that can be received from clients."""
-    APPROVE_CONTEXT = "APPROVE_CONTEXT"
-    APPROVE_RESPONSE = "APPROVE_RESPONSE" 
-    MODIFY_RESPONSE = "MODIFY_RESPONSE"
-    SEND_INPUT = "SEND_INPUT"
-
-@dataclass
-class ServerMessage:
-    """Messages that can be sent to clients."""
-    STATE_CHANGE = "STATE_CHANGE"
-    CONTEXT_UPDATE = "CONTEXT_UPDATE"
-    RESPONSE_UPDATE = "RESPONSE_UPDATE"
-    OUTPUT_UPDATE = "OUTPUT_UPDATE"
-    VALIDATION_ERROR = "VALIDATION_ERROR"
-
-class WebServer:
-    """
-    AIOHTTP-based web server that manages WebSocket connections and 
-    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,
-                 static_dir: Optional[Path] = None):
-        """Initialize the web server."""
-        self.agent = agent
-        self.io_buffer = io_buffer
-        self.host = host
-        self.port = port
-        self.app = web.Application()
-        self.static_dir = static_dir or Path(__file__).parent.parent / "static"
-        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 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."""
-        self.app.middlewares.append(self._cors_middleware)
-        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)
-
-    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."""
-        index_path = self.static_dir / "index.html"
-        if not index_path.exists():
-            raise web.HTTPNotFound()
-
-        with open(index_path, "r") as f:
-            html_content = f.read()
-
-        return web.Response(
-            text=html_content,
-            content_type="text/html"
-        )
-
-    @web.middleware
-    async def _cors_middleware(self, request: web.Request, handler):
-        """Handle CORS headers."""
-        if request.method == "OPTIONS":
-            response = web.Response()
-        else:
-            response = await handler(request)
-
-        response.headers.update({
-            'Access-Control-Allow-Origin': '*',
-            'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
-            'Access-Control-Allow-Headers': 'Content-Type',
-            'Access-Control-Max-Age': '86400',
-        })
-        return response
-
-    @property
-    def clients(self) -> Set[web.WebSocketResponse]:
-        """Get the set of connected clients."""
-        return self._clients
-
-    async def start(self):
-        """Start the web server."""
-        # Verify static directory exists
-        if not self.static_dir.exists():
-            raise FileNotFoundError(f"Static directory not found: {self.static_dir}")
-
-        self.runner = web.AppRunner(self.app)
-        await self.runner.setup()
-        site = web.TCPSite(self.runner, self.host, self.port)
-        await site.start()
-        print(f"Web server started at http://{self.host}:{self.port}")
-        print(f"Serving static files from: {self.static_dir}")
-        print(f"WebSocket endpoint at ws://{self.host}:{self.port}/ws")
-
-    async def stop(self):
-        """Stop the web server and clean up resources."""
-        if self.runner:
-            await self.runner.cleanup()
-
-    async def _send_to_client(self, ws: web.WebSocketResponse, message: dict) -> bool:
-        """Send message to client, return True if successful."""
-        if not ws.closed:
-            try:
-                await ws.send_json(message)
-                return True
-            except Exception as e:
-                print(f"Error sending to client: {e}")
-                return False
-        return False
-            
-    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:
-                    task = asyncio.create_task(self._send_to_client(ws, message))
-                    tasks.append((ws, task))
-                else:
-                    disconnected.add(ws)
-                    
-            if tasks:
-                for ws, task in tasks:
-                    try:
-                        success = await task
-                        if not success:
-                            disconnected.add(ws)
-                    except Exception:
-                        disconnected.add(ws)
-                    
-            self._clients -= disconnected
-
-    async def _handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
-        """Handle new WebSocket connections and messages."""
-        ws = web.WebSocketResponse(heartbeat=30)
-        await ws.prepare(request)
-        
-        async with self._lock:
-            self._clients.add(ws)
-        
-        try:
-            # Send initial state
-            await self._send_to_client(ws, {
-                "type": ServerMessage.STATE_CHANGE,
-                "data": self.agent.current_state.name if self.agent.current_state else "INITIAL"
-            })
-            
-            # Send current context if available
-            if self.agent.context:
-                await self._send_to_client(ws, {
-                    "type": ServerMessage.CONTEXT_UPDATE,
-                    "data": self.agent.context
-                })
-                
-            # Send current response if available
-            if self.agent.response:
-                await self._send_to_client(ws, {
-                    "type": ServerMessage.RESPONSE_UPDATE,
-                    "data": self.agent.response,
-                    "validation_error": self.agent._validation_error
-                })
-            
-            # Handle incoming messages
-            async for msg in ws:
-                if msg.type == WSMsgType.TEXT:
-                    try:
-                        data = json.loads(msg.data)
-                        await self._handle_client_message(data)
-                                
-                    except json.JSONDecodeError:
-                        print(f"Invalid JSON message: {msg.data}")
-                    except Exception as e:
-                        print(f"Error handling message: {e}")
-                        
-                elif msg.type == WSMsgType.ERROR:
-                    print(f"WebSocket connection closed with error: {ws.exception()}")
-                    break
-                    
-        finally:
-            async with self._lock:
-                self._clients.remove(ws)
-                
-        return ws
-
-    async def _handle_client_message(self, data: dict):
-        message_type = data.get("type")
-    
-        if message_type == ClientMessage.APPROVE_CONTEXT:
-            if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
-                await self.agent.approve_context()                
-        elif message_type == ClientMessage.APPROVE_RESPONSE:
-            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
-                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:
-                    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."""
-        await self.broadcast(ServerMessage.STATE_CHANGE, new_state.name)
-        
-    async def _handle_response_change(self, response: str):
-        """Handle response changes from the WebAgent."""
-        await self.broadcast(
-            ServerMessage.RESPONSE_UPDATE,
-            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

sia/web_socket_manager.py

@@ -0,0 +1,114 @@
+from aiohttp import web, WSMsgType
+from typing import Any, Dict, Set
+import asyncio
+import json
+from dataclasses import dataclass
+
+from .web_agent import WebAgent, WebAgentState
+from .web_io_buffer import WebIOBuffer
+
+@dataclass
+class ClientMessage:
+    """Messages that can be received from clients."""
+    APPROVE_CONTEXT = "APPROVE_CONTEXT"
+    APPROVE_RESPONSE = "APPROVE_RESPONSE" 
+    MODIFY_RESPONSE = "MODIFY_RESPONSE"
+    SEND_INPUT = "SEND_INPUT"
+
+@dataclass
+class ServerMessage:
+    """Messages that can be sent to clients."""
+    STATE_CHANGE = "STATE_CHANGE"
+    CONTEXT_UPDATE = "CONTEXT_UPDATE"
+    RESPONSE_UPDATE = "RESPONSE_UPDATE"
+    OUTPUT_UPDATE = "OUTPUT_UPDATE"
+
+class WebSocketManager:
+    """
+    Manages WebSocket connections and integrates with the WebAgent
+    """
+    def __init__(self,
+                agent: WebAgent,
+                io_buffer: WebIOBuffer):
+        """Initialize the web server."""
+        self.agent = agent
+        self.io_buffer = io_buffer
+        self._clients: Set[web.WebSocketResponse] = set()
+
+        self.agent.add_state_change_handler(self._handle_state_change)
+        self.agent.add_response_change_handler(self._handle_response_change)
+
+    async def broadcast_message(self, message: Dict):
+        """Broadcast message to all connected clients."""
+        disconnected = set()
+        for ws in self._clients:
+            try:
+                await ws.send_json(message)
+            except ConnectionResetError:
+                disconnected.add(ws)
+        self._clients -= disconnected
+        
+    def _handle_state_change(self, new_state: WebAgentState):
+        """Handle state changes from the WebAgent."""
+        asyncio.run(self.broadcast_message({
+            "type": ServerMessage.STATE_CHANGE,
+            "state": new_state.name,
+        }))
+
+    def _handle_response_change(self, response: str, validation_error: str):
+        """Handle response changes from the WebAgent."""
+        asyncio.run(self.broadcast_message({
+            "type": ServerMessage.RESPONSE_UPDATE,
+            "response": response,
+            "validation_error": validation_error
+        }))
+
+    async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
+        """Handle new WebSocket connections and messages."""
+        ws = web.WebSocketResponse(heartbeat=30)
+        await ws.prepare(request)
+
+        self._clients.add(ws)
+        
+        try:
+            await ws.send_json({
+                "type": ServerMessage.STATE_CHANGE,
+                "state": self.agent.current_state.name
+            })
+            await ws.send_json({
+                "type": ServerMessage.CONTEXT_UPDATE,
+                "context": self.agent.context
+            })
+            await ws.send_json({
+                "type": ServerMessage.RESPONSE_UPDATE,
+                "response": self.agent.response,
+                "validation_error": self.agent._validation_error
+            })
+            async for msg in ws:
+                if msg.type == WSMsgType.TEXT:
+                    #try:
+                    data = json.loads(msg.data)
+                    await self._handle_client_message(request, data)
+                    #except Exception as e:
+                    #    print(f"Error handling incoming websocket message: {msg}\n{e}")
+                elif msg.type == WSMsgType.ERROR:
+                    print(f"WebSocket connection closed with error: {ws.exception()}")
+        finally:
+            self._clients.remove(ws)
+        return ws
+
+    async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
+        """Handle incoming client message."""
+        message_type = data.get("type")
+        
+        if message_type == ClientMessage.APPROVE_CONTEXT:
+            if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
+                await asyncio.to_thread(self.agent.approve_context)
+        elif message_type == ClientMessage.APPROVE_RESPONSE:
+            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
+                await asyncio.to_thread(self.agent.approve_response)
+        elif message_type == ClientMessage.MODIFY_RESPONSE:
+            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
+                await asyncio.to_thread(self.agent.set_response, data.get("response"))
+        elif message_type == ClientMessage.SEND_INPUT:
+                self.io_buffer.append_stdin(data.get("input"))

test/mock_web_agent.py

@@ -1,53 +0,0 @@
-from typing import Optional, List, Callable
-from sia.web_agent import WebAgentState
-
-class MockWebAgent:
-    """Mock WebAgent for testing."""
-    def __init__(self):
-        self._current_state: Optional[WebAgentState] = None
-        self.context: Optional[str] = None
-        self.response: Optional[str] = None
-        self._validation_error: Optional[str] = None
-        self._state_handlers: List[Callable] = []
-        self._response_handlers: List[Callable] = []
-        
-    @property
-    def current_state(self):
-        return self._current_state
-        
-    def add_state_change_handler(self, handler):
-        self._state_handlers.append(handler)
-        
-    def add_response_change_handler(self, handler):
-        self._response_handlers.append(handler)
-
-    def set_input(self, input_text):
-        pass
-        
-    def _notify_state_handlers(self, new_state: WebAgentState):
-        """Notify all state handlers of a state change."""
-        for handler in self._state_handlers:
-            handler(new_state)
-            
-    def _notify_response_handlers(self, response: str):
-        """Notify all response handlers of a response change."""
-        for handler in self._response_handlers:
-            handler(response)
-            
-    def _set_response(self, response: str):
-        """Set response and notify handlers."""
-        self.response = response
-        self._notify_response_handlers(response)
-            
-    def approve_context(self):
-        """Handle context approval with immediate state updates."""
-        # Change to INFERENCE state and notify
-        self._current_state = WebAgentState.INFERENCE
-        self._notify_state_handlers(WebAgentState.INFERENCE)
-        
-        # Set response and notify
-        self._set_response("<reasoning>test</reasoning>")
-        
-        # Change to RESPONSE_APPROVAL state and notify
-        self._current_state = WebAgentState.RESPONSE_APPROVAL
-        self._notify_state_handlers(WebAgentState.RESPONSE_APPROVAL)

test/web_agent_test.py

@@ -88,7 +88,7 @@ class WebAgentTest(unittest.TestCase):
         """Track state changes for verification."""
         self.state_changes.append(state)
         
-    def response_change_handler(self, response: str):
+    def response_change_handler(self, response: str, validation_error: str):
         """Track response changes for verification."""
         self.response_changes.append(response)
         
@@ -97,7 +97,7 @@ class WebAgentTest(unittest.TestCase):
         self.assertEqual(self.agent.current_state, WebAgentState.CONTEXT_APPROVAL)
         self.assertEqual(self.agent.response, "")
         self.assertIsNotNone(self.agent.context)
-        self.assertIsNone(self.agent._validation_error)
+        self.assertIsNone(self.agent.validation_error)
         self.assertEqual(len(self.agent._state_change_handlers), 0)
         self.assertEqual(len(self.agent._response_change_handlers), 0)
         
@@ -133,7 +133,7 @@ class WebAgentTest(unittest.TestCase):
             self.agent.approve_response()
         self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
         
-    async def test_context_approval_flow(self):
+    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)
@@ -145,27 +145,16 @@ class WebAgentTest(unittest.TestCase):
             yield "<reasoning>test reasoning</reasoning>"
         self.mock_llm.infer.side_effect = mock_infer
         
-        # Set initial state
-        self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
-        
         # 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,
             WebAgentState.RESPONSE_APPROVAL
         ]
         self.assertEqual(self.state_changes, expected_states)
-        
+        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."""
@@ -206,7 +195,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.set_response("<reasoning>test reasoning</reasoning>")
         
         # Approve response
         self.agent.approve_response()
@@ -224,22 +213,13 @@ class WebAgentTest(unittest.TestCase):
         ]
         self.assertEqual(self.state_changes, expected_states)
         
-    async def test_response_validation(self):
+    def test_response_validation(self):
         """Test response validation during response setting."""
-        # Register handlers
         self.agent.add_response_change_handler(self.response_change_handler)
-        
-        # Mock validator to fail
         error_message = "Invalid XML"
         self.mock_validator.validate.return_value = error_message
-        
+        self.agent.set_response("<invalid>")
-        # Set response using async method
+        self.assertEqual(self.agent.validation_error, error_message)
-        await self.agent._set_response("<invalid>")
-        
-        # Verify validation error stored
-        self.assertEqual(self.agent._validation_error, error_message)
-        
-        # Verify response updated
         self.assertEqual(self.response_changes, ["<invalid>"])
         
     def tearDown(self):

test/web_server_test.py

@@ -1,251 +0,0 @@
-from aiohttp import WSMsgType
-from aiohttp import WSMsgType
-from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
-from pathlib import Path
-from sia.web_agent import WebAgentState
-from unittest.mock import patch
-import asyncio
-import asyncio
-import time
-import json
-import json
-import shutil
-import tempfile
-
-from sia.web_agent import WebAgentState
-from sia.web_io_buffer import WebIOBuffer
-from sia.web_server import WebServer, ServerMessage, ClientMessage
-from .mock_web_agent import MockWebAgent
-
-class WebServerTest(AioHTTPTestCase):
-    async def get_application(self):
-        """Create application for testing."""
-        self.io_buffer = WebIOBuffer()
-        
-        self.agent = MockWebAgent()
-
-        # Create temporary directory for static files
-        self.temp_dir = Path(tempfile.mkdtemp())
-        self.static_dir = self.temp_dir / "static"
-        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,
-            "localhost",
-            8080,
-            static_dir=self.static_dir
-        )
-        return self.web_server.app
-
-    def tearDown(self):
-        """Clean up temporary directory."""
-        if hasattr(self, 'temp_dir'):
-            shutil.rmtree(self.temp_dir)
-
-    async def wait_for_messages(self, ws, count=1, timeout=5):
-        """Helper to wait for multiple WebSocket messages."""
-        messages = []
-        try:
-            for _ in range(count):
-                msg = await ws.receive(timeout=timeout)
-                if msg.type == WSMsgType.TEXT:
-                    messages.append(json.loads(msg.data))
-                elif msg.type == WSMsgType.BINARY:
-                    messages.append(json.loads(msg.data.decode()))
-                elif msg.type == WSMsgType.CLOSED:
-                    break
-                elif msg.type == WSMsgType.ERROR:
-                    self.fail(f"WebSocket error while waiting for messages: {msg.data}")
-            return messages
-        except asyncio.TimeoutError:
-            self.fail(f"Timeout waiting for WebSocket message. Got {len(messages)} of {count}")
-        except json.JSONDecodeError as e:
-            self.fail(f"Invalid JSON in message: {e}")
-
-    async def drain_messages(self, ws, timeout=0.1):
-        """Helper to drain any pending messages."""
-        try:
-            while True:
-                msg = await ws.receive(timeout=timeout)
-                if msg.type == WSMsgType.CLOSED:
-                    break  # Normal close, no need to fail
-                elif msg.type == WSMsgType.ERROR:
-                    self.fail(f"WebSocket error while draining messages: {msg.data}")
-        except asyncio.TimeoutError:
-            pass  # Expected when no more messages
-
-    def setUp(self):
-        """Set up test case with initial state."""
-        super().setUp()
-        self.io_buffer = None
-        self.agent = None
-        self.web_server = None
-
-    async def get_application(self):
-        """Create application for testing."""
-        self.io_buffer = WebIOBuffer()
-        self.agent = MockWebAgent()
-        self.agent._current_state = WebAgentState.CONTEXT_APPROVAL  # Set initial state
-        self.web_server = WebServer(self.agent, self.io_buffer, "localhost", 8080)
-        return self.web_server.app
-
-    @unittest_run_loop
-    async def test_approve_context(self):
-        """Test context approval flow."""
-        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
-            
-            # Set agent to CONTEXT_APPROVAL state
-            self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
-            self.agent.context = "<test>context</test>"
-            await asyncio.sleep(0.1)  # Let state changes propagate
-            
-            # Send context approval
-            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)
-            
-            # Verify state change to INFERENCE
-            self.assertEqual(messages[0]["type"], ServerMessage.STATE_CHANGE)
-            self.assertEqual(messages[0]["data"], WebAgentState.INFERENCE.name)
-
-            # Verify response
-            self.assertEqual(messages[1]["type"], ServerMessage.RESPONSE_UPDATE)
-            self.assertEqual(messages[1]["data"], "<reasoning>test</reasoning>")
-            
-            # Verify state change to RESPONSE_APPROVAL
-            self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE)
-            self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name)
-
-    @unittest_run_loop
-    async def test_validation_error(self):
-        """Test handling of validation errors."""
-        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
-            
-            # Set validation error and response
-            self.agent._validation_error = "Invalid XML"
-            self.agent._set_response("<invalid>")  # Now synchronous
-            await asyncio.sleep(0.1)  # Let changes propagate
-            
-            # Should receive response update with error
-            messages = await self.wait_for_messages(ws, count=1)
-            self.assertEqual(messages[0]["type"], ServerMessage.RESPONSE_UPDATE)
-            self.assertEqual(messages[0]["validation_error"], "Invalid XML")
-            
-    @unittest_run_loop
-    async def test_multiple_clients(self):
-        """Test handling multiple client connections."""
-        ws1 = await self.client.ws_connect("/ws")
-        ws2 = await self.client.ws_connect("/ws")
-        
-        try:
-            # Drain initial messages
-            await self.drain_messages(ws1)
-            await self.drain_messages(ws2)
-            await asyncio.sleep(0.1)  # Let connections settle
-            
-            # Verify both clients are tracked
-            self.assertEqual(len(self.web_server.clients), 2)
-            
-        finally:
-            await ws1.close()
-            await ws2.close()
-            
-    @unittest_run_loop
-    async def test_client_disconnect(self):
-        """Test proper cleanup on client disconnect."""
-        ws = await self.client.ws_connect("/ws")
-        await self.drain_messages(ws)
-        await asyncio.sleep(0.1)  # Let connection settle
-        
-        self.assertEqual(len(self.web_server.clients), 1)
-        
-        await ws.close()
-        await asyncio.sleep(0.1)  # Let cleanup complete
-        
-        self.assertEqual(len(self.web_server.clients), 0)
-    
-    @unittest_run_loop
-    async def test_broadcast(self):
-        """Test broadcasting to all clients."""
-        ws1 = await self.client.ws_connect("/ws")
-        ws2 = await self.client.ws_connect("/ws")
-        
-        try:
-            # Drain initial messages
-            await self.drain_messages(ws1)
-            await self.drain_messages(ws2)
-            await asyncio.sleep(0.1)  # Let connections settle
-            
-            # Broadcast test message
-            test_msg = "test broadcast"
-            await self.web_server.broadcast(ServerMessage.OUTPUT_UPDATE, test_msg)
-            
-            # Both clients should receive message
-            msg1 = (await self.wait_for_messages(ws1, count=1))[0]
-            msg2 = (await self.wait_for_messages(ws2, count=1))[0]
-            
-            self.assertEqual(msg1["type"], ServerMessage.OUTPUT_UPDATE)
-            self.assertEqual(msg1["data"], test_msg)
-            self.assertEqual(msg2["type"], ServerMessage.OUTPUT_UPDATE)
-            self.assertEqual(msg2["data"], test_msg)
-            
-        finally:
-            await ws1.close()
-            await ws2.close()
-            
-    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):
-        async with self.client.ws_connect("/ws") as ws:
-            await self.drain_messages(ws)
-            await asyncio.sleep(0.1)
-
-            test_input = "test input"
-            await ws.send_json({
-                "type": ClientMessage.SEND_INPUT,
-                "data": test_input
-            })
-
-            # Wait for buffer to have content
-            self.assertTrue(await self.wait_for_buffer_content())
-            
-            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):
-        """Test handling of invalid JSON messages."""
-        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 invalid JSON
-            await ws.send_str("invalid json")
-            await asyncio.sleep(0.1)  # Let error process
-            
-            # Connection should remain open
-            self.assertEqual(len(self.web_server.clients), 1)

test/web_socket_manager_test.py

@@ -0,0 +1,179 @@
+from aiohttp import web
+from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
+from aiohttp.web import AppRunner
+from unittest.mock import Mock, patch, call
+import asyncio
+import unittest
+import xml.etree.ElementTree as ET
+
+from sia.web_agent import WebAgent, WebAgentState
+from sia.web_io_buffer import WebIOBuffer
+from sia.web_socket_manager import WebSocketManager, ClientMessage, ServerMessage
+from sia.working_memory import WorkingMemory
+from sia.system_metrics import SystemMetrics
+from sia.llm_engine import LlmEngine
+from sia.xml_validator import XMLValidator
+from sia.response_parser import ResponseParser
+
+class WebSocketManagerTest(AioHTTPTestCase):
+    async def get_application(self):
+        """Create test application with WebSocket handler"""
+        # Create real components where beneficial for testing
+        self.io_buffer = WebIOBuffer()
+        self.working_memory = WorkingMemory()
+        
+        # Create mock metrics with proper generate_context implementation
+        self.mock_metrics = Mock(spec=SystemMetrics)
+        metrics_elem = ET.Element("context")
+        metrics_elem.attrib = {
+            "time": "2024-10-31T12:00:00Z",
+            "cpu": "10",
+            "gpu": "20", 
+            "memory_used": "1000",
+            "memory_total": "2000",
+            "disk_used": "5000",
+            "disk_total": "10000",
+            "context": "0",
+            "stdin": "0"
+        }
+        self.mock_metrics.generate_context.return_value = metrics_elem
+        
+        # Create minimal mocks for other components
+        self.mock_llm = Mock(spec=LlmEngine)
+        def mock_infer(prompt: str, context: str):
+            yield "<reasoning>test reasoning</reasoning>"
+        self.mock_llm.infer.side_effect = mock_infer
+        self.mock_validator = Mock(spec=XMLValidator)
+        self.mock_validator.validate.return_value = None
+        
+        # Create parser with real IO buffer
+        self.parser = ResponseParser(self.io_buffer)
+
+        # Create agent with mix of real/mock components
+        self.agent = WebAgent(
+            system_prompt="test prompt",
+            action_schema="test schema",
+            working_memory=self.working_memory,
+            system_metrics=self.mock_metrics,
+            llm=self.mock_llm,
+            validator=self.mock_validator,
+            parser=self.parser
+        )
+        
+        # Create WebSocketManager
+        self.ws_manager = WebSocketManager(self.agent, self.io_buffer)
+        
+        # Create application
+        app = web.Application()
+        app.router.add_get('/ws', self.ws_manager.handle_websocket)
+        return app
+
+    @unittest_run_loop
+    async def test_initial_connection(self):
+        """Test initial WebSocket connection and state messages"""
+        async with self.client.ws_connect('/ws') as ws:
+            # Verify initial state message
+            msg = await ws.receive_json()
+            self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
+            self.assertEqual(msg['state'], WebAgentState.CONTEXT_APPROVAL.name)
+            
+            # Verify initial context message
+            msg = await ws.receive_json()
+            self.assertEqual(msg['type'], ServerMessage.CONTEXT_UPDATE)
+            self.assertIsNotNone(msg['context'])
+            
+            # Verify initial response message
+            msg = await ws.receive_json()
+            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
+            self.assertEqual(msg['response'], '')
+            self.assertIsNone(msg['validation_error'])
+
+    @unittest_run_loop
+    async def test_approve_context(self):
+        """Test context approval flow"""
+        async with self.client.ws_connect('/ws') as ws:
+            for _ in range(3):
+                msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
+            await ws.send_json({
+                'type': ClientMessage.APPROVE_CONTEXT
+            })
+            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
+            self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
+            self.assertEqual(msg['state'], WebAgentState.INFERENCE.name)
+            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
+            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
+            self.assertEqual(msg['response'], "")
+            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
+            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
+            self.assertEqual(msg['response'], "<reasoning>test reasoning</reasoning>")
+            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
+            self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
+            self.assertEqual(msg['state'], WebAgentState.RESPONSE_APPROVAL.name)
+
+    @unittest_run_loop
+    async def test_send_input(self):
+        """Test sending input through WebSocket"""
+        test_input = "test input"
+        
+        async with self.client.ws_connect('/ws') as ws:
+            for _ in range(3):
+                await ws.receive_json()
+                
+            # Send input
+            await ws.send_json({
+                'type': ClientMessage.SEND_INPUT,
+                'input': test_input
+            })
+
+            # Delay for processing
+            await asyncio.sleep(0.5)
+            
+            # Verify input was added to buffer
+            self.assertEqual(self.io_buffer.read(), test_input)
+
+    @unittest_run_loop
+    async def test_modify_response(self):
+        """Test modifying response in RESPONSE_APPROVAL state"""
+        # Set agent to RESPONSE_APPROVAL state
+        async with self.client.ws_connect('/ws') as ws:
+            for _ in range(3):
+                msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
+            await ws.send_json({
+                'type': ClientMessage.APPROVE_CONTEXT
+            })
+            for _ in range(4):
+                msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
+            modified_response = "<reasoning>modified test</reasoning>"
+            await ws.send_json({
+                'type': ClientMessage.MODIFY_RESPONSE,
+                'response': modified_response
+            })
+            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
+            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
+            self.assertEqual(msg['response'], modified_response)
+
+    @unittest_run_loop
+    async def test_multiple_clients(self):
+        """Test broadcasting to multiple WebSocket clients"""
+        # Connect first client
+        async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
+            for _ in range(3):
+                await asyncio.wait_for(ws1.receive_json(), timeout=0.5)
+                await asyncio.wait_for(ws2.receive_json(), timeout=0.5)
+            await ws1.send_json({
+                'type': ClientMessage.APPROVE_CONTEXT
+            })
+            msg1 = await ws1.receive_json()
+            msg2 = await ws2.receive_json()
+            self.assertEqual(msg1['type'], ServerMessage.STATE_CHANGE)
+            self.assertEqual(msg1['state'], WebAgentState.INFERENCE.name)
+            self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE)
+            self.assertEqual(msg2['state'], WebAgentState.INFERENCE.name)
+
+    def tearDown(self):
+        """Clean up resources"""
+        super().tearDown()
+        self.io_buffer = None
+        self.working_memory = None
+        self.agent = None
+        self.ws_manager = None