diff --git a/README.md b/README.md index 009ca19..196e05e 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,12 @@ Cogstack ModelServe (CMS) is a model-serving and model-governance system created for a range of CogStack NLP tasks. Targeting language models with NER and entity linking capabilities, CMS provides a one-stop shop for serving and fine-tuning models, training lifecycle management, as well as monitoring and end-to-end observability. +[![build](https://img.shields.io/github/actions/workflow/status/CogStack/CogStack-ModelServe/main.yaml)](https://github.com/CogStack/CogStack-ModelServe/actions) +[![release](https://img.shields.io/github/v/release/CogStack/CogStack-ModelServe)](https://github.com/CogStack/CogStack-ModelServe/releases) +[![docker](https://img.shields.io/docker/v/cogstacksystems/cogstack-modelserve?sort=semver&label=docker)](https://hub.docker.com/r/cogstacksystems/cogstack-modelserve) +[![python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue)](https://www.python.org/) +[![license](https://img.shields.io/badge/license-Apache%20License%202.0-blue)](https://github.com/CogStack/CogStack-ModelServe/blob/main/LICENSE) + ## Install Dependencies A virtual environment is highly recommended prior to installation. To install the dependencies, run: ```commandline @@ -254,6 +260,15 @@ Note that to enable quantization and training features, you need to install the ```commandline pip install '.[llm]' ``` + +### CMS MCP server +You can run an MCP server to expose local or remote CMS capabilities to your preferred MCP client. +To that end, install the following extra dependencies: +```commandline +pip install '.[mcp]' +``` +For detailed configuration, please refer to the [CMS MCP Server docs](./app/mcp/README.md). + #### Chat with served models You can also "chat" with the running model using the `/stream/ws` endpoint. For example: ```html diff --git a/app/api/utils.py b/app/api/utils.py index 017cb12..7316f9f 100644 --- a/app/api/utils.py +++ b/app/api/utils.py @@ -323,8 +323,8 @@ async def init_vllm_engine(app: FastAPI, ) from vllm import SamplingParams, TokensPrompt except ImportError: - logger.error("Cannot import the vLLM engine. Please install it with `pip install cms[llm]`.") - raise ExtraDependencyRequiredException("Cannot import the vLLM engine. Please install it with `pip install cms[llm]`.") + logger.error("Cannot import the vLLM engine. Please install it with `pip install '.[llm]'`.") + raise ExtraDependencyRequiredException("Cannot import the vLLM engine. Please install it with `pip install '.[llm]'`.") parser = FlexibleArgumentParser() parser = make_arg_parser(parser) diff --git a/app/cli/cli.py b/app/cli/cli.py index f354f50..2177463 100644 --- a/app/cli/cli.py +++ b/app/cli/cli.py @@ -599,7 +599,7 @@ def run_mcp_server( transport: str = typer.Option("http", help="The transport type (either 'stdio', 'sse' or 'http')"), cms_base_url: str = typer.Option("http://127.0.0.1:8000", help="The base URL of the CMS API"), cms_api_key: str = typer.Option("Bearer", help="The API key for authenticating with the CMS API"), - mcp_api_keys: str = typer.Option("", help="Comma-separated API keys for authenticating MCP clients"), + cms_mcp_api_keys: str = typer.Option("", help="Comma-separated API keys for authenticating CMS MCP clients"), cms_mcp_oauth_enabled: Optional[bool] = typer.Option(None, help="Whether to enable OAuth2 authentication for MCP clients"), github_client_id: str = typer.Option("", help="The GitHub OAuth2 client ID"), github_client_secret: str = typer.Option("", help="The GitHub OAuth2 client secret"), @@ -618,6 +618,13 @@ def run_mcp_server( port (int): The port of the MCP server. Defaults to 8080. transport (str): The transport type for the MCP server. Can be "stdio" or "http". Defaults to "stdio". cms_base_url (str): The base URL of the CMS API endpoint. Defaults to "http://localhost:8000". + cms_api_key (str): The API key for authenticating with the CMS API. Defaults to "Bearer". + cms_mcp_api_keys (str): Comma-separated API keys for authenticating CMS MCP clients. Defaults to "". + cms_mcp_oauth_enabled (Optional[bool]): Whether to enable OAuth2 authentication for MCP clients. Defaults to None. + github_client_id (str): The GitHub OAuth2 client ID, required if cms_mcp_oauth_enabled is True. Defaults to "". + github_client_secret (str): The GitHub OAuth2 client secret, required if cms_mcp_oauth_enabled is True. Defaults to an "". + google_client_id (str): The Google OAuth2 client ID, required if cms_mcp_oauth_enabled is True. Defaults to an "". + google_client_secret (str): The Google OAuth2 client secret, required if cms_mcp_oauth_enabled is True. Defaults to an "". debug (Optional[bool]): Run in debug mode if set to True. """ @@ -629,7 +636,7 @@ def run_mcp_server( os.environ["CMS_MCP_SERVER_PORT"] = str(port) os.environ["CMS_MCP_TRANSPORT"] = transport.lower() os.environ["CMS_API_KEY"] = cms_api_key - os.environ["MCP_API_KEYS"] = mcp_api_keys + os.environ["CMS_MCP_API_KEYS"] = cms_mcp_api_keys os.environ["CMS_MCP_OAUTH_ENABLED"] = "true" if cms_mcp_oauth_enabled else "false" os.environ["GITHUB_CLIENT_ID"] = github_client_id os.environ["GITHUB_CLIENT_SECRET"] = github_client_secret @@ -645,9 +652,9 @@ def run_mcp_server( logger.info(f"Connected to CMS API at {cms_base_url}") main() except ImportError as e: - logger.error(f"Cannot import MCP. Please install it with `pip install cms[mcp]`: {e}") + logger.error(f"Cannot import MCP. Please install it with `pip install '.[mcp]'`: {e}") typer.echo(f"ERROR: Cannot import MCP: {e}") - typer.echo("Please install it with `pip install cms[mcp]`.") + typer.echo("Please install it with `pip install '.[mcp]'`.") raise typer.Exit(code=1) except KeyboardInterrupt: logger.info("MCP server stopped by the user") diff --git a/app/mcp/README.md b/app/mcp/README.md index a61fde4..9d48aab 100644 --- a/app/mcp/README.md +++ b/app/mcp/README.md @@ -12,7 +12,7 @@ pip install '.[mcp]' ### 2. Set Environment Variables ```bash export CMS_BASE_URL="http://127.0.0.1:8000" # CogStack ModelServe API base URL -export MCP_API_KEYS="key1,key2,...keyN" # Optional: The API key(s) for authentication +export CMS_MCP_API_KEYS="key1,key2,...keyN" # Optional: The API key(s) for authentication ``` ### 3. Run the Server @@ -23,10 +23,84 @@ cms mcp run ```bash # HTTP transport -export CMS_MCP_TRANSPORT=http cms mcp run --transport http ``` -Once the above succeeds, the MCP server will be running on http://127.0.0.1:8080/mcp +Once the above succeeds, the MCP server will be running at http://127.0.0.1:8080/mcp + +```bash +# SSE transport +cms mcp run --transport sse +``` +Once the above succeeds, the MCP server will be running at http://127.0.0.1:8080/sse + +## Claude Desktop Configuration + +To use this MCP server with Claude Desktop, add the following configuration to your `claude_desktop_config.json` file: + +```json +{ + "mcpServers": { + "cms-mcp-server": { + "command": "npx", + "args": [ + "mcp-remote", + "http://127.0.0.1:8080/sse" + ] + } + } +} +``` + +With API-key-based authentication: +```bash +cms mcp run --transport sse --cms-mcp-api-keys "key1,key2,...keyN" +``` +```json +{ + "mcpServers": { + "cms-mcp-server": { + "command": "npx", + "args": [ + "mcp-remote", + "http://127.0.0.1:8080/sse", + "--header", + "X-API-Key:${API_KEY_HEADER}" + ], + "env": { + "API_KEY_HEADER": "ONE_OF_THE_API_KEYS" + } + } + } +} +``` + +With OAuth2-based authentication: +```bash +cms mcp run --transport sse + --cms-mcp-oauth-enabled \ + --github-client-id \ + --github-client-secret \ + --google-client-id \ + --google-client-secret +``` +```json +{ + "mcpServers": { + "cms-mcp-server": { + "command": "npx", + "args": [ + "mcp-remote", + "http://127.0.0.1:8080/sse", + "--header", + "X-API-Key:${AUTH_HEADER}" + ], + "env": { + "AUTH_HEADER": "Bearer " + } + } + } +} +``` ## Available Tools @@ -40,18 +114,18 @@ Once the above succeeds, the MCP server will be running on http://127.0.0.1:8080 ## Configuration -| Environment Variable | Default | Description | -|---------------------|---------|-------------| +| Environment Variable | Default | Description | +|---------------------|-------------------------|-------------| | `CMS_BASE_URL` | `http://127.0.0.1:8000` | ModelServe API base URL | | `CMS_MCP_SERVER_HOST` | `127.0.0.1` | MCP server host | | `CMS_MCP_SERVER_PORT` | `8080` | MCP server port | | `CMS_MCP_TRANSPORT` | `stdio` | Transport type (`stdio`, `http` or `sse`) | | `CMS_ACCESS_TOKEN` | Empty | Bearer token for ModelServe API | | `CMS_API_KEY` | `Bearer` | API key for ModelServe API | -| `MCP_API_KEYS` | None | Comma-separated API keys for authentication | -| `CMS_MCP_OAUTH_ENABLED` | `true` | Enable OAuth authentication | -| `CMS_MCP_BASE_URL` | `http://:` | Base URL for OAuth callback | -| `CMS_MCP_DEV` | `0` | Run in development mode (creates server instance) | +| `CMS_MCP_API_KEYS` | None | Comma-separated API keys for authentication | +| `CMS_MCP_OAUTH_ENABLED` | `false` | Enable OAuth authentication | +| `CMS_MCP_BASE_URL` | `http://:` | Base URL for OAuth callback | +| `CMS_MCP_DEV` | `0` | Run in development mode | ## Authentication @@ -59,8 +133,8 @@ Once the above succeeds, the MCP server will be running on http://127.0.0.1:8080 The server supports two authentication methods: ### 1. API Key Authentication -When `MCP_API_KEYS` is set, clients must authenticate using: -- **Header**: `x-api-key: your-key` +When `CMS_MCP_API_KEYS` is set, clients must authenticate using: +- **Header**: `X-API-Key: your-key` ### 2. OAuth Authentication (SSE Transport) When `CMS_MCP_OAUTH_ENABLED=true`, the server provides a built-in OAuth 2.0 login flow for SSE transport. @@ -80,7 +154,7 @@ When `CMS_MCP_OAUTH_ENABLED=true`, the server provides a built-in OAuth 2.0 logi | `GOOGLE_CLIENT_ID` | Google OAuth client ID | | `GOOGLE_CLIENT_SECRET` | Google OAuth client secret | -**Note:** OAuth credentials can also be set via environment variables or `.env` file. If not configured, the server will log a warning but continue running. +**Note:** If OAuth credentials are not configured, the server will log a warning but continue running. **Session Authentication:** After OAuth login, a session cookie (`cms_mcp_session`) is set. Subsequent MCP requests should include this cookie for authentication. diff --git a/app/mcp/oauth/routes.py b/app/mcp/oauth/routes.py deleted file mode 100644 index b43127c..0000000 --- a/app/mcp/oauth/routes.py +++ /dev/null @@ -1,406 +0,0 @@ -import os -import secrets -from typing import Optional -from fastapi import Request, Response, HTTPException, Cookie -from fastapi.responses import HTMLResponse, RedirectResponse -from fastapi import FastAPI -from app.mcp.oauth.oauth import OAuthManager -from app.mcp.logger import get_logger - -logger = get_logger(__name__) - - -def register_oauth_routes(app: FastAPI, oauth_manager: OAuthManager) -> None: - - @app.get("/oauth/login", response_class=HTMLResponse) - async def oauth_login() -> HTMLResponse: - html_content = """ - - - - CMS MCP Server - Login - - - - - - - """ - return HTMLResponse(content=html_content) - - @app.get("/oauth/authorize/{provider}") - async def oauth_authorize(provider: str, response: Response) -> RedirectResponse: - oauth_provider = oauth_manager.get_provider(provider) - if not oauth_provider: - raise HTTPException(status_code=400, detail=f"Unknown provider: {provider}") - - auth_url, state = oauth_provider.generate_authorization_url() - - response = RedirectResponse(url=auth_url) - response.set_cookie( - key=f"oauth_state_{provider}", - value=state, - max_age=300, - httponly=True, - samesite="lax", - secure=os.getenv("CMS_MCP_SECURE_COOKIES", "false").lower() == "true" - ) - - logger.info(f"Initiating OAuth flow for provider: {provider}") - return response - - @app.get("/oauth/callback/{provider}") - async def oauth_callback( - request: Request, - provider: str, - code: Optional[str] = None, - state: Optional[str] = None, - error: Optional[str] = None, - ) -> HTMLResponse: - if error: - logger.error(f"OAuth error for {provider}: {error}") - return HTMLResponse( - content=f""" - - - Authentication Error - -

❌ Authentication Failed

-

Error: {error}

- Try again - - - """, - status_code=400 - ) - - if not code or not state: - raise HTTPException(status_code=400, detail="Missing code or state") - - if request is None: - raise HTTPException(status_code=400, detail="Request object is required") - - oauth_provider = oauth_manager.get_provider(provider) - if not oauth_provider: - raise HTTPException(status_code=400, detail=f"Unknown provider: {provider}") - - stored_state = request.cookies.get(f"oauth_state_{provider}") - if not stored_state or stored_state != state: - logger.error(f"State mismatch for {provider}") - raise HTTPException(status_code=400, detail="Invalid state parameter") - - if not oauth_provider.verify_state(state): - logger.error(f"State verification failed for {provider}") - raise HTTPException(status_code=400, detail="State verification failed") - - try: - token = await oauth_provider.exchange_code_for_token(code) - user_info = await oauth_provider.get_user_info(token.access_token) - session_id = secrets.token_urlsafe(32) - oauth_manager.store_token(session_id, token) - user_email = user_info.get("email", "N/A") - user_name = user_info.get("name") or user_info.get("login", "User") - - html_content = f""" - - - - Authentication Successful - - - -
-
✅
-

Authentication Successful!

- -
-

Session ID:

-

{session_id}

-

- Use this session ID in your MCP client configuration or include the session cookie. -

-
-

- You can now use the MCP server with your authenticated session. -

- Check Session Status -
- - - """ - - response = HTMLResponse(content=html_content) - response.set_cookie( - key="cms_mcp_session", - value=session_id, - max_age=86400, - httponly=True, - samesite="lax", - secure=os.getenv("CMS_MCP_SECURE_COOKIES", "false").lower() == "true" - ) - - response.delete_cookie(f"oauth_state_{provider}") - - logger.info(f"Successfully authenticated user via {provider}: {user_email}") - return response - - except Exception as e: - logger.error(f"OAuth callback error for {provider}: {str(e)}") - raise HTTPException(status_code=500, detail=f"Authentication failed: {str(e)}") - - @app.get("/oauth/status") - async def oauth_status( - request: Request, - cms_mcp_session: Optional[str] = Cookie(None) - ) -> HTMLResponse: - if not cms_mcp_session: - return HTMLResponse( - content=""" - - - No Session - -

🔒 No Active Session

-

You are not currently authenticated.

- Login - - - """ - ) - - token = await oauth_manager.get_valid_token(cms_mcp_session) - - if not token: - return HTMLResponse( - content=""" - - - Session Expired - -

⏰ Session Expired

-

Your session has expired. Please login again.

- Login - - - """ - ) - - is_expired = "Yes" if token.is_expired() else "No" - - html_content = f""" - - - - Session Status - - - -
-

✅ Active Session

-
- Session ID: - {cms_mcp_session[:16]}... -
-
- Token Type: - {token.token_type} -
-
- Expires In: - {token.expires_in} seconds -
-
- Is Expired: - {is_expired} -
-
- Has Refresh Token: - {"Yes" if token.refresh_token else "No"} -
- Logout -
- - - """ - - return HTMLResponse(content=html_content) - - @app.get("/oauth/logout") - async def oauth_logout(cms_mcp_session: Optional[str] = Cookie(None)) -> RedirectResponse: - if cms_mcp_session: - oauth_manager.remove_token(cms_mcp_session) - - response = RedirectResponse(url="/oauth/login") - response.delete_cookie("cms_mcp_session") - - logger.info("User logged out") - return response diff --git a/app/mcp/oauth/templates/callback.html b/app/mcp/oauth/templates/callback.html index 67f1499..9d0c4a3 100644 --- a/app/mcp/oauth/templates/callback.html +++ b/app/mcp/oauth/templates/callback.html @@ -68,4 +68,4 @@

Authentication Successful!

Check Session Status - \ No newline at end of file + diff --git a/app/mcp/utils.py b/app/mcp/utils.py index 10d5050..84b5c23 100644 --- a/app/mcp/utils.py +++ b/app/mcp/utils.py @@ -15,7 +15,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Union[Any, Awaitable[Any]]: ) else: api_key = None - api_keys_env = os.environ.get("MCP_API_KEYS") + api_keys_env = os.environ.get("CMS_MCP_API_KEYS") if not api_keys_env: # No API-key-based authentication required return tool_fn(*args, **kwargs) diff --git a/app/trainers/huggingface_llm_trainer.py b/app/trainers/huggingface_llm_trainer.py index 8596bde..41900ff 100644 --- a/app/trainers/huggingface_llm_trainer.py +++ b/app/trainers/huggingface_llm_trainer.py @@ -310,8 +310,8 @@ def run( from trl import GRPOConfig, GRPOTrainer # , PPOConfig, PPOTrainer except ImportError as e: logger.exception(e) - logger.error("Cannot import the GRPO Trainer. Please install it with `pip install cms[llm]`.") - raise ExtraDependencyRequiredException("Cannot import the GRPO Trainer. Please install it with `pip install cms[llm]`.") + logger.error("Cannot import the GRPO Trainer. Please install it with `pip install '.[llm]'`.") + raise ExtraDependencyRequiredException("Cannot import the GRPO Trainer. Please install it with `pip install '.[llm]'`.") trained_model_pack_path = None redeploy = self._config.REDEPLOY_TRAINED_MODEL == "true" diff --git a/tests/app/mcp/oauth/test_routes.py b/tests/app/mcp/oauth/test_routes.py deleted file mode 100644 index 9845b15..0000000 --- a/tests/app/mcp/oauth/test_routes.py +++ /dev/null @@ -1,176 +0,0 @@ -"""Tests for OAuth routes""" -import pytest -from unittest.mock import Mock, patch, AsyncMock -from fastapi import FastAPI -from fastapi.testclient import TestClient -from app.mcp.oauth.routes import register_oauth_routes -from app.mcp.oauth.oauth import OAuthManager, OAuthToken - - -@pytest.fixture -def app(): - """Create a test FastAPI app""" - return FastAPI() - - -@pytest.fixture -def oauth_manager(): - """Create an OAuthManager instance""" - return OAuthManager(base_url="http://localhost:8080") - - -@pytest.fixture -def test_client(app, oauth_manager): - """Create a test client with OAuth routes registered""" - register_oauth_routes(app, oauth_manager) - return TestClient(app) - - -class TestOAuthRoutes: - """Tests for OAuth route handlers""" - - def test_oauth_login_returns_html(self, test_client): - """Test that /oauth/login returns HTML content""" - response = test_client.get("/oauth/login") - - assert response.status_code == 200 - assert "text/html" in response.headers["content-type"] - assert "CMS MCP Server" in response.text - assert "Continue with Google" in response.text - assert "Continue with GitHub" in response.text - - def test_oauth_authorize_unknown_provider(self, test_client): - """Test that unknown provider returns 400""" - response = test_client.get("/oauth/authorize/unknown_provider") - - assert response.status_code == 400 - assert "Unknown provider" in response.text - - def test_oauth_authorize_google_redirects(self, test_client): - """Test that Google authorization redirects""" - response = test_client.get("/oauth/authorize/google", follow_redirects=False) - - assert response.status_code == 307 # Temporary redirect - assert response.headers["location"].startswith("https://accounts.google.com") - - def test_oauth_authorize_github_redirects(self, test_client): - """Test that GitHub authorization redirects""" - response = test_client.get("/oauth/authorize/github", follow_redirects=False) - - assert response.status_code == 307 - assert response.headers["location"].startswith("https://github.com/login/oauth/authorize") - - def test_oauth_callback_missing_code_or_state(self, test_client): - """Test that missing code or state returns 400""" - response = test_client.get("/oauth/callback/google") - - assert response.status_code == 400 - assert "Missing code or state" in response.text - - def test_oauth_callback_unknown_provider(self, test_client): - """Test that unknown provider returns 400""" - response = test_client.get( - "/oauth/callback/unknown_provider", - params={"code": "test_code", "state": "test_state"} - ) - - assert response.status_code == 400 - assert "Unknown provider" in response.text - - def test_oauth_status_no_session(self, test_client): - """Test that /oauth/status returns login page when no session""" - response = test_client.get("/oauth/status") - - assert response.status_code == 200 - assert "No Active Session" in response.text or "not authenticated" in response.text.lower() - - def test_oauth_logout_redirects(self, test_client): - """Test that /oauth/login redirects to login page""" - response = test_client.get("/oauth/logout", follow_redirects=False) - - assert response.status_code == 307 - assert response.headers["location"] == "/oauth/login" - - -class TestRegisterOAuthRoutes: - """Tests for register_oauth_routes function""" - - def test_register_routes_adds_all_endpoints(self, app, oauth_manager): - """Test that all OAuth routes are registered""" - register_oauth_routes(app, oauth_manager) - - # Check that routes are added - route_paths = [route.path for route in app.routes] - assert "/oauth/login" in route_paths - assert "/oauth/authorize/{provider}" in route_paths - assert "/oauth/callback/{provider}" in route_paths - assert "/oauth/status" in route_paths - assert "/oauth/logout" in route_paths - - def test_register_routes_with_different_oauth_manager(self, app): - """Test registering routes with a custom OAuth manager""" - custom_manager = OAuthManager(base_url="http://custom:9000") - register_oauth_routes(app, custom_manager) - - # Verify the routes were added - route_paths = [route.path for route in app.routes] - assert "/oauth/login" in route_paths - assert "/oauth/authorize/{provider}" in route_paths - - -class TestOAuthCallbackWithMockedProvider: - """Tests for OAuth callback with mocked provider""" - - @pytest.fixture - def mock_oauth_manager(self): - """Create a mock OAuth manager""" - manager = Mock(spec=OAuthManager) - return manager - - @pytest.fixture - def client_with_mock(self, app, mock_oauth_manager): - """Create a test client with mocked OAuth manager""" - register_oauth_routes(app, mock_oauth_manager) - return TestClient(app) - - def test_oauth_callback_with_error(self, client_with_mock): - """Test OAuth callback handles errors""" - response = client_with_mock.get( - "/oauth/callback/google", - params={"error": "access_denied", "code": None, "state": None} - ) - - assert response.status_code == 400 - assert "Authentication Error" in response.text or "Authentication Failed" in response.text - - def test_oauth_callback_state_mismatch(self, app, mock_oauth_manager, client_with_mock): - """Test OAuth callback rejects state mismatch""" - # Setup mock to return a provider - mock_provider = Mock() - mock_provider.generate_authorization_url.return_value = ("https://auth.url", "correct_state") - mock_provider.verify_state.return_value = True - mock_oauth_manager.get_provider.return_value = mock_provider - - response = client_with_mock.get( - "/oauth/callback/google", - params={"code": "test_code", "state": "wrong_state"}, - cookies={"oauth_state_google": "correct_state"} - ) - - assert response.status_code == 400 - assert "Invalid state" in response.text - - def test_oauth_callback_verification_failure(self, app, mock_oauth_manager, client_with_mock): - """Test OAuth callback fails when state verification fails""" - mock_provider = Mock() - mock_provider.verify_state.return_value = False - mock_oauth_manager.get_provider.return_value = mock_provider - - response = client_with_mock.get( - "/oauth/callback/google", - params={"code": "test_code", "state": "test_state"}, - cookies={"oauth_state_google": "test_state"} - ) - - assert response.status_code == 400 - assert "State verification failed" in response.text diff --git a/tests/app/mcp/test_server.py b/tests/app/mcp/test_server.py index 421a3f6..135e09c 100644 --- a/tests/app/mcp/test_server.py +++ b/tests/app/mcp/test_server.py @@ -19,16 +19,16 @@ def test_server_version(self): @patch.dict( os.environ, - {"MCP_API_KEYS": "key1,key2,key3", "CMS_MCP_TRANSPORT": "http"}, + {"CMS_MCP_API_KEYS": "key1,key2,key3", "CMS_MCP_TRANSPORT": "http"}, clear=True ) def test_create_server_with_api_keys(self): app = create_server() assert app is not None - assert os.environ.get("MCP_API_KEYS") == "key1,key2,key3" + assert os.environ.get("CMS_MCP_API_KEYS") == "key1,key2,key3" @patch.dict(os.environ, {"CMS_MCP_TRANSPORT": "http"}, clear=True) def test_create_server_without_api_keys(self): app = create_server() assert app is not None - assert os.environ.get("MCP_API_KEYS") is None + assert os.environ.get("CMS_MCP_API_KEYS") is None diff --git a/tests/app/mcp/test_utils.py b/tests/app/mcp/test_utils.py index 50ed358..c55091a 100644 --- a/tests/app/mcp/test_utils.py +++ b/tests/app/mcp/test_utils.py @@ -11,13 +11,13 @@ async def test_require_api_key(): async def mock_tool(anystr, ctx): return "success" - with patch.dict(os.environ, {"MCP_API_KEYS": "key1,key2"}, clear=True): + with patch.dict(os.environ, {"CMS_MCP_API_KEYS": "key1,key2"}, clear=True): ctx = Mock() ctx.request_context.request.headers.get.return_value = "key1" result = await mock_tool("test", ctx=ctx) assert result == "success" - with patch.dict(os.environ, {"MCP_API_KEYS": "key1,key2"}, clear=True): + with patch.dict(os.environ, {"CMS_MCP_API_KEYS": "key1,key2"}, clear=True): ctx = Mock() ctx.request_context.request.headers.get.return_value = "invalid" with pytest.raises(PermissionError):