This commit is contained in:
Yug 2026-05-01 08:09:05 +05:30
parent 12f1b36089
commit 8e5074e59a
3 changed files with 40 additions and 23 deletions

View file

@ -318,6 +318,11 @@ class MCPClient:
if key in os.environ:
safe_env[key] = os.environ[key]
if "NPM_CONFIG_CACHE" not in safe_env:
from litellm.constants import MCP_NPM_CACHE_DIR
safe_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR
return safe_env
async def _execute_session_operation(

View file

@ -2570,7 +2570,7 @@ if MCP_AVAILABLE:
status_code = 500
detail = str(e)
headers = {}
headers: Dict[str, str] = {}
if isinstance(e, HTTPException):
status_code = e.status_code
@ -2744,15 +2744,17 @@ if MCP_AVAILABLE:
# (e.g. _receive_loop), so tool handlers can't read auth_context_var reliably.
# Storing it in session-isolated storage ensures it's isolated per SSE session.
# We store it before calling server.run so it's available for the session's lifespan.
auth = MCPAuthenticatedUser( # type: ignore
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=_sse_client_ip,
)
auth = auth_context_var.get()
if not auth or not isinstance(auth, MCPAuthenticatedUser):
auth = MCPAuthenticatedUser( # type: ignore
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=_sse_client_ip,
)
_session_auth_storage[streams[0]] = auth
if session_id:
_session_id_auth_storage[session_id] = auth
@ -3056,15 +3058,19 @@ if MCP_AVAILABLE:
"_session_auth_storage; creating standalone entry "
"(cleanup may not remove it via identity check)"
)
_session_obj_auth_storage[id(session)] = MCPAuthenticatedUser(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=_client_ip,
)
auth = auth_context_var.get()
if auth and isinstance(auth, MCPAuthenticatedUser):
_session_obj_auth_storage[id(session)] = auth
else:
_session_obj_auth_storage[id(session)] = MCPAuthenticatedUser(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=_client_ip,
)
else:
# Fallback: try session-object-identity lookup first (robust)
stored: Optional[MCPAuthenticatedUser] = None

View file

@ -532,11 +532,17 @@ async def test_sse_post_messages_auth_failure():
handle_sse_post_messages,
)
with pytest.raises(HTTPException) as exc_info:
await handle_sse_post_messages(mock_scope, mock_receive, mock_send)
await handle_sse_post_messages(mock_scope, mock_receive, mock_send)
assert exc_info.value.status_code == 401
assert exc_info.value.detail == "Unauthorized"
# Verify JSONResponse 401 was sent via ASGI send
assert mock_send.called
# Extract status code from mock_send
response_start = next(
call.args[0]
for call in mock_send.mock_calls
if call.args[0].get("type") == "http.response.start"
)
assert response_start["status"] == 401
def test_generate_stable_server_id():