diff --git a/qiskit-code-assistant-mcp-server/src/qiskit_code_assistant_mcp_server/server.py b/qiskit-code-assistant-mcp-server/src/qiskit_code_assistant_mcp_server/server.py index 15d89e7..551f19b 100644 --- a/qiskit-code-assistant-mcp-server/src/qiskit_code_assistant_mcp_server/server.py +++ b/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") @@ -224,28 +253,8 @@ async def accept_completion_tool(completion_id: str) -> dict[str, Any]: 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 diff --git a/qiskit-code-assistant-mcp-server/src/qiskit_code_assistant_mcp_server/utils.py b/qiskit-code-assistant-mcp-server/src/qiskit_code_assistant_mcp_server/utils.py index 387cf0a..c7bf5c6 100644 --- a/qiskit-code-assistant-mcp-server/src/qiskit_code_assistant_mcp_server/utils.py +++ b/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 diff --git a/qiskit-code-assistant-mcp-server/tests/conftest.py b/qiskit-code-assistant-mcp-server/tests/conftest.py index d13bc7b..f1643fb 100644 --- a/qiskit-code-assistant-mcp-server/tests/conftest.py +++ b/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 diff --git a/qiskit-code-assistant-mcp-server/tests/test_utils.py b/qiskit-code-assistant-mcp-server/tests/test_utils.py index d0b280c..1d1ab52 100644 --- a/qiskit-code-assistant-mcp-server/tests/test_utils.py +++ b/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."""