scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-21 05:15:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 161

fix: close streaming response in finally block to avoid connection leak on 401

Browse source
This commit is contained in:
yash025 2026-07-12 08:44:49 +05:30
parent 5ef6f4b579
commit dc4fa82e6f
2 changed files with 36 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -2925,8 +2925,10 @@ class MCPServerManager:
handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
response = await handler.get(server_url, stream=True)
status_code = response.status_code
response.raise_for_status()
await response.aclose()
try:
response.raise_for_status()
finally:
await response.aclose()
(
authorization_servers,
resource_scopes,

32
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -7097,6 +7097,38 @@ class TestOBOEndpointDiscovery:
mock_fetch_auth.assert_not_awaited()
assert result is None or result.token_url is None
@pytest.mark.asyncio
async def test_descovery_metadata_closes_stream_on_401_oauth_challenge(self):
manager = MCPServerManager()
server_url = "https://example.com/mcp"
request = httpx.Request("GET", server_url)
error_response = MagicMock()
error_response.headers = {}
streaming_response = MagicMock()
streaming_response.status_code = 401
streaming_response.raise_for_status = MagicMock(
side_effect=httpx.HTTPStatusError("unauthorized", request=request, response=error_response)
)
streaming_response.aclose = AsyncMock()
mock_handler = MagicMock()
mock_handler.get = AsyncMock(return_value=streaming_response)
with (
patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
return_value=mock_handler,
),
patch.object(manager, "_attempt_well_known_discovery", AsyncMock(return_value=([], None))),
patch.object(manager, "_fetch_oauth_metadata_from_resource", AsyncMock(return_value=([], None))),
patch.object(manager, "_fetch_authorization_server_metadata", AsyncMock(return_value=None)),
):
result = await manager._descovery_metadata(server_url)
streaming_response.aclose.assert_awaited_once()
assert result is None
@pytest.mark.asyncio
async def test_descovery_metadata_returns_none_for_no_auth_sse_server(self):
manager = MCPServerManager()