Follow redirects only within the MCP endpoint's origin (#3397)
This commit is contained in:
Max
GitHub
parent
5fd3abc548
commit
c6762e821c
3 changed files with 1029 additions and 205 deletions
|
|
@ -1,16 +1,27 @@
|
|||
"""Tests for httpx2 utility functions."""
|
||||
"""Tests for the httpx2 helpers the client transports are built on."""
|
||||
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import httpx2
|
||||
import pytest
|
||||
|
||||
from mcp.shared._httpx_utils import create_mcp_http_client
|
||||
from mcp.shared._httpx_utils import (
|
||||
create_mcp_http_client,
|
||||
request_within_origin,
|
||||
sse_within_origin,
|
||||
stream_within_origin,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def test_default_settings():
|
||||
"""Test that default settings are applied correctly."""
|
||||
def test_default_client_uses_mcp_timeouts_and_httpx_redirect_default():
|
||||
"""The factory applies the transports' timeouts and leaves redirect following to the transports."""
|
||||
client = create_mcp_http_client()
|
||||
|
||||
assert client.follow_redirects is True
|
||||
assert client.timeout.connect == 30.0
|
||||
assert client.follow_redirects is False
|
||||
assert client.timeout == httpx2.Timeout(30.0, read=300.0)
|
||||
|
||||
|
||||
def test_custom_parameters():
|
||||
|
|
@ -22,3 +33,239 @@ def test_custom_parameters():
|
|||
|
||||
assert client.headers["Authorization"] == "Bearer token"
|
||||
assert client.timeout.connect == 60.0
|
||||
|
||||
|
||||
class _Body(httpx2.AsyncByteStream):
|
||||
"""A response body served as a real stream, recording whether the client closed it."""
|
||||
|
||||
def __init__(self, data: bytes, closed: list[bool]) -> None:
|
||||
self._data = data
|
||||
self._closed = closed
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
yield self._data
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self._closed.append(True)
|
||||
|
||||
|
||||
def _recording_client(
|
||||
redirects: dict[str, tuple[int, str]], **client_kwargs: Any
|
||||
) -> tuple[httpx2.AsyncClient, list[str], list[bool]]:
|
||||
"""A client whose server redirects each URL in `redirects` (status, Location) and answers 200
|
||||
to anything else; plus the `METHOD url` lines the server received and one entry per redirect
|
||||
response body the client closed."""
|
||||
received: list[str] = []
|
||||
closed: list[bool] = []
|
||||
|
||||
def serve(request: httpx2.Request) -> httpx2.Response:
|
||||
received.append(f"{request.method} {request.url}")
|
||||
if str(request.url) in redirects:
|
||||
status, location = redirects[str(request.url)]
|
||||
return httpx2.Response(status, headers={"location": location}, stream=_Body(b"moved", closed))
|
||||
return httpx2.Response(200, text=request.content.decode() or "ok")
|
||||
|
||||
return httpx2.AsyncClient(transport=httpx2.MockTransport(serve), **client_kwargs), received, closed
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("url", "location"),
|
||||
[
|
||||
("http://mcp.example/mcp", "http://mcp.example/mcp/"),
|
||||
("http://mcp.example/mcp", "/other/path"),
|
||||
("http://mcp.example:8080/mcp", "http://mcp.example:8080/v2/mcp"),
|
||||
("http://mcp.example/mcp", "http://MCP.EXAMPLE:80/mcp/"),
|
||||
("http://mcp.example/mcp", "https://mcp.example:443/mcp"),
|
||||
],
|
||||
)
|
||||
async def test_redirect_within_origin_is_followed_with_method_and_body(url: str, location: str):
|
||||
"""A redirect that stays on the request's origin (or upgrades it to https) is followed, and a
|
||||
307 keeps the method and body (SDK-defined policy; the re-send itself is httpx2's)."""
|
||||
client, received, closed = _recording_client({url: (307, location)})
|
||||
|
||||
async with client, stream_within_origin(client, "POST", url, content=b"payload") as response:
|
||||
await response.aread()
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.text == "payload"
|
||||
assert received == [f"POST {url}", f"POST {httpx2.URL(url).join(location)}"]
|
||||
assert closed == [True]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"location",
|
||||
[
|
||||
"http://other.example/mcp",
|
||||
"http://mcp.example:8080/mcp",
|
||||
"http://sub.mcp.example/mcp",
|
||||
"https://mcp.example:8443/mcp",
|
||||
"ftp://mcp.example/mcp",
|
||||
],
|
||||
)
|
||||
async def test_redirect_outside_origin_is_not_followed(location: str):
|
||||
"""A redirect to another origin is handed back unfollowed, the way httpx2 hands back a redirect
|
||||
with following off, and the location is never requested (SDK-defined policy)."""
|
||||
url = "http://mcp.example/mcp"
|
||||
client, received, closed = _recording_client({url: (307, location)})
|
||||
|
||||
async with client, stream_within_origin(client, "POST", url, content=b"payload") as response:
|
||||
pass
|
||||
|
||||
assert response.status_code == 307
|
||||
assert response.next_request is not None
|
||||
assert response.next_request.url == location
|
||||
assert received == [f"POST {url}"]
|
||||
assert closed == [True]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [301, 302, 303])
|
||||
async def test_method_changing_redirect_of_a_post_is_not_followed(status: int):
|
||||
"""httpx2 turns a POST into a body-less GET for 301/302/303, which would drop the message, so a
|
||||
same-origin redirect with one of those codes is handed back unfollowed (SDK-defined)."""
|
||||
url = "http://mcp.example/mcp"
|
||||
client, received, _ = _recording_client({url: (status, "/mcp/")})
|
||||
|
||||
async with client, stream_within_origin(client, "POST", url, content=b"payload") as response:
|
||||
pass
|
||||
|
||||
assert response.status_code == status
|
||||
assert received == [f"POST {url}"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [301, 302, 303, 307, 308])
|
||||
async def test_same_origin_redirect_of_a_get_is_followed_for_every_redirect_status(status: int):
|
||||
"""A GET keeps its method under every redirect status, so the SSE GET follows all of them
|
||||
within the origin (SDK-defined policy over httpx2's method rules)."""
|
||||
url = "http://mcp.example/sse"
|
||||
client, received, _ = _recording_client({url: (status, "/sse/")})
|
||||
|
||||
async with client, stream_within_origin(client, "GET", url) as response:
|
||||
await response.aread()
|
||||
|
||||
assert response.status_code == 200
|
||||
assert received == [f"GET {url}", "GET http://mcp.example/sse/"]
|
||||
|
||||
|
||||
async def test_https_to_http_on_same_host_is_outside_origin():
|
||||
"""Only the upgrade direction counts as staying on the origin; a downgrade is not followed."""
|
||||
url = "https://mcp.example/mcp"
|
||||
client, received, _ = _recording_client({url: (302, "http://mcp.example/mcp")})
|
||||
|
||||
async with client, stream_within_origin(client, "GET", url) as response:
|
||||
pass
|
||||
|
||||
assert response.status_code == 302
|
||||
assert received == [f"GET {url}"]
|
||||
|
||||
|
||||
async def test_client_configured_to_follow_redirects_is_still_scoped_to_origin():
|
||||
"""The client's own follow_redirects=True does not widen the policy: the transport helper
|
||||
decides per request (SDK-defined)."""
|
||||
url = "http://mcp.example/mcp"
|
||||
client, received, _ = _recording_client({url: (307, "http://other.example/mcp")}, follow_redirects=True)
|
||||
|
||||
async with client, stream_within_origin(client, "POST", url) as response:
|
||||
pass
|
||||
|
||||
assert response.status_code == 307
|
||||
assert received == [f"POST {url}"]
|
||||
|
||||
|
||||
async def test_redirect_past_the_client_max_redirects_budget_is_handed_back_unfollowed():
|
||||
"""Same-origin hops are bounded by the client's max_redirects; the redirect after that is not
|
||||
followed but handed back like any other, so a loop fails the one call rather than raising
|
||||
(SDK-defined; max_redirects=0 therefore means "follow none")."""
|
||||
url = "http://mcp.example/a"
|
||||
client, received, closed = _recording_client(
|
||||
{
|
||||
"http://mcp.example/a": (307, "/b"),
|
||||
"http://mcp.example/b": (307, "/c"),
|
||||
"http://mcp.example/c": (307, "/d"),
|
||||
},
|
||||
max_redirects=2,
|
||||
)
|
||||
|
||||
async with client:
|
||||
response = await request_within_origin(client, "GET", url)
|
||||
|
||||
assert response.status_code == 307
|
||||
assert response.next_request is not None
|
||||
assert response.next_request.url == "http://mcp.example/d"
|
||||
assert received == ["GET http://mcp.example/a", "GET http://mcp.example/b", "GET http://mcp.example/c"]
|
||||
assert closed == [True, True, True]
|
||||
|
||||
|
||||
async def test_redirect_location_with_userinfo_is_not_followed():
|
||||
"""A Location carrying user:password is handed back unfollowed even within the origin, since
|
||||
httpx2 would otherwise send that userinfo as Basic auth (SDK-defined)."""
|
||||
url = "http://mcp.example/mcp"
|
||||
client, received, _ = _recording_client({url: (307, "http://user:secret@mcp.example/mcp/")})
|
||||
|
||||
async with client, stream_within_origin(client, "POST", url) as response:
|
||||
pass
|
||||
|
||||
assert response.status_code == 307
|
||||
assert received == [f"POST {url}"]
|
||||
|
||||
|
||||
async def test_request_within_origin_returns_a_read_response():
|
||||
"""The non-streaming form hands back a response whose body is already read."""
|
||||
url = "http://mcp.example/mcp"
|
||||
client, received, _ = _recording_client({url: (307, "/mcp/")})
|
||||
|
||||
async with client:
|
||||
response = await request_within_origin(client, "DELETE", url)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.text == "ok"
|
||||
assert received == [f"DELETE {url}", "DELETE http://mcp.example/mcp/"]
|
||||
|
||||
|
||||
async def test_sse_within_origin_sends_event_stream_headers_and_caller_headers():
|
||||
"""The SSE form asks for an event stream exactly as client.sse() does, merged case-insensitively
|
||||
with the caller's headers, and yields an EventSource over the final response."""
|
||||
seen: list[httpx2.Headers] = []
|
||||
|
||||
def serve(request: httpx2.Request) -> httpx2.Response:
|
||||
seen.append(request.headers)
|
||||
return httpx2.Response(200, headers={"content-type": "text/event-stream"}, text="data: hello\n\n")
|
||||
|
||||
client = httpx2.AsyncClient(transport=httpx2.MockTransport(serve))
|
||||
async with client:
|
||||
async with sse_within_origin(client, "http://mcp.example/sse") as source:
|
||||
events = [event.data async for event in source]
|
||||
async with sse_within_origin(client, "http://mcp.example/sse", headers={"accept": "x/y", "k": "v"}):
|
||||
pass
|
||||
|
||||
assert events == ["hello"]
|
||||
assert seen[0]["accept"] == "text/event-stream"
|
||||
assert seen[0]["cache-control"] == "no-store"
|
||||
assert seen[1].get_list("accept") == ["x/y"]
|
||||
assert seen[1]["cache-control"] == "no-store"
|
||||
assert seen[1]["k"] == "v"
|
||||
|
||||
|
||||
async def test_auth_flow_requests_are_not_redirected():
|
||||
"""Requests an httpx2 Auth flow issues while a transport request is in flight (a token refresh,
|
||||
say) inherit the per-request no-follow setting, so a redirect on them is handed back to the
|
||||
auth flow rather than followed (httpx2 behaviour the transports rely on)."""
|
||||
received: list[str] = []
|
||||
|
||||
class TokenThenRequest(httpx2.Auth):
|
||||
async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
|
||||
token_response = yield httpx2.Request("POST", "http://mcp.example/token", content=b"grant")
|
||||
request.headers["x-token-status"] = str(token_response.status_code)
|
||||
yield request
|
||||
|
||||
def serve(request: httpx2.Request) -> httpx2.Response:
|
||||
received.append(f"{request.method} {request.url}")
|
||||
if request.url.path == "/token":
|
||||
return httpx2.Response(307, headers={"location": "http://other.example/token"})
|
||||
return httpx2.Response(200, text=request.headers["x-token-status"])
|
||||
|
||||
client = httpx2.AsyncClient(transport=httpx2.MockTransport(serve), auth=TokenThenRequest(), follow_redirects=True)
|
||||
async with client:
|
||||
response = await request_within_origin(client, "POST", "http://mcp.example/mcp")
|
||||
|
||||
assert response.text == "307"
|
||||
assert received == ["POST http://mcp.example/token", "POST http://mcp.example/mcp"]
|
||||
|
|
|
|||
|
|
@ -3,14 +3,13 @@
|
|||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock
|
||||
from unittest.mock import Mock
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import anyio
|
||||
import httpx2
|
||||
import mcp_types as types
|
||||
import pytest
|
||||
from httpx2 import ServerSentEvent
|
||||
from inline_snapshot import snapshot
|
||||
from mcp_types import (
|
||||
CallToolRequestParams,
|
||||
|
|
@ -41,6 +40,7 @@ from mcp.server.sse import SseServerTransport
|
|||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared._httpx_utils import McpHttpClientFactory
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.shared.message import SessionMessage
|
||||
from tests.interaction.transports import StreamingASGITransport
|
||||
|
||||
SERVER_NAME = "test_server_for_SSE"
|
||||
|
|
@ -58,15 +58,13 @@ def in_process_client_factory(app: Starlette) -> McpHttpClientFactory:
|
|||
auth: httpx2.Auth | None = None,
|
||||
) -> httpx2.AsyncClient:
|
||||
# The SSE GET runs until it observes a disconnect, so the bridge must let the
|
||||
# application drain on close rather than cancelling it. follow_redirects matches
|
||||
# create_mcp_http_client, the factory this one stands in for.
|
||||
# application drain on close rather than cancelling it.
|
||||
return httpx2.AsyncClient(
|
||||
transport=StreamingASGITransport(app, cancel_on_close=False),
|
||||
base_url=BASE_URL,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
auth=auth,
|
||||
follow_redirects=True,
|
||||
)
|
||||
|
||||
return factory
|
||||
|
|
@ -400,10 +398,10 @@ async def test_sse_client_handles_empty_keepalive_pings() -> None:
|
|||
send an SSE event consisting of an event ID and an empty data field in order
|
||||
to prime the client to reconnect."
|
||||
|
||||
This test mocks the SSE event stream to include empty "message" events and
|
||||
verifies the client skips them without crashing.
|
||||
The event stream served here carries an endpoint event, an empty "message"
|
||||
event (the case under test), then a real response; the client must skip the
|
||||
empty one and deliver the response.
|
||||
"""
|
||||
# Build a proper JSON-RPC response using types (not hardcoded strings)
|
||||
init_result = InitializeResult(
|
||||
protocol_version="2024-11-05",
|
||||
capabilities=ServerCapabilities(),
|
||||
|
|
@ -415,41 +413,86 @@ async def test_sse_client_handles_empty_keepalive_pings() -> None:
|
|||
result=init_result.model_dump(by_alias=True, exclude_none=True),
|
||||
)
|
||||
response_json = response.model_dump_json(by_alias=True, exclude_none=True)
|
||||
event_stream = (
|
||||
"event: endpoint\ndata: /messages/?session_id=abc123\n\n"
|
||||
"event: message\ndata: \n\n"
|
||||
f"event: message\ndata: {response_json}\n\n"
|
||||
)
|
||||
|
||||
# Mock SSE events using httpx2's ServerSentEvent: an endpoint event, an
|
||||
# empty keep-alive ping (the case under test), then a real response.
|
||||
mock_event_source = MagicMock()
|
||||
mock_event_source.__aiter__.return_value = [
|
||||
ServerSentEvent(event="endpoint", data="/messages/?session_id=abc123"),
|
||||
ServerSentEvent(event="message", data=""),
|
||||
ServerSentEvent(event="message", data=response_json),
|
||||
]
|
||||
mock_event_source.response.raise_for_status = MagicMock()
|
||||
def serve(request: httpx2.Request) -> httpx2.Response:
|
||||
assert request.url.path == "/sse"
|
||||
return httpx2.Response(200, headers={"content-type": "text/event-stream"}, text=event_stream)
|
||||
|
||||
mock_sse = MagicMock()
|
||||
mock_sse.__aenter__ = AsyncMock(return_value=mock_event_source)
|
||||
mock_sse.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_client.sse = MagicMock(return_value=mock_sse)
|
||||
mock_client.post = AsyncMock(return_value=MagicMock(status_code=200, raise_for_status=MagicMock()))
|
||||
|
||||
def mock_factory(
|
||||
def factory(
|
||||
headers: dict[str, str] | None = None,
|
||||
timeout: httpx2.Timeout | None = None,
|
||||
auth: httpx2.Auth | None = None,
|
||||
) -> httpx2.AsyncClient:
|
||||
return mock_client
|
||||
return httpx2.AsyncClient(transport=httpx2.MockTransport(serve))
|
||||
|
||||
async with sse_client("http://test/sse", httpx_client_factory=mock_factory) as (read_stream, _):
|
||||
# Read the message - should skip the empty one and get the real response
|
||||
msg = await read_stream.receive()
|
||||
# If we get here without error, the empty message was skipped successfully
|
||||
assert not isinstance(msg, Exception)
|
||||
assert isinstance(msg.message, types.JSONRPCResponse)
|
||||
assert msg.message.id == 1
|
||||
with anyio.fail_after(5):
|
||||
async with sse_client("http://test/sse", httpx_client_factory=factory) as (read_stream, _):
|
||||
msg = await read_stream.receive()
|
||||
|
||||
assert isinstance(msg, SessionMessage)
|
||||
assert isinstance(msg.message, types.JSONRPCResponse)
|
||||
assert msg.message.id == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_client_follows_redirect_within_origin_on_connect() -> None:
|
||||
"""SDK-defined: a redirect of the SSE GET that stays on the endpoint's origin is followed by
|
||||
the transport itself, with a client left at httpx2's no-follow default."""
|
||||
received: list[str] = []
|
||||
|
||||
def serve(request: httpx2.Request) -> httpx2.Response:
|
||||
received.append(str(request.url))
|
||||
if request.url.path == "/sse":
|
||||
return httpx2.Response(307, headers={"location": "/sse/"})
|
||||
assert request.url.path == "/sse/"
|
||||
return httpx2.Response(
|
||||
200, headers={"content-type": "text/event-stream"}, text="event: endpoint\ndata: /messages/\n\n"
|
||||
)
|
||||
|
||||
def factory(
|
||||
headers: dict[str, str] | None = None,
|
||||
timeout: httpx2.Timeout | None = None,
|
||||
auth: httpx2.Auth | None = None,
|
||||
) -> httpx2.AsyncClient:
|
||||
return httpx2.AsyncClient(transport=httpx2.MockTransport(serve))
|
||||
|
||||
with anyio.fail_after(5):
|
||||
async with sse_client("http://test/sse", httpx_client_factory=factory):
|
||||
pass
|
||||
|
||||
assert received == ["http://test/sse", "http://test/sse/"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_client_does_not_follow_redirect_to_another_origin_on_connect() -> None:
|
||||
"""SDK-defined: a redirect of the SSE GET to another origin is not followed, even with a client
|
||||
configured to follow redirects: connecting fails with HTTPStatusError for the redirect response
|
||||
and that origin is never contacted."""
|
||||
received: list[str] = []
|
||||
|
||||
def serve(request: httpx2.Request) -> httpx2.Response:
|
||||
received.append(str(request.url))
|
||||
return httpx2.Response(307, headers={"location": "http://other.example/sse"})
|
||||
|
||||
def factory(
|
||||
headers: dict[str, str] | None = None,
|
||||
timeout: httpx2.Timeout | None = None,
|
||||
auth: httpx2.Auth | None = None,
|
||||
) -> httpx2.AsyncClient:
|
||||
return httpx2.AsyncClient(transport=httpx2.MockTransport(serve), follow_redirects=True)
|
||||
|
||||
with anyio.fail_after(5):
|
||||
with pytest.raises(httpx2.HTTPStatusError) as exc_info:
|
||||
async with sse_client("http://test/sse", httpx_client_factory=factory):
|
||||
pytest.fail("should not connect") # pragma: no cover
|
||||
|
||||
assert exc_info.value.response.status_code == 307
|
||||
assert received == ["http://test/sse"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
|
|||
|
|
@ -358,10 +358,11 @@ async def running_app(
|
|||
|
||||
|
||||
def make_client(app: Starlette, headers: dict[str, str] | None = None) -> httpx2.AsyncClient:
|
||||
"""An httpx2 client served in process by `app`, with create_mcp_http_client's redirect default.
|
||||
"""An httpx2 client served in process by `app`.
|
||||
|
||||
(Starlette's Mount 307-redirects the bare /mcp path to /mcp/, which the SDK's own client
|
||||
factory follows.)
|
||||
Starlette's Mount 307-redirects the bare /mcp path to /mcp/. The MCP transport follows that
|
||||
same-origin redirect itself; `follow_redirects=True` is here for the tests in this file that
|
||||
POST to /mcp with this client directly.
|
||||
"""
|
||||
return httpx2.AsyncClient(
|
||||
transport=StreamingASGITransport(app), base_url=BASE_URL, headers=headers, follow_redirects=True
|
||||
|
|
@ -1089,7 +1090,9 @@ async def test_streamable_http_client_get_stream(basic_app: Starlette) -> None:
|
|||
assert resource_update_found, "ResourceUpdatedNotification not received via GET stream"
|
||||
|
||||
|
||||
def create_session_id_capturing_client(app: Starlette) -> tuple[httpx2.AsyncClient, list[str]]:
|
||||
def create_session_id_capturing_client(
|
||||
app: Starlette, transport: httpx2.AsyncBaseTransport | None = None
|
||||
) -> tuple[httpx2.AsyncClient, list[str]]:
|
||||
"""Create an in-process httpx2 client that captures the session ID from responses."""
|
||||
captured_ids: list[str] = []
|
||||
|
||||
|
|
@ -1099,9 +1102,8 @@ def create_session_id_capturing_client(app: Starlette) -> tuple[httpx2.AsyncClie
|
|||
captured_ids.append(session_id)
|
||||
|
||||
client = httpx2.AsyncClient(
|
||||
transport=StreamingASGITransport(app),
|
||||
transport=transport or StreamingASGITransport(app),
|
||||
base_url=BASE_URL,
|
||||
follow_redirects=True,
|
||||
event_hooks={"response": [capture_session_id]},
|
||||
)
|
||||
return client, captured_ids
|
||||
|
|
@ -1145,36 +1147,34 @@ async def test_streamable_http_client_session_termination(basic_app: Starlette)
|
|||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_streamable_http_client_session_termination_204(
|
||||
basic_app: Starlette, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
async def test_streamable_http_client_session_termination_204(basic_app: Starlette) -> None:
|
||||
"""Session termination also succeeds when the server answers the DELETE with 204.
|
||||
|
||||
This test patches the httpx2 client to return a 204 response for DELETEs.
|
||||
The in-process server answers the DELETE with 200; a wrapping HTTP transport rewrites that to
|
||||
204 on the way back, which is what some servers send.
|
||||
"""
|
||||
|
||||
# Save the original delete method to restore later
|
||||
original_delete = httpx2.AsyncClient.delete
|
||||
class AnswerDeleteWith204(httpx2.AsyncBaseTransport):
|
||||
def __init__(self, inner: StreamingASGITransport) -> None:
|
||||
self.inner = inner
|
||||
|
||||
# Mock the client's delete method to return a 204
|
||||
async def mock_delete(self: httpx2.AsyncClient, *args: Any, **kwargs: Any) -> httpx2.Response:
|
||||
# Call the original method to get the real response
|
||||
response = await original_delete(self, *args, **kwargs)
|
||||
async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
|
||||
response = await self.inner.handle_async_request(request)
|
||||
if request.method != "DELETE" or response.status_code != 200:
|
||||
return response
|
||||
await response.aread()
|
||||
return httpx2.Response(204, headers=response.headers, request=request)
|
||||
|
||||
# Create a new response with 204 status code but same headers
|
||||
mocked_response = httpx2.Response(
|
||||
204,
|
||||
headers=response.headers,
|
||||
content=response.content,
|
||||
request=response.request,
|
||||
)
|
||||
return mocked_response
|
||||
async def __aenter__(self) -> AnswerDeleteWith204:
|
||||
await self.inner.__aenter__()
|
||||
return self
|
||||
|
||||
# Apply the patch to the httpx2 client
|
||||
monkeypatch.setattr(httpx2.AsyncClient, "delete", mock_delete)
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
await self.inner.__aexit__(*args)
|
||||
|
||||
# Use httpx2 client with event hooks to capture session ID
|
||||
httpx_client, captured_ids = create_session_id_capturing_client(basic_app)
|
||||
httpx_client, captured_ids = create_session_id_capturing_client(
|
||||
basic_app, transport=AnswerDeleteWith204(StreamingASGITransport(basic_app))
|
||||
)
|
||||
|
||||
async with httpx_client:
|
||||
async with streamable_http_client(f"{BASE_URL}/mcp", http_client=httpx_client) as (
|
||||
|
|
|
|||
Loading…
Reference in a new issue