Stub web project

sia/web_server.py

@@ -1,12 +1,19 @@
+from aiohttp import web, WSMsgType
+from pathlib import Path
+from typing import Set, Optional
 import asyncio
 import json
-from typing import Set, Optional
-from aiohttp import web, WSMsgType
 from dataclasses import dataclass
+import mimetypes
 
 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."""
@@ -27,16 +34,18 @@ class ServerMessage:
 class WebServer:
     """
     AIOHTTP-based web server that manages WebSocket connections and 
-    integrates with the WebAgent.
+    integrates with the WebAgent. Serves static files from a web directory.
     """
-    
-    def __init__(self, agent: WebAgent, io_buffer: WebIOBuffer, host: str = "localhost", port: int = 8080):
+    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()
@@ -48,8 +57,48 @@ class WebServer:
 
     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 _serve_index(self, request: web.Request) -> web.Response:
+        """
+        Serve the React application HTML for any unmatched routes to support
+        client-side routing.
+        """
+        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."""
@@ -57,11 +106,17 @@ class WebServer:
 
     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 ws://{self.host}:{self.port}/ws")
+        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."""
@@ -70,11 +125,13 @@ class WebServer:
 
     async def _send_to_client(self, ws: web.WebSocketResponse, message: dict) -> bool:
         """Send message to client, return True if successful."""
-        try:
-            await ws.send_json(message)
-            return True
-        except Exception:
-            return False
+        if not ws.closed:
+            try:
+                await ws.send_json(message)
+                return True
+            except Exception:
+                return False
+        return False
             
     async def broadcast(self, message_type: str, data: str, **kwargs):
         """Broadcast a message to all connected clients."""
@@ -91,21 +148,24 @@ class WebServer:
             for ws in self._clients:
                 if not ws.closed:
                     task = asyncio.create_task(self._send_to_client(ws, message))
-                    tasks.append(task)
+                    tasks.append((ws, task))
                 else:
                     disconnected.add(ws)
                     
             if tasks:
-                results = await asyncio.gather(*tasks, return_exceptions=True)
-                for ws, success in zip(self._clients - disconnected, results):
-                    if not success:
+                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()
+        ws = web.WebSocketResponse(heartbeat=30)
         await ws.prepare(request)
         
         async with self._lock:
@@ -161,7 +221,7 @@ class WebServer:
         
         if message_type == ClientMessage.APPROVE_CONTEXT:
             if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
-                await self.agent.approve_context()  # Make sure this is awaited
+                await self.agent.approve_context()
                 
         elif message_type == ClientMessage.APPROVE_RESPONSE:
             if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
@@ -188,4 +248,4 @@ class WebServer:
             ServerMessage.RESPONSE_UPDATE,
             response,
             validation_error=self.agent._validation_error
-        )
+        )