improve multi-server shutdown

Author: krisztianfekete

src/agentevals/api/app.py

@@ -65,6 +65,8 @@ async def event_generator():
             try:
                 while True:
                     event = await queue.get()
+                    if event is None:
+                        break
                     yield f"data: {json.dumps(event)}\n\n"
             except asyncio.CancelledError:
                 pass
@@ -125,4 +127,4 @@ async def on_startup():
 @app.on_event("shutdown")
 async def on_shutdown():
     if _trace_manager:
-        await _trace_manager.stop_cleanup_task()
+        await _trace_manager.shutdown()

src/agentevals/cli.py

@@ -155,6 +155,27 @@ def list_metrics() -> None:
         click.echo()
 
 
+def _link_server_shutdown(*servers) -> None:
+    """Link multiple uvicorn servers so a single SIGINT shuts down all of them.
+
+    Uvicorn installs per-server signal handlers; the last server's handler
+    overwrites earlier ones.  This replaces handle_exit on every server with
+    a shared callback that sets should_exit / force_exit on all of them.
+    """
+    import signal as _signal
+
+    def _shared_exit(sig, frame):
+        force = all(s.should_exit for s in servers)
+        for s in servers:
+            if force and sig == _signal.SIGINT:
+                s.force_exit = True
+            else:
+                s.should_exit = True
+
+    for s in servers:
+        s.handle_exit = _shared_exit
+
+
 @main.command("serve")
 @click.option(
     "--dev",
@@ -258,6 +279,7 @@ async def _run_dev_servers():
             )
             main_server = uvicorn.Server(main_config)
             otlp_server = uvicorn.Server(otlp_config)
+            _link_server_shutdown(main_server, otlp_server)
             await asyncio.gather(main_server.serve(), otlp_server.serve())
 
         asyncio.run(_run_dev_servers())
@@ -283,6 +305,7 @@ async def _run_ui_servers():
             )
             main_server = uvicorn.Server(main_config)
             otlp_server = uvicorn.Server(otlp_config)
+            _link_server_shutdown(main_server, otlp_server)
             await asyncio.gather(main_server.serve(), otlp_server.serve())
 
         asyncio.run(_run_ui_servers())

src/agentevals/streaming/ws_server.py

@@ -61,8 +61,10 @@ def start_cleanup_task(self) -> None:
             self._cleanup_task = asyncio.create_task(self._cleanup_old_sessions_loop())
             logger.info("Started session cleanup task (TTL: %s, max: %d)", self.session_ttl, self.max_sessions)
 
-    async def stop_cleanup_task(self) -> None:
-        """Stop the background cleanup task."""
+    async def shutdown(self) -> None:
+        """Gracefully shut down: close SSE clients and cancel background tasks."""
+        for queue in self.sse_queues:
+            queue.put_nowait(None)
         if self._cleanup_task:
             self._cleanup_task.cancel()
             try: