From 9b8cdc3d72a1ad25093ff322fc4757a06f7e7394 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 9 Feb 2026 17:19:56 -0800 Subject: [PATCH] fix: replace asserts with proper guards, wrap HTTP errors with context - Replace `assert` statements with `if/raise ValueError` (asserts can be disabled with python -O in production) - Wrap `httpx.HTTPStatusError` to provide a clear error message with server_id and status code - Add tests for HTTP error and non-dict JSON response error paths - Remove unused imports Co-Authored-By: Claude Opus 4.6 --- .../mcp_server/oauth2_token_cache.py | 23 ++++++++--- .../mcp_server/test_oauth2_token_cache.py | 39 ++++++++++++++++++- 2 files changed, 55 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index f898447f24d..0de381ee1df 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -8,6 +8,8 @@ with ``client_id``, ``client_secret``, and ``token_url``. import asyncio from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union +import httpx + from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( @@ -74,9 +76,13 @@ class MCPOAuth2TokenCache(InMemoryCache): """ client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - assert server.client_id is not None, "client_id must be set" - assert server.client_secret is not None, "client_secret must be set" - assert server.token_url is not None, "token_url must be set" + if not server.client_id or not server.client_secret or not server.token_url: + raise ValueError( + f"MCP server '{server.server_id}' missing required OAuth2 fields: " + f"client_id={bool(server.client_id)}, " + f"client_secret={bool(server.client_secret)}, " + f"token_url={bool(server.token_url)}" + ) data: Dict[str, str] = { "grant_type": "client_credentials", @@ -91,8 +97,15 @@ class MCPOAuth2TokenCache(InMemoryCache): server.server_id, ) - response = await client.post(server.token_url, data=data) - response.raise_for_status() + try: + response = await client.post(server.token_url, data=data) + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise ValueError( + f"OAuth2 token request for MCP server '{server.server_id}' " + f"failed with status {exc.response.status_code}" + ) from exc + body = response.json() if not isinstance(body, dict): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py index 2e0f78df605..55735dca98e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py @@ -5,10 +5,9 @@ Covers the critical path: resolve_mcp_auth(), token caching, auth priority, fallback to static token, and the skip-condition property. """ -import asyncio -import time from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( @@ -120,3 +119,39 @@ def test_needs_user_oauth_token_property(): # Non-OAuth2 → never needs user OAuth token assert _server(auth_type=MCPAuth.bearer_token).needs_user_oauth_token is False + + +@pytest.mark.asyncio +async def test_http_error_raises_value_error(): + """HTTP errors from the token endpoint are wrapped in a clear ValueError.""" + server = _server() + mock_response = MagicMock() + mock_response.status_code = 401 + mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( + "Unauthorized", request=MagicMock(), response=mock_response, + ) + mock_client = AsyncMock() + mock_client.post.return_value = mock_response + + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", + return_value=mock_client, + ), pytest.raises(ValueError, match="failed with status 401"): + await resolve_mcp_auth(server) + + +@pytest.mark.asyncio +async def test_non_dict_response_raises_value_error(): + """A non-dict JSON response raises a clear ValueError.""" + server = _server() + resp = MagicMock() + resp.json.return_value = ["not", "a", "dict"] + resp.raise_for_status = MagicMock() + mock_client = AsyncMock() + mock_client.post.return_value = resp + + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", + return_value=mock_client, + ), pytest.raises(ValueError, match="non-object JSON"): + await resolve_mcp_auth(server)