Merge pull request #198 from Qiskit/feat/197-code-assistant-httpx-lifespan

feat(code-assistant): add lifespan for httpx client lifecycle management

Author: cbjuan

qiskit-code-assistant-mcp-server/src/qiskit_code_assistant_mcp_server/server.py

@@ -23,12 +23,17 @@
 """
 
 import logging
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from typing import Any
 
+import httpx
 from fastmcp import FastMCP
 
 from qiskit_code_assistant_mcp_server.constants import (
     QCA_MCP_DEBUG_LEVEL,
+    QCA_REQUEST_TIMEOUT,
+    QCA_TOOL_X_CALLER,
     validate_configuration,
 )
 from qiskit_code_assistant_mcp_server.qca import (
@@ -41,13 +46,36 @@
     get_service_status,
     list_models,
 )
-from qiskit_code_assistant_mcp_server.utils import close_http_client
+from qiskit_code_assistant_mcp_server.utils import (
+    _get_token,
+    clear_http_client,
+    set_http_client,
+)
 
 
 # Configure logging
 logging.basicConfig(level=getattr(logging, QCA_MCP_DEBUG_LEVEL, logging.INFO))
 logger = logging.getLogger(__name__)
 
+
+@asynccontextmanager
+async def lifespan(server: FastMCP) -> AsyncIterator[None]:
+    """Manage the httpx client lifecycle."""
+    headers = {
+        "x-caller": QCA_TOOL_X_CALLER,
+        "Accept": "application/json",
+        "Authorization": f"Bearer {_get_token()}",
+    }
+    async with httpx.AsyncClient(
+        headers=headers,
+        timeout=httpx.Timeout(QCA_REQUEST_TIMEOUT),
+        limits=httpx.Limits(max_keepalive_connections=5, max_connections=10),
+    ) as client:
+        set_http_client(client)
+        yield
+    clear_http_client()
+
+
 # Initialize FastMCP server
 mcp = FastMCP(
     "Qiskit Code Assistant",
@@ -72,6 +100,7 @@
 Browse qca://models to discover available models and qca://status to check \
 service health.\
 """,
+    lifespan=lifespan,
 )
 
 logger.info("Qiskit Code Assistant MCP Server initialized")
@@ -267,28 +296,8 @@ def setup_model(model_id: str) -> str:
 
 
 if __name__ == "__main__":
-    import atexit
-
     logger.info("Starting Qiskit Code Assistant MCP Server")
-
-    # Register cleanup function
-    def cleanup() -> None:
-        import asyncio
-
-        try:
-            asyncio.run(close_http_client())
-            logger.info("HTTP client closed successfully")
-        except Exception as e:
-            logger.error(f"Error closing HTTP client: {e}")
-
-    atexit.register(cleanup)
-
-    try:
-        mcp.run(transport="stdio", show_banner=False)
-    except KeyboardInterrupt:
-        logger.info("Server interrupted, shutting down...")
-    finally:
-        cleanup()
+    mcp.run(transport="stdio", show_banner=False)
 
 
 # Assisted by watsonx Code Assistant

qiskit-code-assistant-mcp-server/src/qiskit_code_assistant_mcp_server/utils.py

@@ -155,6 +155,18 @@ def _get_token() -> str:
 _client: httpx.AsyncClient | None = None
 
 
+def set_http_client(client: httpx.AsyncClient) -> None:
+    """Set the shared HTTP client (called by server lifespan)."""
+    global _client
+    _client = client
+
+
+def clear_http_client() -> None:
+    """Clear the shared HTTP client (called on server shutdown)."""
+    global _client
+    _client = None
+
+
 def get_http_client() -> httpx.AsyncClient:
     """Get or create the shared HTTP client."""
     global _client

qiskit-code-assistant-mcp-server/tests/conftest.py

@@ -12,6 +12,7 @@
 
 """Test configuration and fixtures for Qiskit Code Assistant MCP Server tests."""
 
+import contextlib
 import os
 from unittest.mock import patch
 
@@ -42,21 +43,16 @@ async def reset_http_client():
 
     # Reset before test - force clear without trying to close
     # (might already be closed or in invalid state)
-    utils_module._client = None
+    utils_module.clear_http_client()
     utils_module._cached_token = None  # Reset cached token for fresh state
     utils_module._token_checked = False  # Reset token check flag
 
     yield
 
     # Reset after test - properly close if still open
-    if utils_module._client is not None:
-        try:
-            if not utils_module._client.is_closed:
-                await utils_module._client.aclose()
-        except Exception:
-            pass  # Ignore errors during cleanup
-        finally:
-            utils_module._client = None
+    with contextlib.suppress(Exception):
+        await utils_module.close_http_client()
+    utils_module.clear_http_client()  # Ensure cleared even if close failed
     utils_module._cached_token = None  # Reset cached token after test
     utils_module._token_checked = False  # Reset token check flag after test
 
@@ -85,13 +81,14 @@ async def http_client_for_tests(mock_env_vars):
         "Accept": "application/json",
         "Authorization": f"Bearer {token}",
     }
-    utils_module._client = httpx.AsyncClient(
+    client = httpx.AsyncClient(
         headers=headers,
         timeout=httpx.Timeout(test_timeout),
         limits=httpx.Limits(max_keepalive_connections=5, max_connections=10),
     )
+    utils_module.set_http_client(client)
 
-    yield utils_module._client
+    yield client
 
     # Cleanup is handled by reset_http_client autouse fixture
 

qiskit-code-assistant-mcp-server/tests/test_utils.py

@@ -20,10 +20,12 @@
 import respx
 
 from qiskit_code_assistant_mcp_server.utils import (
+    clear_http_client,
     close_http_client,
     get_error_message,
     get_http_client,
     make_qca_request,
+    set_http_client,
 )
 
 
@@ -181,6 +183,22 @@ async def test_close_http_client(self, mock_env_vars, mock_async_client):
             await close_http_client()
             assert client.is_closed
 
+    @pytest.mark.asyncio
+    async def test_set_http_client(self, mock_async_client):
+        """Test setting the shared HTTP client."""
+        set_http_client(mock_async_client)
+        client = get_http_client()
+        assert client is mock_async_client
+
+    def test_clear_http_client(self, mock_env_vars, mock_async_client):
+        """Test clearing the shared HTTP client."""
+        set_http_client(mock_async_client)
+        clear_http_client()
+        # After clearing, get_http_client should create a new instance
+        with patch.object(httpx, "AsyncClient", return_value=mock_async_client):
+            client = get_http_client()
+            assert client is mock_async_client
+
     @pytest.mark.asyncio
     async def test_get_client_after_close(self, mock_env_vars):
         """Test getting client after closure creates new instance."""