Fix store port conflict handling (issue #221) (#227) (#235)

(cherry picked from commit 01955ae)

Co-authored-by: ddsfda99 <168000702+ddsfda99@users.noreply.github.com>

Author: agent-lightning-bot

agentlightning/cli/store.py

@@ -6,12 +6,15 @@
 
 import argparse
 import asyncio
+import logging
 from typing import Iterable
 
 from agentlightning.logging import configure_logger
 from agentlightning.store.client_server import LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
 
+logger = logging.getLogger(__name__)
+
 
 def main(argv: Iterable[str] | None = None) -> int:
     parser = argparse.ArgumentParser(description="Run a LightningStore server")
@@ -22,7 +25,11 @@ def main(argv: Iterable[str] | None = None) -> int:
 
     store = InMemoryLightningStore()
     server = LightningStoreServer(store, host="0.0.0.0", port=args.port)
-    asyncio.run(server.run_forever())
+    try:
+        asyncio.run(server.run_forever())
+    except RuntimeError as exc:
+        logger.error("LightningStore server failed to start: %s", exc, exc_info=True)
+        return 1
     return 0
 
 

agentlightning/store/client_server.py

@@ -109,6 +109,7 @@ def __init__(self, store: LightningStore, host: str, port: int):
         self._uvicorn_server: uvicorn.Server | None = uvicorn.Server(self._uvicorn_config)
 
         self._serving_thread: Optional[threading.Thread] = None
+        self._server_start_exception: Optional[BaseException] = None
 
         # Process-awareness:
         # LightningStoreServer holds a plain Python object (self.store) in one process
@@ -167,17 +168,45 @@ async def start(self):
         logger.info(f"Starting server at {self.endpoint}")
 
         uvicorn_server = self._uvicorn_server
+        self._server_start_exception = None
 
         def run_server_forever():
-            asyncio.run(uvicorn_server.serve())
-
-        self._serving_thread = threading.Thread(target=run_server_forever, daemon=True)
-        self._serving_thread.start()
+            try:
+                asyncio.run(uvicorn_server.serve())
+            except (SystemExit, Exception) as exc:
+                logger.debug("LightningStore server thread exiting due to %s", exc, exc_info=exc)
+                self._server_start_exception = exc
+
+        serving_thread = threading.Thread(target=run_server_forever, daemon=True)
+        self._serving_thread = serving_thread
+        serving_thread.start()
+
+        # Wait for uvicorn to report that it has started before pinging /health.
+        start_deadline = time.time() + 10
+        while time.time() < start_deadline:
+            if uvicorn_server.started:
+                break
+            if self._server_start_exception is not None or not serving_thread.is_alive():
+                self._handle_failed_start()
+                raise RuntimeError(self._format_start_failure_reason())
+            await asyncio.sleep(0.05)
+        else:
+            self._handle_failed_start()
+            raise RuntimeError("Server failed to start within the 10 seconds.")
 
-        # Wait for /health to be available
+        # Wait for /health to be available once uvicorn reports started.
         if not await self._server_health_check():
+            self._handle_failed_start()
             raise RuntimeError("Server failed to start within the 10 seconds.")
 
+        # If startup failed (e.g. port already in use), uvicorn never flips `started`
+        # and the worker thread stops immediately. Guard against latching on to a
+        # different process that happened to satisfy the health check.
+        if not uvicorn_server.started or not serving_thread.is_alive() or self._server_start_exception is not None:
+            self._handle_failed_start()
+            failure_reason = self._format_start_failure_reason()
+            raise RuntimeError(failure_reason)
+
     async def _server_health_check(self) -> bool:
         """Checks if the server is healthy."""
         current_time = time.time()
@@ -190,22 +219,63 @@ async def _server_health_check(self) -> bool:
             await asyncio.sleep(0.1)
         return False
 
+    def _handle_failed_start(self) -> None:
+        """Clean up thread state when startup fails."""
+        if self._uvicorn_server is not None:
+            self._uvicorn_server.should_exit = True
+        if self._serving_thread is not None:
+            # Thread already exited in most failure scenarios; join defensively.
+            self._serving_thread.join(timeout=0.1)
+            self._serving_thread = None
+
+    def _format_start_failure_reason(self) -> str:
+        base_message = f"LightningStore server failed to start on {self.endpoint}."
+        if isinstance(self._server_start_exception, SystemExit):
+            return f"{base_message} Another process may already be using this port."
+        if isinstance(self._server_start_exception, OSError):
+            return f"{base_message} {self._server_start_exception.strerror}."
+        if self._server_start_exception is not None:
+            return f"{base_message} Reason: {self._server_start_exception}."
+        return f"{base_message} Another process may already be using this port."
+
     async def run_forever(self):
         """Runs the FastAPI server indefinitely.
 
         You need to call this method in the same process as the server was created in.
         """
         assert self._uvicorn_server is not None
+        uvicorn_server = self._uvicorn_server
 
         async def _wait_till_healthy():
             health = await self._server_health_check()
             if not health:
                 raise RuntimeError("Server did not become healthy within the 10 seconds.")
             logger.info("Store server is online at %s", self.endpoint)
 
+        async def _serve_capture():
+            try:
+                await uvicorn_server.serve()
+            except KeyboardInterrupt:
+                raise
+            except (SystemExit, Exception) as exc:
+                logger.debug("LightningStore server serve() raised %s", exc, exc_info=exc)
+                self._server_start_exception = exc
+                raise RuntimeError("LightningStore server failed to serve") from exc
+
         # We run _wait_till_healthy and self._uvicorn_server.serve in parallel
         # until one of them raises an exception.
-        await asyncio.gather(_wait_till_healthy(), self._uvicorn_server.serve())
+        try:
+            await asyncio.gather(_wait_till_healthy(), _serve_capture())
+        except BaseException as exc:
+            if isinstance(exc, KeyboardInterrupt):
+                raise
+            startup_failed = not uvicorn_server.started or isinstance(
+                self._server_start_exception, (SystemExit, OSError)
+            )
+            if startup_failed:
+                self._handle_failed_start()
+                raise RuntimeError(self._format_start_failure_reason())
+            raise
 
     async def stop(self):
         """Gracefully stops the running FastAPI server.

tests/store/test_client_server.py

@@ -76,6 +76,40 @@ async def server_client() -> AsyncGenerator[Tuple[LightningStoreServer, Lightnin
         await server.stop()
 
 
+@pytest.mark.asyncio
+async def test_server_start_rejects_port_conflict() -> None:
+    """Ensure startup fails loudly when the port is already owned by another store."""
+    store_a = InMemoryLightningStore()
+    port = _get_free_port()
+    server_a = LightningStoreServer(store_a, "127.0.0.1", port)
+    await server_a.start()
+
+    store_b = InMemoryLightningStore()
+    server_b = LightningStoreServer(store_b, "127.0.0.1", port)
+
+    with pytest.raises(RuntimeError, match="Another process may already be using this port"):
+        await server_b.start()
+
+    await server_a.stop()
+
+
+@pytest.mark.asyncio
+async def test_run_forever_rejects_port_conflict() -> None:
+    """Ensure run_forever also reports port conflicts with the friendly message."""
+    store_a = InMemoryLightningStore()
+    port = _get_free_port()
+    server_a = LightningStoreServer(store_a, "127.0.0.1", port)
+    await server_a.start()
+
+    store_b = InMemoryLightningStore()
+    server_b = LightningStoreServer(store_b, "127.0.0.1", port)
+
+    with pytest.raises(RuntimeError, match="Another process may already be using this port"):
+        await server_b.run_forever()
+
+    await server_a.stop()
+
+
 @pytest.mark.asyncio
 async def test_add_resources_via_server(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
     """Test that add_resources works correctly via server."""