Update FastAPI dependency to version 0.116.1 and refactor LogsServer to use async context manager for file watching lifecycle. Enhance WebSocket connection handling in the frontend with exponential backoff for reconnection attempts.

Author: Dylan Huang

eval_protocol/utils/logs_server.py

@@ -1,4 +1,5 @@
 import asyncio
+from contextlib import asynccontextmanager
 import json
 import logging
 import os
@@ -8,7 +9,7 @@
 from watchdog.observers import Observer
 
 import uvicorn
-from fastapi import WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 
 from eval_protocol.dataset_logger import default_logger
 from eval_protocol.utils.vite_server import ViteServer
@@ -85,6 +86,7 @@ def broadcast_file_update(self, update_type: str, file_path: str):
         """Broadcast file update to all connected clients."""
         if not file_path.startswith(default_logger.datasets_dir):
             return
+        logger.info(f"Broadcasting file update: {update_type} {file_path}")
 
         message = {"type": update_type, "path": file_path, "timestamp": time.time()}
         # Include file contents for created and modified events
@@ -126,19 +128,29 @@ def __init__(
             os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "vite-app", "dist")
         ),
         host: str = "localhost",
-        port: int = 4789,
+        port: Optional[int] = None,
         index_file: str = "index.html",
         watch_paths: Optional[List[str]] = None,
     ):
-        super().__init__(build_dir, host, port, index_file)
-
         # Initialize WebSocket manager
         self.websocket_manager = WebSocketManager()
 
         # Set up file watching
         self.watch_paths = watch_paths or [os.getcwd()]
         self.observer = Observer()
         self.file_watcher = FileWatcher(self.websocket_manager)
+        self._file_watching_started = False
+
+        @asynccontextmanager
+        async def lifespan(app: FastAPI):
+            print("before")
+            self.start_file_watching()
+            self.websocket_manager._loop = asyncio.get_running_loop()
+            yield
+            print("after")
+            self.stop_file_watching()
+
+        super().__init__(build_dir, host, port, index_file, lifespan=lifespan)
 
         # Add WebSocket endpoint
         self._setup_websocket_routes()
@@ -174,6 +186,11 @@ async def status():
 
     def start_file_watching(self):
         """Start watching file system for changes."""
+        # Check if file watching has already been started
+        if self._file_watching_started:
+            logger.info("File watching already started, skipping")
+            return
+
         for path in self.watch_paths:
             if os.path.exists(path):
                 self.observer.schedule(self.file_watcher, path, recursive=True)
@@ -182,15 +199,18 @@ def start_file_watching(self):
                 logger.warning(f"Watch path does not exist: {path}")
 
         self.observer.start()
+        self._file_watching_started = True
         logger.info("File watching started")
 
     def stop_file_watching(self):
         """Stop watching file system."""
-        self.observer.stop()
-        self.observer.join()
-        logger.info("File watching stopped")
+        if self._file_watching_started:
+            self.observer.stop()
+            self.observer.join()
+            self._file_watching_started = False
+            logger.info("File watching stopped")
 
-    async def run_async(self, reload: bool = False):
+    async def run_async(self):
         """
         Run the logs server asynchronously with file watching.
 
@@ -212,10 +232,9 @@ async def run_async(self, reload: bool = False):
                 self.app,
                 host=self.host,
                 port=self.port,
-                reload=reload,
-                reload_dirs=self.watch_paths,
                 log_level="info",
             )
+
             server = uvicorn.Server(config)
             await server.serve()
 
@@ -224,22 +243,21 @@ async def run_async(self, reload: bool = False):
         finally:
             self.stop_file_watching()
 
-    def run(self, reload: bool = False):
+    def run(self):
         """
         Run the logs server with file watching.
 
         Args:
             reload: Whether to enable auto-reload (default: False)
         """
-        asyncio.run(self.run_async(reload))
+        asyncio.run(self.run_async())
+
+
+server = LogsServer()
+app = server.app
 
 
-def serve_logs(
-    host: str = "localhost",
-    port: int = 4789,
-    watch_paths: Optional[List[str]] = None,
-    reload: bool = False,
-):
+def serve_logs():
     """
     Convenience function to create and run a LogsServer.
 
@@ -251,9 +269,8 @@ def serve_logs(
         watch_paths: List of paths to watch for file changes
         reload: Whether to enable auto-reload
     """
-    server = LogsServer(host=host, port=port, watch_paths=watch_paths)
-    server.run(reload=reload)
+    server.run()
 
 
 if __name__ == "__main__":
-    serve_logs(reload=True)
+    serve_logs()

eval_protocol/utils/vite_server.py

@@ -1,7 +1,7 @@
 import os
 import logging
 from pathlib import Path
-from typing import Optional
+from typing import AsyncGenerator, Callable, Optional
 
 import uvicorn
 from fastapi import FastAPI, HTTPException
@@ -27,13 +27,18 @@ class ViteServer:
     """
 
     def __init__(
-        self, build_dir: str = "dist", host: str = "localhost", port: int = 8000, index_file: str = "index.html"
+        self,
+        build_dir: str = "dist",
+        host: str = "localhost",
+        port: int = 8000,
+        index_file: str = "index.html",
+        lifespan: Optional[Callable[[FastAPI], AsyncGenerator[None, None]]] = None,
     ):
         self.build_dir = Path(build_dir)
         self.host = host
         self.port = port
         self.index_file = index_file
-        self.app = FastAPI(title="Vite SPA Server")
+        self.app = FastAPI(title="Vite SPA Server", lifespan=lifespan)
 
         # Validate build directory exists
         if not self.build_dir.exists():

pyproject.toml

@@ -20,7 +20,6 @@ dependencies = [
     "requests>=2.25.0",
     "pydantic>=2.0.0",
     "dataclasses-json>=0.5.7",
-    "fastapi>=0.68.0",
     "uvicorn>=0.15.0",
     "python-dotenv>=0.19.0",
     "openai==1.78.1",
@@ -50,6 +49,7 @@ dependencies = [
     "watchdog>=2.1.0",
     "websockets>=15.0.1",
     "fireworks-ai>=0.19.12",
+    "fastapi>=0.116.1",
 ]
 
 [project.urls]
@@ -143,6 +143,7 @@ tau2 = { git = "https://github.com/sierra-research/tau2-bench.git" }
 
 [dependency-groups]
 dev = [
+    "fastapi[standard]>=0.116.1",
     "fastmcp>=2.10.6",
     "haikus==0.3.8",
     "pytest>=8.4.1",