diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 57eb63752f5..31df1a31b91 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -289,7 +289,7 @@ if MCP_AVAILABLE: now: Optional[float] = None, ) -> None: """Terminate expired stateful sessions and drop their auth contexts.""" - now = now or time.monotonic() + now = time.monotonic() if now is None else now server_instances = getattr(session_manager_stateful, "_server_instances", {}) expired_session_ids = [] for session_id, last_seen in _stateful_session_auth_context_last_seen.items(): @@ -2684,8 +2684,8 @@ if MCP_AVAILABLE: ) -> str: """ Stable, non-reversible identifier for the caller used to bind an - mcp-session-id to its creator. ``api_key`` on UserAPIKeyAuth is - already hashed at construction time, so we can use it directly. + mcp-session-id to its creator. Hash the resolved credential before + using it so custom key formats are never stored in cleartext. For OAuth2 passthrough (``UserAPIKeyAuth()`` with no key/user_id), the caller's identity is the upstream OAuth bearer; hash it so two @@ -2701,7 +2701,10 @@ if MCP_AVAILABLE: """ if user_api_key_auth is not None: if user_api_key_auth.api_key: - return f"key:{user_api_key_auth.api_key}" + api_key_hash = hashlib.sha256( + user_api_key_auth.api_key.encode("utf-8") + ).hexdigest() + return f"key:{api_key_hash}" if user_api_key_auth.user_id: return f"user:{user_api_key_auth.user_id}" if oauth2_headers: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index fb400d8cc39..49f95575bf1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1476,6 +1476,45 @@ async def test_stateful_mcp_auth_contexts_do_not_expire_active_sessions(): mcp_server._stateful_session_active_request_counts.pop(session_id, None) +@pytest.mark.asyncio +async def test_stateful_mcp_auth_context_cleanup_respects_zero_now(): + """Explicit now=0 should be used as-is instead of falling back to monotonic.""" + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + session_id = "zero-now-stateful-session" + auth_user = UserAPIKeyAuth(api_key="zero-now-key", user_id="zero-now-user") + transport = MagicMock() + transport.terminate = AsyncMock() + + mcp_server._stateful_session_auth_contexts[session_id] = auth_user + mcp_server._stateful_session_auth_context_last_seen[session_id] = 0.0 + + try: + with ( + patch.object( + mcp_server.session_manager_stateful, + "_server_instances", + {session_id: transport}, + ), + patch.object( + mcp_server.time, + "monotonic", + return_value=mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + 1, + ), + ): + await mcp_server._purge_expired_stateful_session_auth_contexts(now=0.0) + + assert session_id in mcp_server._stateful_session_auth_contexts + assert session_id in mcp_server._stateful_session_auth_context_last_seen + transport.terminate.assert_not_awaited() + finally: + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_auth_context_last_seen.pop(session_id, None) + + @pytest.mark.asyncio async def test_owner_fingerprint_distinguishes_oauth_callers(): """ @@ -1517,6 +1556,26 @@ async def test_owner_fingerprint_distinguishes_oauth_callers(): assert "10.0.0.1" not in fp_ip_a +@pytest.mark.asyncio +async def test_owner_fingerprint_hashes_custom_api_keys(): + """Custom API key formats should not appear in owner fingerprints.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _owner_fingerprint_for, + ) + except ImportError: + pytest.skip("MCP server not available") + + auth = UserAPIKeyAuth(api_key="custom-master-key") + fp = _owner_fingerprint_for(auth) + fp_again = _owner_fingerprint_for(auth) + + assert fp == fp_again + assert fp.startswith("key:") + assert "custom-master-key" not in fp + assert fp != "key:custom-master-key" + + @pytest.mark.asyncio async def test_stateful_mcp_session_owner_mismatch_returns_403(): """