From 12b97c834a2d22a6c3f8959f5a2b0f59a5508bf8 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Mon, 14 Sep 2026 15:15:39 +0200 Subject: [PATCH 1/4] Python: Make HTTP response-cookie persistence explicit Configure internally created MCP, A2A, and AG-UI HTTP clients to reject response-cookie persistence while preserving caller-provided clients. Retain connection reuse and cleanup, and document the migration path for applications requiring persistent cookies. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/a2a/README.md | 9 + .../a2a/agent_framework_a2a/_agent.py | 10 +- python/packages/a2a/tests/test_a2a_agent.py | 85 ++++++- python/packages/ag-ui/README.md | 14 ++ .../ag-ui/agent_framework_ag_ui/_client.py | 7 +- .../agent_framework_ag_ui/_http_service.py | 15 +- .../ag-ui/tests/ag_ui/test_ag_ui_client.py | 216 ++++++++++++++++++ python/packages/core/AGENTS.md | 1 + python/packages/core/README.md | 15 ++ python/packages/core/agent_framework/_mcp.py | 37 +-- python/packages/core/tests/core/test_mcp.py | 127 +++++----- .../core/tests/core/test_mcp_http_auth.py | 113 ++++++++- 12 files changed, 558 insertions(+), 91 deletions(-) diff --git a/python/packages/a2a/README.md b/python/packages/a2a/README.md index 0f73d0fb423..0fe9a961560 100644 --- a/python/packages/a2a/README.md +++ b/python/packages/a2a/README.md @@ -22,6 +22,15 @@ a2a_agent = A2AAgent(url="http://remote-agent/a2a") response = await a2a_agent.run("Hello!") ``` +**Breaking change:** HTTP clients created by `A2AAgent` no longer persist response cookies. This prevents +one run's cookies from being reused by another run through the same agent. +Connection pooling and the client's existing lifetime are unchanged. + +If the remote service requires cookies for authentication, sessions, or load-balancer affinity, explicitly supply an +`httpx.AsyncClient` using `http_client=`. Its cookie behavior is preserved. +Scope that client to the intended authenticated user or security context and +manage its lifetime; do not share a user-specific cookie jar across users. + ### A2AExecutor (Hosting) The `A2AExecutor` class bridges local AI agents built with the `agent_framework` library to the A2A protocol, allowing them to be hosted and accessed by other A2A-compliant clients. diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 1a143aca33f..30912ae3b86 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -7,6 +7,7 @@ import uuid import warnings from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence +from http.cookiejar import CookieJar, DefaultCookiePolicy from typing import Any, Final, Literal, TypeAlias, cast, overload import httpx @@ -244,7 +245,8 @@ def __init__( agent_card: The agent card for the agent. url: The URL for the A2A server. client: The A2A client for the agent. - http_client: Optional httpx.AsyncClient to use. + http_client: Optional httpx.AsyncClient to use. Agent-created clients do not persist + response cookies; supplied clients retain their configured cookie behavior. auth_interceptor: Optional authentication interceptor for secured endpoints. timeout: Request timeout configuration. Can be a float (applied to all timeout components), httpx.Timeout object (for full control), or None (uses 10.0s connect, 60.0s read, @@ -282,7 +284,11 @@ def __init__( # Create or use provided httpx client if http_client is None: headers = prepend_agent_framework_to_user_agent() - http_client = httpx.AsyncClient(timeout=self._timeout_config, headers=headers) + http_client = httpx.AsyncClient( + timeout=self._timeout_config, + headers=headers, + cookies=CookieJar(policy=DefaultCookiePolicy(allowed_domains=[])), + ) self._http_client = http_client # Store for cleanup self._close_http_client = True diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 049d2a45b0f..7c8dc8d5b57 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import json from collections.abc import AsyncIterator from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -30,6 +31,7 @@ ) from agent_framework.a2a import A2AAgent from agent_framework.exceptions import AgentInvalidRequestException +from google.protobuf.json_format import MessageToDict from pytest import fixture, mark, raises, warns from agent_framework_a2a import A2AAgentSession, A2AContinuationToken, A2AServiceSessionId @@ -130,6 +132,70 @@ def a2a_agent(mock_a2a_client: MockA2AClient) -> A2AAgent: return A2AAgent(name="Test Agent", id="test-agent", client=cast(Any, mock_a2a_client), http_client=None) +@mark.parametrize("set_cookie", [True, False], ids=["set-cookie", "no-set-cookie-control"]) +async def test_run_does_not_replay_upstream_cookies_between_sessions(set_cookie: bool) -> None: + """A later session must not receive private data through an earlier session's response cookie.""" + private_record = "First caller's private record" + upstream_requests: list[httpx.Request] = [] + upstream_responses: list[httpx.Response] = [] + + def handle_upstream_request(request: httpx.Request) -> httpx.Response: + upstream_requests.append(request) + payload = json.loads(request.content) + assert request.method == "POST" + assert payload["method"] == "SendMessage" + prompt = payload["params"]["message"]["parts"][-1]["text"] + headers = {} + if prompt == "sign-in": + text = "Signed in" + if set_cookie: + headers["set-cookie"] = "remote_session=first-session; Path=/; HttpOnly; Secure" + else: + assert prompt == "read-record" + text = ( + private_record + if "remote_session=first-session" in request.headers.get("cookie", "") + else "Access denied" + ) + response = httpx.Response( + 200, + headers=headers, + json={ + "jsonrpc": "2.0", + "id": payload["id"], + "result": { + "message": MessageToDict( + A2AMessage( + message_id=str(uuid4()), + role=A2ARole.ROLE_AGENT, + parts=[Part(text=text)], + ) + ) + }, + }, + ) + upstream_responses.append(response) + return response + + async_client_type = httpx.AsyncClient + + def create_http_client(**kwargs: Any) -> httpx.AsyncClient: + # Keep the agent-owned HTTPX client and its cookie handling; replace only network I/O. + return async_client_type(transport=httpx.MockTransport(handle_upstream_request), **kwargs) + + with patch("agent_framework_a2a._agent.httpx.AsyncClient", side_effect=create_http_client): + async with A2AAgent(url="https://a2a.example.test/") as agent: + first_response = await agent.run("sign-in", session=agent.create_session()) + second_response = await agent.run("read-record", session=agent.create_session()) + + assert len(upstream_requests) == 2 + assert "cookie" not in upstream_requests[0].headers + assert ("set-cookie" in upstream_responses[0].headers) is set_cookie + assert first_response.text == "Signed in" + assert second_response.text == "Access denied" + assert "cookie" not in upstream_requests[1].headers + + def test_a2a_agent_initialization_with_client(mock_a2a_client: MockA2AClient) -> None: """Test A2AAgent initialization with provided client.""" # Use model_construct to bypass Pydantic validation for mock objects @@ -466,13 +532,22 @@ async def test_context_manager_no_cleanup_when_no_http_client() -> None: async def test_context_manager_does_not_close_caller_supplied_http_client() -> None: - """A caller-supplied http_client (without client=) must survive __aexit__. See issue #7950.""" - mock_http_client = AsyncMock() + """A caller-supplied HTTP client must stay open and retain its cookie behavior after __aexit__.""" - async with A2AAgent(url="http://localhost:9999/", http_client=cast(Any, mock_http_client)): - pass + def handle_request(request: httpx.Request) -> httpx.Response: + assert request.headers["cookie"] == "caller_cookie=kept" + return httpx.Response(200, headers={"set-cookie": "server_cookie=retained; Path=/"}) - mock_http_client.aclose.assert_not_called() + async with httpx.AsyncClient( + transport=httpx.MockTransport(handle_request), cookies={"caller_cookie": "kept"} + ) as http_client: + async with A2AAgent(url="https://a2a.example.test/", http_client=http_client): + pass + + assert not http_client.is_closed + await http_client.get("https://a2a.example.test/") + assert http_client.cookies["caller_cookie"] == "kept" + assert http_client.cookies["server_cookie"] == "retained" def test_prepare_message_for_a2a_with_multiple_contents() -> None: diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 6130f0c7588..ee9c0af4edf 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -105,6 +105,20 @@ The `AGUIChatClient` supports: - Integration with `Agent` for client-side history management - Canonical interrupt/resume passthrough (`availableInterrupts` and `resume`) +#### HTTP client ownership and cookies + +**Breaking change:** `AGUIChatClient` and `AGUIHttpService` reuse an internally owned HTTP client +for connection pooling but no longer persist response cookies. This applies to all runs, +including repeated requests with the same thread ID. Closing the AG-UI client or leaving +its async context manager closes its internally owned HTTP client. + +Applications that need cookies for upstream authentication, sessions, or load-balancer affinity, including those migrating from +the previous default of retaining response cookies, must supply an `httpx.AsyncClient` +through `http_client`. Supplied clients retain their headers, cookies, timeout, transport, +and response-cookie handling, and must be closed by the caller. Scope a cookie-bearing +client to a single authenticated principal; do not share it across users. AG-UI thread IDs +are correlation identifiers, not authentication boundaries. + ## Tool Return Helpers Use `state_update` when a backend tool needs to send different payloads to the model, the UI, and shared state. The `text` value remains the LLM-bound tool result, `tool_result` becomes the AG-UI `ToolCallResultEvent.content` for frontend rendering, and `state` is merged into durable shared state. diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index 7aa43062427..f15fc367a8d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -279,7 +279,10 @@ def __init__( Args: endpoint: The AG-UI server endpoint URL (e.g., "http://localhost:8888/") - http_client: Optional httpx.AsyncClient instance. If None, one will be created. + http_client: Optional httpx.AsyncClient instance. If None, creates a client that does not + persist response cookies. Supplied clients retain their cookie behavior and remain + caller-owned; scope cookie-based authentication to a single authenticated principal + rather than sharing it across users. timeout: Request timeout in seconds (default: 60.0) additional_properties: Additional properties to store middleware: Optional middleware to apply to the client. @@ -297,7 +300,7 @@ def __init__( ) async def close(self) -> None: - """Close the HTTP client.""" + """Close the internally created HTTP client, leaving supplied clients open.""" await self._http_service.close() async def __aenter__(self) -> Self: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_http_service.py b/python/packages/ag-ui/agent_framework_ag_ui/_http_service.py index 4fb02ef8b19..43461e5cc5e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_http_service.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_http_service.py @@ -7,6 +7,7 @@ import json import logging from collections.abc import AsyncIterable, Mapping, Sequence +from http.cookiejar import CookieJar, DefaultCookiePolicy from typing import Any, cast import httpx @@ -151,12 +152,22 @@ def __init__( Args: endpoint: AG-UI server endpoint URL (e.g., "http://localhost:8888/") - http_client: Optional httpx AsyncClient. If None, creates a new one. + http_client: Optional httpx AsyncClient. If None, creates a client that does not + persist response cookies. Supplied clients retain their cookie behavior and + remain caller-owned; scope cookie-based authentication to a single authenticated + principal rather than sharing it across users. timeout: Request timeout in seconds (default: 60.0) """ self.endpoint = endpoint.rstrip("/") self._owns_client = http_client is None - self.http_client = http_client or httpx.AsyncClient(timeout=timeout) + if http_client is not None: + self.http_client = http_client + else: + # An empty domain allowlist prevents response cookies from persisting across runs. + self.http_client = httpx.AsyncClient( + timeout=timeout, + cookies=CookieJar(policy=DefaultCookiePolicy(allowed_domains=[])), + ) async def post_run( self, diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index acb6af67e8a..65298dde9c6 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -2,8 +2,10 @@ """Tests for AGUIChatClient.""" +import asyncio import json from collections.abc import AsyncGenerator, Awaitable, Mapping, MutableSequence +from contextlib import nullcontext from typing import Any, cast import httpx @@ -70,6 +72,220 @@ async def test_client_context_manager(self) -> None: async with StubAGUIChatClient(endpoint="http://localhost:8888/") as client: assert client is not None + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + @pytest.mark.parametrize("set_cookie", [False, True], ids=["no_cookie_control", "response_cookie"]) + @pytest.mark.parametrize("first_status", [200, 503], ids=["successful_response", "failed_response"]) + @pytest.mark.parametrize("second_thread_id", ["thread_a", "thread_b"], ids=["same_thread", "different_thread"]) + async def test_default_client_does_not_replay_response_cookies_between_runs( + self, monkeypatch: MonkeyPatch, stream: bool, set_cookie: bool, first_status: int, second_thread_id: str + ) -> None: + """Response cookies must not persist between runs, regardless of thread ID.""" + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + payload = json.loads(await request.aread()) + headers = {"content-type": "text/event-stream"} + if set_cookie and len(requests) == 1: + headers["set-cookie"] = "session=thread-a-only; Path=/; HttpOnly; Secure" + events = [ + {"type": "RUN_STARTED", "threadId": payload["thread_id"], "runId": payload["run_id"]}, + {"type": "TEXT_MESSAGE_START", "messageId": "msg_1", "role": "assistant"}, + {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "ok"}, + {"type": "TEXT_MESSAGE_END", "messageId": "msg_1"}, + {"type": "RUN_FINISHED", "threadId": payload["thread_id"], "runId": payload["run_id"]}, + ] + return httpx.Response( + first_status if len(requests) == 1 else 200, + headers=headers, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(), + ) + + async with httpx.MockTransport(handler) as transport: + + async def handle_async_request( + _transport: httpx.AsyncHTTPTransport, request: httpx.Request + ) -> httpx.Response: + return await transport.handle_async_request(request) + + # Replace only network I/O, preserving the default client's real cookie handling. + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", handle_async_request) + async with StubAGUIChatClient(endpoint="https://agui.example.test/") as client: + http_client = client.http_service.http_client + for index, thread_id in enumerate(("thread_a", second_thread_id)): + messages = [Message(role="user", contents=["Hello"])] + expected_error = ( + pytest.raises(httpx.HTTPStatusError, match="503") + if index == 0 and first_status == 503 + else nullcontext() + ) + with expected_error: + if stream: + updates = [ + update + async for update in client.get_response( + messages, stream=True, options={"metadata": {"thread_id": thread_id}} + ) + ] + assert "".join(update.text for update in updates) == "ok" + else: + response = await client.get_response( + messages, options={"metadata": {"thread_id": thread_id}} + ) + assert response.text == "ok" + assert client.http_service.http_client is http_client + + assert len(requests) == 2 + assert [json.loads(request.content)["thread_id"] for request in requests] == ["thread_a", second_thread_id] + assert all(request.method == "POST" for request in requests) + assert all(request.headers["accept"] == "text/event-stream" for request in requests) + assert requests[0].headers.get("cookie") is None + assert requests[1].headers.get("cookie") is None + assert not http_client.cookies + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + async def test_default_client_does_not_replay_cookies_while_another_response_is_open( + self, monkeypatch: MonkeyPatch, stream: bool + ) -> None: + """Response cookies must be rejected before an overlapping request is built.""" + requests: list[httpx.Request] = [] + first_headers_received = asyncio.Event() + second_request_received = asyncio.Event() + response_body = ( + b'data: {"type":"TEXT_MESSAGE_START","messageId":"msg_1","role":"assistant"}\n\n' + b'data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"msg_1","delta":"ok"}\n\n' + b'data: {"type":"TEXT_MESSAGE_END","messageId":"msg_1"}\n\n' + ) + + class PendingResponseStream(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncGenerator[bytes, None]: + # HTTPX has processed response headers before it starts reading this body. + first_headers_received.set() + await second_request_received.wait() + yield response_body + + async def handle_async_request(_transport: httpx.AsyncHTTPTransport, request: httpx.Request) -> httpx.Response: + requests.append(request) + await request.aread() + if len(requests) == 1: + return httpx.Response( + 200, + headers={ + "content-type": "text/event-stream", + "set-cookie": "session=thread-a-only; Path=/; HttpOnly; Secure", + }, + stream=PendingResponseStream(), + ) + second_request_received.set() + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=response_body) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", handle_async_request) + async with StubAGUIChatClient(endpoint="https://agui.example.test/") as client: + http_client = client.http_service.http_client + + async def run(thread_id: str) -> str: + if thread_id == "thread_b": + await first_headers_received.wait() + messages = [Message(role="user", contents=["Hello"])] + if stream: + updates = [ + update + async for update in client.get_response( + messages, stream=True, options={"metadata": {"thread_id": thread_id}} + ) + ] + return "".join(update.text for update in updates) + response = await client.get_response(messages, options={"metadata": {"thread_id": thread_id}}) + return response.text + + tasks = [asyncio.create_task(run(thread_id)) for thread_id in ("thread_a", "thread_b")] + try: + responses = await asyncio.wait_for(asyncio.gather(*tasks), timeout=10) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + assert responses == ["ok", "ok"] + assert client.http_service.http_client is http_client + + assert len(requests) == 2 + assert [json.loads(request.content)["thread_id"] for request in requests] == ["thread_a", "thread_b"] + assert requests[0].headers.get("cookie") is None + assert requests[1].headers.get("cookie") is None + assert not http_client.cookies + + @pytest.mark.parametrize("raise_in_context", [False, True], ids=["normal_exit", "exception_exit"]) + async def test_context_manager_closes_owned_http_client(self, raise_in_context: bool) -> None: + """Owned HTTP clients close on both normal and exceptional context exits.""" + client = StubAGUIChatClient(endpoint="https://agui.example.test/", timeout=17.0) + http_client = client.http_service.http_client + try: + with pytest.raises(RuntimeError, match="context failed") if raise_in_context else nullcontext(): + async with client: + assert not http_client.is_closed + assert http_client.timeout == httpx.Timeout(17.0) + if raise_in_context: + raise RuntimeError("context failed") + + assert http_client.is_closed + await client.close() + assert http_client.is_closed + finally: + await http_client.aclose() + + @pytest.mark.parametrize("raise_in_context", [False, True], ids=["normal_exit", "exception_exit"]) + async def test_context_manager_preserves_supplied_http_client(self, raise_in_context: bool) -> None: + """Caller-owned clients retain their configuration and remain usable after cleanup.""" + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + headers={ + "content-type": "text/event-stream", + "set-cookie": "response=session-only; Path=/; HttpOnly; Secure", + }, + content=( + b'data: {"type":"RUN_STARTED","threadId":"thread_1","runId":"run_1"}\n\n' + b'data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"msg_1","delta":"ok"}\n\n' + b'data: {"type":"RUN_FINISHED","threadId":"thread_1","runId":"run_1"}\n\n' + ), + ) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), + headers={"Authorization": "Bearer test-token"}, + cookies={"configured": "caller-owned"}, + timeout=17.0, + ) as http_client: + client = StubAGUIChatClient(endpoint="https://agui.example.test/", http_client=http_client) + with pytest.raises(RuntimeError, match="context failed") if raise_in_context else nullcontext(): + async with client: + assert client.http_service.http_client is http_client + response = await client.get_response( + [Message(role="user", contents=["Hello"])], + options={"metadata": {"thread_id": "thread_1"}}, + ) + assert response.text == "ok" + if raise_in_context: + raise RuntimeError("context failed") + + await client.close() + assert not http_client.is_closed + assert http_client.headers["authorization"] == "Bearer test-token" + assert http_client.cookies["configured"] == "caller-owned" + assert http_client.cookies["response"] == "session-only" + assert http_client.timeout == httpx.Timeout(17.0) + response = await http_client.get("https://agui.example.test/health") + assert response.status_code == 200 + + assert [request.method for request in requests] == ["POST", "GET"] + assert all(request.headers["authorization"] == "Bearer test-token" for request in requests) + assert requests[0].headers["cookie"] == "configured=caller-owned" + assert set(requests[1].headers["cookie"].split("; ")) == {"configured=caller-owned", "response=session-only"} + async def test_extract_state_from_messages_no_state(self) -> None: """Test state extraction when no state is present.""" client = StubAGUIChatClient(endpoint="http://localhost:8888/") diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 08d6790bd9c..8b724f9bc32 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -166,6 +166,7 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID - **`additional_tool_argument_names`** (constructor arg on all `MCPTool` subclasses) - Opt extra argument names back into the allowlist. Accepts a `Sequence[str]` (applied to every tool) or a `Mapping[str, Sequence[str]]` keyed by **remote tool name**, where the reserved key `"*"` denotes global extras. It is configured only in user code at construction; there is **no per-call/runtime override**, so a model-issued tool call cannot change which names pass through — but note this constrains the *model*, not the *server*, which still widens the effective allowlist through its schema. To use a server that accepts `additionalProperties: true`, list the extra names here and then either (1) manually extend that tool's `inputSchema` (via the `.functions` list after connecting) so the model is prompted to supply them, or (2) supply the values yourself via `function_invocation_kwargs`. If a normal forwarded argument name is supplied by both the model and `function_invocation_kwargs`, the model-supplied value wins; `_meta` is the exception and only trusted runtime/caller metadata is used. - **MCP HTTP header request scoping** - `static_headers` supplies fixed, origin-scoped headers without serializing concurrent calls; `header_provider` resolves dynamic per-call headers under the instance call lock. When both are configured, dynamic values override fixed values with the same name. With a shared `http_client`, processing stays scoped to the originating `MCPStreamableHTTPTool`, and every injected header is stripped from cross-origin redirects. The session exit stack removes its request hook after transport shutdown, including failed initialization or discovery, and closes framework-created HTTP clients. Failed discovery resets the connection and discovery flags and rolls back partial function/metadata additions. Framework-created sessions are discarded; constructor-supplied sessions remain caller-owned and reusable across cleanup, close, and reset. - **MCP lifecycle caller cancellation** - The lifecycle owner skips cancelled queued connect requests and waits for the caller to acknowledge successful setup. If the caller cancels before accepting a newly established connection, the owner tears it down before processing the next request. Cancelling a close waiter does not interrupt teardown, and cancelling a redundant connect does not discard a previously established session. +- **Streamable HTTP cookies** - Framework-created clients reject response-cookie persistence with or without a `header_provider`; explicit `Cookie` headers from `static_headers` or a provider remain supported. Caller-provided clients retain their cookie behavior and ownership. Applications requiring cookie persistence must scope clients and MCP sessions to one authenticated principal. Cookie rejection does not isolate MCP protocol sessions or other server-side state. - **`function_invocation_kwargs` and MCP servers** - That dict is shared across every tool in the run, including every attached `MCPTool`, and any name in it reaches a server that declares a matching `inputSchema` property. `header_provider` does not mitigate this — it reads the kwargs without consuming them. To keep a credential out of tool arguments, source it outside `function_invocation_kwargs`: read a `ContextVar` inside the provider (this still allows a different value per request), configure a custom `http_client`, or use `env` for `MCPStdioTool`. - **Sampling guardrails** (`sampling_callback`) - Passing `client=` advertises `SamplingCapability` so the server can send `sampling/createMessage`. Because remote servers are untrusted (confused-deputy risk), the default `sampling_callback` is **deny-by-default** and applies, in order: a per-session rate limit (`sampling_max_requests`, default `_DEFAULT_SAMPLING_MAX_REQUESTS`), an approval gate (`sampling_approval_callback`), and a `maxTokens` cap (`sampling_max_tokens`, default `_DEFAULT_SAMPLING_MAX_TOKENS`). The approval callback (constructor arg on all subclasses; exported type alias `SamplingApprovalCallback`) receives the raw `CreateMessageRequestParams`, may be sync or async, and must return truthy to approve. When it is `None` (the default) every sampling request is denied; pass `lambda params: True` to restore legacy auto-approve as an explicit opt-in. Requests and denials are logged at WARNING (content is not logged). The per-session counter resets in `_reset_session_state`. - **`MCPTaskOptions`** (experimental, `MCP_LONG_RUNNING_TASKS` feature, **frozen**) - Per-tool-instance options controlling the SEP-2663 long-running task lifecycle. When the server advertises a tool with `execution.taskSupport == "required"`, `MCPTool.call_tool` transparently routes through `call_tool_as_task`, which sends an augmented `tools/call`, polls `tasks/get` until terminal, and reinterprets `tasks/result` as a normal `CallToolResult`. Instances are immutable; replace via `MCPTool.task_options = MCPTaskOptions(...)`. Fields: diff --git a/python/packages/core/README.md b/python/packages/core/README.md index cc145b0aa3b..6af0ecaf7d9 100644 --- a/python/packages/core/README.md +++ b/python/packages/core/README.md @@ -168,6 +168,21 @@ asyncio.run(main()) You can explore additional agent samples [here](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents). +### MCP HTTP client ownership and cookies + +**Breaking change:** HTTP clients created by `MCPStreamableHTTPTool` no longer persist +response cookies. This applies with or without a `header_provider`, including repeated +calls within the same MCP session. Connection reuse, timeout and redirect defaults, +and cleanup on close or reset are preserved. Explicit `Cookie` headers supplied through +`static_headers` or a `header_provider` are still supported. + +Applications relying on cookie-based authentication, sessions, or load-balancer affinity +must supply an `httpx.AsyncClient` through `http_client=`. Supplied clients retain their +configuration and response-cookie handling and must be closed by the caller. Scope +cookie-bearing clients and MCP tool sessions to one authenticated principal; do not +share a user-specific cookie jar across users. Rejecting cookies does not partition +MCP protocol sessions or other server-side state between principals. + ## 5. Multi-Agent Orchestration Coordinate multiple agents to collaborate on complex tasks using orchestration patterns: diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index d683e99493e..0ec17ecb076 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -18,6 +18,7 @@ from dataclasses import dataclass from datetime import date, datetime, timedelta from functools import partial +from http.cookiejar import CookieJar, DefaultCookiePolicy from inspect import isawaitable from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict, cast @@ -3494,7 +3495,7 @@ def __init__( The arguments are used to create a streamable HTTP client using the new ``mcp.client.streamable_http.streamable_http_client`` API. If an asyncClient is provided via ``http_client``, it will be used as the underlying transport client. - Otherwise, the ``streamable_http_client`` API will create and manage a default client. + Otherwise, the tool creates and manages a default client without response-cookie persistence. Args: name: The name of the tool. @@ -3551,8 +3552,13 @@ def __init__( (``min(requested, cap)``); ``None`` disables it. sampling_max_requests: Per-session cap on the number of sampling requests; further requests are rejected. Resets on reconnect. ``None`` disables it. - http_client: Optional asyncClient to use. If not provided, the - ``streamable_http_client`` API will create and manage a default client. + http_client: Optional asyncClient to use. If not provided, the tool creates and manages + a client that does not persist response cookies, with or without a ``header_provider``. + Explicit ``Cookie`` headers from ``static_headers`` or a ``header_provider`` are supported. + Supplied clients retain their cookie behavior and remain caller-owned. Applications + requiring cookie-based sessions must supply a client scoped to one authenticated + principal and manage its lifetime. Cookie rejection does not partition MCP protocol + sessions or other server-side state between principals. Use ``static_headers`` for fixed headers. To configure timeouts or other HTTP client settings, create and pass your own ``asyncClient`` instance. Security: when you attach sensitive headers (e.g. authentication tokens) @@ -3699,17 +3705,20 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: """ from httpx import URL, AsyncClient, Timeout + target_origin = ( + _url_origin(URL(self.url)) if self._static_headers or self._header_provider is not None else None + ) http_client = self._httpx_client - if self._static_headers or self._header_provider is not None: - target_origin = _url_origin(URL(self.url)) - if http_client is None: - http_client = AsyncClient( - follow_redirects=True, - timeout=Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT), - ) - self._httpx_client = http_client - self._exit_stack.push_async_callback(self._close_owned_http_client, http_client) + if http_client is None: + http_client = AsyncClient( + follow_redirects=True, + timeout=Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT), + cookies=CookieJar(policy=DefaultCookiePolicy(allowed_domains=[])), + ) + self._httpx_client = http_client + self._exit_stack.push_async_callback(self._close_owned_http_client, http_client) + if target_origin is not None: if not hasattr(self, "_inject_headers_hook"): async def _inject_headers(request: Request) -> None: # ruff:ignore[unused-async] @@ -3775,9 +3784,7 @@ async def _inject_headers(request: Request) -> None: # ruff:ignore[unused-async # while successful sessions keep the hook through transport shutdown. self._exit_stack.callback(self._remove_header_hook) - transport_http_client = ( - _MCPHeaderScopedClient(http_client, self._header_request_owner) if http_client is not None else None - ) + transport_http_client = _MCPHeaderScopedClient(http_client, self._header_request_owner) return streamable_http_client( url=self.url, diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 7102dbd7e35..6af4b809823 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -4672,24 +4672,32 @@ def test_mcp_stdio_tool_get_mcp_client_with_env_and_kwargs(): ) -def test_mcp_streamable_http_tool_get_mcp_client_all_params(): +async def test_mcp_streamable_http_tool_get_mcp_client_all_params(): """Test MCPStreamableHTTPTool.get_mcp_client() with all parameters.""" + from agent_framework._mcp import MCP_DEFAULT_SSE_READ_TIMEOUT, MCP_DEFAULT_TIMEOUT + tool = MCPStreamableHTTPTool( name="test", url="http://example.com", terminate_on_close=True, ) - with patch("mcp.client.streamable_http.streamable_http_client") as mock_http_client: - tool.get_mcp_client() - - # Verify streamable_http_client was called with None for http_client - # (since we didn't provide one, the API will create its own) - mock_http_client.assert_called_once_with( - url="http://example.com", - http_client=None, - terminate_on_close=True, - ) + try: + with patch("mcp.client.streamable_http.streamable_http_client") as mock_http_client: + tool.get_mcp_client() + http_client = tool._httpx_client + assert http_client is not None + assert http_client.follow_redirects is True + assert http_client.timeout.connect == MCP_DEFAULT_TIMEOUT + assert http_client.timeout.read == MCP_DEFAULT_SSE_READ_TIMEOUT + assert not http_client.event_hooks["request"] + assert mock_http_client.call_args.kwargs["http_client"]._client is http_client + assert mock_http_client.call_args.kwargs["url"] == "http://example.com" + assert mock_http_client.call_args.kwargs["terminate_on_close"] is True + mock_http_client.assert_called_once() + finally: + await tool.close() + assert http_client.is_closed def test_mcp_websocket_tool_get_mcp_client_with_kwargs(): @@ -4853,7 +4861,7 @@ async def test_load_prompts_prevents_multiple_calls(): async def test_mcp_streamable_http_tool_httpx_client_cleanup(): - """Test that MCPStreamableHTTPTool delegates to caller-provided httpx clients.""" + """Owned clients are closed, while caller-provided clients remain caller-owned.""" from unittest.mock import AsyncMock, Mock, patch from agent_framework import MCPStreamableHTTPTool @@ -4876,7 +4884,6 @@ async def test_mcp_streamable_http_tool_httpx_client_cleanup(): mock_session_class.return_value.__aenter__ = AsyncMock(return_value=mock_session) mock_session_class.return_value.__aexit__ = AsyncMock(return_value=None) - # Test 1: Tool without provided client (passes None to streamable_http_client) tool1 = MCPStreamableHTTPTool( name="test", url="http://localhost:8081/mcp", @@ -4884,11 +4891,13 @@ async def test_mcp_streamable_http_tool_httpx_client_cleanup(): load_prompts=False, terminate_on_close=False, ) - await tool1.connect() - # When no client is provided, _httpx_client should be None - assert tool1._httpx_client is None, "httpx client should be None when not provided" + async with tool1: + owned_client = tool1._httpx_client + assert owned_client is not None + assert not owned_client.is_closed + assert owned_client.is_closed + assert tool1._httpx_client is None - # Test 2: Tool with user-provided client user_client = Mock() tool2 = MCPStreamableHTTPTool( name="test", @@ -4898,15 +4907,10 @@ async def test_mcp_streamable_http_tool_httpx_client_cleanup(): terminate_on_close=False, http_client=user_client, ) - await tool2.connect() - - # Verify the user-provided client was stored - assert tool2._httpx_client is user_client, "User-provided client should be stored" - - # Verify the transport wrapper delegates to the user's client. - # Get the last call (should be from tool2.connect()) - call_args = mock_client.call_args - assert call_args.kwargs["http_client"]._client is user_client + async with tool2: + assert tool2._httpx_client is user_client + assert mock_client.call_args.kwargs["http_client"]._client is user_client + user_client.aclose.assert_not_called() async def test_load_tools_with_pagination(): @@ -7219,39 +7223,56 @@ async def test_mcp_streamable_http_tool_replaces_headers_on_same_origin_redirect await tool._httpx_client.aclose() # type: ignore[union-attr] -async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client(): - """Test that header_provider works when the user provides their own httpx client.""" +@pytest.mark.parametrize("use_header_provider", [False, True]) +async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client(use_header_provider: bool): + """Supplied clients preserve configuration, cookie persistence, and ownership.""" import httpx from agent_framework._mcp import _mcp_call_headers - user_client = httpx.AsyncClient(headers={"X-Base": "static"}) - - tool = MCPStreamableHTTPTool( - name="test", - url="http://example.com/mcp", - http_client=user_client, - header_provider=lambda kw: {"X-Dynamic": kw.get("dynamic", "")}, - ) - - with patch("agent_framework._mcp.streamable_http_client"): - tool.get_mcp_client() - - # The user's client should still be used - assert tool._httpx_client is user_client - hooks = user_client.event_hooks.get("request", []) - assert len(hooks) == 1 - - # Verify the hook injects headers - token = _mcp_call_headers.set({"X-Dynamic": "per-request"}) + def handle(request: httpx.Request) -> httpx.Response: + assert request.headers["X-Base"] == "static" + assert request.headers["Authorization"].startswith("Basic ") + assert request.headers["Cookie"] == "caller_session=kept" + return httpx.Response(200, headers={"Set-Cookie": "response_session=retained; Path=/"}) + + auth = httpx.BasicAuth("caller", "test-password") + async with httpx.AsyncClient( + transport=httpx.MockTransport(handle), + headers={"X-Base": "static"}, + cookies={"caller_session": "kept"}, + auth=auth, + timeout=17.0, + ) as user_client: + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + http_client=user_client, + header_provider=(lambda kw: {"X-Dynamic": kw.get("dynamic", "")}) if use_header_provider else None, + ) try: - request = _request_for_mcp_tool(tool) - await hooks[0](request) - assert request.headers.get("X-Dynamic") == "per-request" + with patch("agent_framework._mcp.streamable_http_client"): + tool.get_mcp_client() + assert tool._httpx_client is user_client + hooks = user_client.event_hooks["request"] + assert len(hooks) == int(use_header_provider) + if use_header_provider: + token = _mcp_call_headers.set({"X-Dynamic": "per-request"}) + try: + request = _request_for_mcp_tool(tool) + await hooks[0](request) + assert request.headers.get("X-Dynamic") == "per-request" + finally: + _mcp_call_headers.reset(token) finally: - _mcp_call_headers.reset(token) - - await user_client.aclose() + await tool.close() + assert not user_client.is_closed + assert user_client.auth is auth + assert user_client.timeout == httpx.Timeout(17.0) + assert not user_client.event_hooks["request"] + await user_client.get("http://example.com/unrelated") + assert user_client.cookies["caller_session"] == "kept" + assert user_client.cookies["response_session"] == "retained" async def test_mcp_streamable_http_tool_header_provider_isolated_on_shared_httpx_client(): diff --git a/python/packages/core/tests/core/test_mcp_http_auth.py b/python/packages/core/tests/core/test_mcp_http_auth.py index 31bb5afae5c..ffb4e8348c9 100644 --- a/python/packages/core/tests/core/test_mcp_http_auth.py +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -6,7 +6,8 @@ import contextlib import json import sys -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Callable +from contextvars import ContextVar from typing import Any, Literal, TypeAlias from unittest.mock import AsyncMock, Mock, patch @@ -17,10 +18,11 @@ from agent_framework.exceptions import ToolException, ToolExecutionException MCPHTTPServer: TypeAlias = tuple[httpx.AsyncClient, list[httpx.Request], dict[str, list[str]]] +MCPHTTPClientFactory: TypeAlias = tuple[Callable[..., httpx.AsyncClient], list[httpx.Request], dict[str, list[str]]] @pytest.fixture -async def mcp_http_server() -> AsyncIterator[MCPHTTPServer]: +async def mcp_http_client_factory() -> AsyncIterator[MCPHTTPClientFactory]: requests: list[httpx.Request] = [] writes: dict[str, list[str]] = {"token-a": [], "token-b": [], "token-c": []} @@ -69,9 +71,34 @@ async def handle(request: httpx.Request) -> httpx.Response: return httpx.Response(202) return httpx.Response(200, headers=headers, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) - async with httpx.AsyncClient( - transport=httpx.MockTransport(handle), event_hooks={"request": [record_request]} - ) as client: + async_client = httpx.AsyncClient + clients: list[httpx.AsyncClient] = [] + + def create_client(*, response_cookies: dict[str, str] | None = None, **kwargs: Any) -> httpx.AsyncClient: + async def handle_with_cookies(request: httpx.Request) -> httpx.Response: + response = await handle(request) + if response_cookies and request.method == "POST": + method = json.loads(request.content).get("method") + if method in response_cookies: + response.headers["Set-Cookie"] = response_cookies[method] + return response + + client = async_client(transport=httpx.MockTransport(handle_with_cookies), **kwargs) + client.event_hooks["request"].insert(0, record_request) + clients.append(client) + return client + + try: + yield create_client, requests, writes + finally: + for client in clients: + await client.aclose() + + +@pytest.fixture +async def mcp_http_server(mcp_http_client_factory: MCPHTTPClientFactory) -> AsyncIterator[MCPHTTPServer]: + create_client, requests, writes = mcp_http_client_factory + async with create_client() as client: yield client, requests, writes @@ -93,6 +120,38 @@ def _calls(requests: list[httpx.Request]) -> list[httpx.Request]: ] +async def test_owned_client_does_not_replay_response_cookie_across_principals( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + create_client, requests, _ = mcp_http_client_factory + principal = ContextVar("principal", default="token-a") + response_cookies = {"tools/call": "regression_session=token-a; Path=/; Secure; HttpOnly"} + + def create_owned_client(**kwargs: Any) -> httpx.AsyncClient: + return create_client(response_cookies=response_cookies, **kwargs) + + tool = MCPStreamableHTTPTool( + name="owned", + url="https://mcp.example/mcp", + load_prompts=False, + header_provider=lambda _: {"Authorization": principal.get()}, + ) + with patch("httpx.AsyncClient", side_effect=create_owned_client): + async with tool: + await tool.call_tool("record") + response_cookies.clear() + token = principal.set("token-b") + try: + await tool.call_tool("record") + finally: + principal.reset(token) + + calls = _calls(requests) + assert [request.headers["Authorization"] for request in calls] == ["token-a", "token-b"] + assert "Cookie" not in calls[0].headers + assert "regression_session=token-a" not in calls[1].headers.get("Cookie", "") + + @pytest.mark.parametrize("principals", [("token-a", "token-b"), ("token-b", "token-a")]) async def test_shared_client_keeps_tools_and_caller_requests_isolated( mcp_http_server: MCPHTTPServer, principals: tuple[str, str] @@ -206,20 +265,50 @@ async def transport(**kwargs: Any) -> AsyncGenerator[tuple[()]]: await tool.close() -async def test_owned_client_is_closed_after_successful_session(mcp_http_server: MCPHTTPServer) -> None: - client, _, _ = mcp_http_server - original_hooks = list(client.event_hooks["request"]) +@pytest.mark.parametrize("header_source", ["none", "static", "provider"]) +async def test_owned_client_is_closed_after_successful_session( + mcp_http_client_factory: MCPHTTPClientFactory, header_source: str +) -> None: + create_client, requests, _ = mcp_http_client_factory + clients: list[httpx.AsyncClient] = [] + + def create_owned_client(**kwargs: Any) -> httpx.AsyncClient: + client = create_client( + response_cookies={"tools/call": "implicit_session=previous; Path=/; Secure; HttpOnly"}, + **kwargs, + ) + client.headers["Authorization"] = "token-a" + clients.append(client) + return client + tool = MCPStreamableHTTPTool( name="owned", url="https://mcp.example/mcp", load_prompts=False, - header_provider=lambda _: {"Authorization": "token-a"}, + static_headers={"Cookie": "explicit_session=caller"} if header_source == "static" else None, + header_provider=(lambda _: {"Authorization": "token-a", "Cookie": "explicit_session=caller"}) + if header_source == "provider" + else None, ) - with patch("httpx.AsyncClient", return_value=client): + with patch("httpx.AsyncClient", side_effect=create_owned_client): async with tool: await tool.call_tool("record") - assert client.is_closed - assert client.event_hooks["request"] == original_hooks + await tool.call_tool("record") + assert len(clients) == 1 + assert tool._httpx_client is clients[0] + await tool.connect(reset=True) + assert len(clients) == 2 + assert clients[0].is_closed + await tool.call_tool("record") + assert len(_calls(requests)) == 3 + assert all( + request.headers.get("Cookie") == ("explicit_session=caller" if header_source != "none" else None) + for request in _calls(requests) + ) + assert all(not client.cookies for client in clients) + assert all(client.is_closed for client in clients) + assert all(len(client.event_hooks["request"]) == 1 for client in clients) + assert tool._httpx_client is None async def test_connecting_another_tool_during_a_call_does_not_capture_its_headers( From d40fda5790f09380483e7d3abfd59d580093a596 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Mon, 14 Sep 2026 15:30:24 +0200 Subject: [PATCH 2/4] Python: Align declarative HTTP cookie persistence defaults Apply the default response-cookie policy to internally managed declarative HTTP clients. Preserve supplied and provider-returned clients, client selection and lifecycle, explicit request headers, and response headers. Document the configuration path and cover reuse, failure responses, and overlapping requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/declarative/README.md | 19 ++ .../_workflows/_http_handler.py | 17 +- .../test_default_http_request_handler.py | 199 ++++++++++++++---- .../test_http_request_yaml_integration.py | 46 ++++ 4 files changed, 233 insertions(+), 48 deletions(-) diff --git a/python/packages/declarative/README.md b/python/packages/declarative/README.md index 0ab3f8c9fc1..c536403aba4 100644 --- a/python/packages/declarative/README.md +++ b/python/packages/declarative/README.md @@ -20,3 +20,22 @@ This package ships at two different stability levels: ## Declarative features The declarative packages provides support for building agents based on a declarative yaml specification. + +## HTTP request client ownership and cookies + +**Breaking change:** The HTTP client created by `DefaultHttpRequestHandler` no longer +persists response cookies. The handler still creates its client lazily, reuses it across +requests and workflows, and closes it on `aclose()` or async context-manager exit. +Timeout and redirect defaults are unchanged. + +Applications requiring cookies for authentication, session continuity, or load-balancer +affinity must supply an `httpx.AsyncClient` through `client=` or `client_provider=`. +Supplied and provider-returned clients retain their configuration and cookie behavior +and must be closed by the caller. Scope cookie-bearing clients to one authenticated +principal; sharing a handler across workflows does not partition a client's cookies. +A provider returning `None` falls back to `client=`, if supplied, then to the internally +owned client with the default cookie policy. + +Explicit `Cookie` request headers remain supported. Response `Set-Cookie` headers +remain available in `HttpRequestResult.headers`; disabling persistence does not redact +the response. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_http_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_http_handler.py index f52cc73eff0..3dff14d17d6 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_http_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_http_handler.py @@ -23,6 +23,7 @@ import asyncio from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, field +from http.cookiejar import CookieJar, DefaultCookiePolicy from typing import Any, Protocol, runtime_checkable from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit @@ -122,12 +123,18 @@ class DefaultHttpRequestHandler: Construction modes: 1. ``DefaultHttpRequestHandler()`` — owns an internal client created lazily - on first ``send()``. Closed by :meth:`aclose`. + on first ``send()`` without response-cookie persistence. Closed by :meth:`aclose`. 2. ``DefaultHttpRequestHandler(client=existing)`` — caller-owned client. - Not closed by :meth:`aclose`. + Retains its cookie behavior and is not closed by :meth:`aclose`. 3. ``DefaultHttpRequestHandler(client_provider=cb)`` — per-request client lookup (parity with .NET's ``httpClientProvider`` callback). The - provider may return ``None`` to fall back to the owned/default client. + provider's clients retain their cookie behavior and are not closed by :meth:`aclose`. + Returning ``None`` falls back to ``client``, if supplied, then to the owned client. + + Applications requiring persistent cookies must supply a client through ``client`` + or ``client_provider`` scoped to one authenticated principal and manage its lifetime. + Explicit outbound ``Cookie`` headers and response ``Set-Cookie`` headers are preserved; + the owned-client policy only prevents automatic cookie persistence. .. warning:: @@ -262,7 +269,9 @@ async def _resolve_client(self, info: HttpRequestInfo) -> httpx.AsyncClient: # one of them. async with self._owned_client_lock: if self._owned_client is None: - self._owned_client = httpx.AsyncClient() + self._owned_client = httpx.AsyncClient( + cookies=CookieJar(policy=DefaultCookiePolicy(allowed_domains=[])), + ) return self._owned_client async def __aenter__(self) -> DefaultHttpRequestHandler: diff --git a/python/packages/declarative/tests/test_default_http_request_handler.py b/python/packages/declarative/tests/test_default_http_request_handler.py index 1d4e393560d..3db80fd943e 100644 --- a/python/packages/declarative/tests/test_default_http_request_handler.py +++ b/python/packages/declarative/tests/test_default_http_request_handler.py @@ -11,7 +11,11 @@ from __future__ import annotations +import asyncio import sys +from collections.abc import AsyncIterator +from typing import Any +from unittest.mock import patch import httpx import pytest @@ -387,16 +391,87 @@ def respond(request: httpx.Request) -> httpx.Response: class TestClientOwnership: - @pytest.mark.asyncio async def test_owned_client_is_closed_on_aclose(self) -> None: + requests: list[httpx.Request] = [] + clients: list[httpx.AsyncClient] = [] + original_ctor = httpx.AsyncClient + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, headers={"Set-Cookie": "session=previous; Path=/; Secure"}) + + def create_client(**kwargs: Any) -> httpx.AsyncClient: + client = original_ctor(transport=httpx.MockTransport(respond), **kwargs) + clients.append(client) + return client + handler = DefaultHttpRequestHandler() - # Inject a MockTransport-backed client into the owned slot and verify - # aclose() releases it. Avoids real network access. - owned = httpx.AsyncClient(transport=httpx.MockTransport(lambda r: httpx.Response(200, text="ok"))) - handler._owned_client = owned - assert not owned.is_closed - await handler.aclose() - assert owned.is_closed + info = HttpRequestInfo(method="GET", url="https://api.example.test/x") + with patch("httpx.AsyncClient", side_effect=create_client): + for index in range(2): + async with handler: + assert handler._owned_client is None + await handler.send(info) + await handler.send(info) + assert len(clients) == index + 1 + assert handler._owned_client is clients[index] + assert not clients[index].is_closed + assert clients[index].timeout == httpx.Timeout(5.0) + assert not clients[index].follow_redirects + assert handler._owned_client is None + assert clients[index].is_closed + await handler.aclose() + assert clients[0].cookies.jar is not clients[1].cookies.jar + assert all(not client.cookies for client in clients) + assert all("Cookie" not in request.headers for request in requests) + + @pytest.mark.parametrize("provider_fallback", [False, True]) + @pytest.mark.parametrize("first_status", [200, 503]) + async def test_owned_client_does_not_replay_response_cookies( + self, provider_fallback: bool, first_status: int + ) -> None: + requests: list[httpx.Request] = [] + response_cookies = ["session=first; Path=/; HttpOnly; Secure", "affinity=first; Path=/; Secure"] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + first_status if len(requests) == 1 else 200, + text="ok", + headers=[("Set-Cookie", cookie) for cookie in response_cookies] if len(requests) == 1 else [], + ) + + async def provider(info: HttpRequestInfo) -> httpx.AsyncClient | None: + return None + + original_ctor = httpx.AsyncClient + + def create_client(**kwargs: Any) -> httpx.AsyncClient: + return original_ctor(transport=httpx.MockTransport(respond), **kwargs) + + with patch("httpx.AsyncClient", side_effect=create_client): + async with DefaultHttpRequestHandler(client_provider=provider if provider_fallback else None) as handler: + first = await handler.send( + HttpRequestInfo( + method="GET", url="https://api.example.test/x", headers={"Authorization": "caller-a"} + ) + ) + second = await handler.send( + HttpRequestInfo( + method="GET", url="https://api.example.test/x", headers={"Authorization": "caller-b"} + ) + ) + await handler.send( + HttpRequestInfo( + method="GET", url="https://api.example.test/x", headers={"Cookie": "explicit=caller"} + ) + ) + assert first.status_code == first_status + assert first.is_success_status_code is (first_status == 200) + assert first.headers["set-cookie"] == response_cookies + assert second.is_success_status_code + assert [request.headers.get("Authorization") for request in requests[:2]] == ["caller-a", "caller-b"] + assert [request.headers.get("Cookie") for request in requests] == [None, None, "explicit=caller"] @pytest.mark.asyncio async def test_caller_owned_client_is_not_closed(self) -> None: @@ -407,7 +482,6 @@ async def test_caller_owned_client_is_not_closed(self) -> None: assert not client.is_closed await client.aclose() # cleanup - @pytest.mark.asyncio async def test_concurrent_first_send_creates_single_owned_client(self) -> None: """Concurrent first-send calls must not race-leak duplicate clients. @@ -417,38 +491,50 @@ async def test_concurrent_first_send_creates_single_owned_client(self) -> None: is serialized: all concurrent sends end up using the same client and ``aclose()`` cleanly closes it. """ - import asyncio - - # Patch httpx.AsyncClient to count constructions, but only when called - # from inside _resolve_client (no transport=) so we don't break the - # MockTransport-backed clients used elsewhere. original_ctor = httpx.AsyncClient - construction_count = 0 + clients: list[httpx.AsyncClient] = [] + requests: list[httpx.Request] = [] + first_headers_received = asyncio.Event() + second_request_received = asyncio.Event() + + class PendingResponse(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + # HTTPX processes response cookies before reading the body. + first_headers_received.set() + await second_request_received.wait() + yield b"ok" - def counting_ctor(*args, **kwargs): # type: ignore[no-untyped-def] - nonlocal construction_count - if not args and not kwargs: - construction_count += 1 - return original_ctor(transport=httpx.MockTransport(lambda r: httpx.Response(200, text="ok"))) - return original_ctor(*args, **kwargs) + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + if len(requests) == 1: + return httpx.Response( + 200, + headers={"Set-Cookie": "session=first; Path=/; HttpOnly; Secure"}, + stream=PendingResponse(), + ) + second_request_received.set() + return httpx.Response(200, text="ok") - import agent_framework_declarative._workflows._http_handler as hh + def create_client(**kwargs: Any) -> httpx.AsyncClient: + client = original_ctor(transport=httpx.MockTransport(respond), **kwargs) + clients.append(client) + return client - hh.httpx.AsyncClient = counting_ctor # type: ignore[assignment] # ty: ignore[invalid-assignment] - try: - handler = DefaultHttpRequestHandler() - try: - await asyncio.gather(*[ - handler.send(HttpRequestInfo(method="GET", url="https://api.example.test/x")) for _ in range(8) - ]) - finally: - await handler.aclose() - finally: - hh.httpx.AsyncClient = original_ctor # type: ignore[assignment] + with patch("httpx.AsyncClient", side_effect=create_client): + async with DefaultHttpRequestHandler() as handler: - assert construction_count == 1, ( - f"Expected exactly 1 owned client to be lazily created but got {construction_count}" - ) + async def send(index: int) -> None: + if index: + await first_headers_received.wait() + result = await handler.send(HttpRequestInfo(method="GET", url="https://api.example.test/x")) + assert result.body == "ok" + + await asyncio.wait_for(asyncio.gather(*(send(index) for index in range(8))), timeout=5) + assert len(clients) == 1 + assert handler._owned_client is clients[0] + assert clients[0].is_closed + assert len(requests) == 8 + assert all("Cookie" not in request.headers for request in requests) class TestClientProvider: @@ -477,6 +563,8 @@ async def provider(info: HttpRequestInfo) -> httpx.AsyncClient: assert captured["transport"] == "provided" finally: await handler.aclose() + assert not primary_client.is_closed + assert not provided_client.is_closed await primary_client.aclose() await provided_client.aclose() @@ -634,23 +722,33 @@ def respond(request: httpx.Request) -> httpx.Response: assert captured["req"].url.raw_path == (b"/s?in-url=fromurl&in-explicit=fromexplicit&unique=kept") - @pytest.mark.asyncio - async def test_client_headers_cookies_auth_and_timeout_survive(self) -> None: + @pytest.mark.parametrize("client_source", ["caller", "provider", "provider_none"]) + async def test_client_headers_cookies_auth_and_timeout_survive(self, client_source: str) -> None: """Building through the client keeps its configuration, params merge bypassed.""" - captured: dict[str, httpx.Request] = {} + requests: list[httpx.Request] = [] def respond(request: httpx.Request) -> httpx.Response: - captured["req"] = request - return httpx.Response(200, text="ok") + requests.append(request) + return httpx.Response(200, text="ok", headers={"Set-Cookie": "response=retained; Path=/; Secure"}) + auth = httpx.BasicAuth("user", "pass") client = httpx.AsyncClient( params={"api-version": "2024-01-01"}, headers={"X-Client-Header": "client-value"}, cookies={"session": "cookie-value"}, - auth=("user", "pass"), + auth=auth, + timeout=17.0, + follow_redirects=True, transport=httpx.MockTransport(respond), ) - handler = DefaultHttpRequestHandler(client=client) + + async def provider(info: HttpRequestInfo) -> httpx.AsyncClient | None: + return client if client_source == "provider" else None + + handler = DefaultHttpRequestHandler( + client=client if client_source != "provider" else None, + client_provider=provider if client_source != "caller" else None, + ) try: await handler.send( HttpRequestInfo( @@ -659,11 +757,24 @@ def respond(request: httpx.Request) -> httpx.Response: timeout_ms=4500, ) ) + await handler.send(HttpRequestInfo(method="GET", url="https://api.example.test/s")) + await handler.aclose() + assert not client.is_closed + assert handler._owned_client is None + assert client.cookies["session"] == "cookie-value" + assert client.cookies["response"] == "retained" + assert client.auth is auth + assert client.timeout == httpx.Timeout(17.0) + assert client.follow_redirects + assert set(requests[1].headers["Cookie"].split("; ")) == { + "session=cookie-value", + "response=retained", + } finally: await handler.aclose() await client.aclose() - request = captured["req"] + request = requests[0] assert request.headers["X-Client-Header"] == "client-value" assert request.headers["Cookie"] == "session=cookie-value" assert request.headers["Authorization"].startswith("Basic ") diff --git a/python/packages/declarative/tests/test_http_request_yaml_integration.py b/python/packages/declarative/tests/test_http_request_yaml_integration.py index 68fb22cb050..1887b6787db 100644 --- a/python/packages/declarative/tests/test_http_request_yaml_integration.py +++ b/python/packages/declarative/tests/test_http_request_yaml_integration.py @@ -12,7 +12,9 @@ import sys from pathlib import Path from typing import Any +from unittest.mock import patch +import httpx import pytest try: @@ -36,6 +38,7 @@ from agent_framework_declarative import WorkflowFactory # noqa: E402 from agent_framework_declarative._workflows import DECLARATIVE_STATE_KEY # noqa: E402 from agent_framework_declarative._workflows._http_handler import ( # noqa: E402 + DefaultHttpRequestHandler, HttpRequestInfo, HttpRequestResult, ) @@ -98,6 +101,49 @@ async def test_http_request_yaml_roundtrip() -> None: assert sent.headers["User-Agent"] == "agent-framework-integration-test" +async def test_http_request_yaml_shared_handler_does_not_replay_cookies() -> None: + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + text="ok", + headers={"Set-Cookie": "session=first; Path=/; HttpOnly; Secure"} if len(requests) == 1 else {}, + ) + + original_ctor = httpx.AsyncClient + + def create_client(**kwargs: Any) -> httpx.AsyncClient: + return original_ctor(transport=httpx.MockTransport(respond), **kwargs) + + definition = """ +kind: Workflow +trigger: + kind: OnConversationStart + id: http_cookie_test + actions: + - kind: HttpRequestAction + id: request + method: GET + url: https://api.example.test/x + headers: + Authorization: caller-a + response: Local.Response +""" + with patch("httpx.AsyncClient", side_effect=create_client): + async with DefaultHttpRequestHandler() as handler: + factory = WorkflowFactory(http_request_handler=handler) + first = factory.create_workflow_from_yaml(definition) + second = factory.create_workflow_from_yaml(definition.replace("caller-a", "caller-b")) + await first.run({}) + client = handler._owned_client + await second.run({}) + assert handler._owned_client is client + assert [request.headers["Authorization"] for request in requests] == ["caller-a", "caller-b"] + assert all("Cookie" not in request.headers for request in requests) + + @pytest.mark.asyncio async def test_http_request_yaml_missing_handler_fails_at_build_time() -> None: """Without an http_request_handler, building the workflow must raise.""" From 2eb7dd96a0f5b87ae5569dfd9e31bbbe92a0fc5f Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Tue, 15 Sep 2026 10:37:09 +0200 Subject: [PATCH 3/4] Python: Preserve caller-owned A2A HTTP clients Leave HTTP clients supplied alongside an A2A protocol client open on agent exit. Extend the existing ownership regression, retain internally owned cleanup coverage, and document caller cleanup responsibility. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/a2a/README.md | 6 ++++- .../a2a/agent_framework_a2a/_agent.py | 7 +++-- python/packages/a2a/tests/test_a2a_agent.py | 27 +++++-------------- 3 files changed, 15 insertions(+), 25 deletions(-) diff --git a/python/packages/a2a/README.md b/python/packages/a2a/README.md index 0fe9a961560..b21836751c8 100644 --- a/python/packages/a2a/README.md +++ b/python/packages/a2a/README.md @@ -24,13 +24,17 @@ response = await a2a_agent.run("Hello!") **Breaking change:** HTTP clients created by `A2AAgent` no longer persist response cookies. This prevents one run's cookies from being reused by another run through the same agent. -Connection pooling and the client's existing lifetime are unchanged. +Connection pooling and cleanup of internally created clients are unchanged. If the remote service requires cookies for authentication, sessions, or load-balancer affinity, explicitly supply an `httpx.AsyncClient` using `http_client=`. Its cookie behavior is preserved. Scope that client to the intended authenticated user or security context and manage its lifetime; do not share a user-specific cookie jar across users. +Supplied HTTP clients remain caller-owned even when `client=` is also provided. +Applications that previously relied on the agent closing a supplied HTTP client +in that combination must now close it explicitly. + ### A2AExecutor (Hosting) The `A2AExecutor` class bridges local AI agents built with the `agent_framework` library to the A2A protocol, allowing them to be hosted and accessed by other A2A-compliant clients. diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 30912ae3b86..5dca644f184 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -246,7 +246,8 @@ def __init__( url: The URL for the A2A server. client: The A2A client for the agent. http_client: Optional httpx.AsyncClient to use. Agent-created clients do not persist - response cookies; supplied clients retain their configured cookie behavior. + response cookies; supplied clients retain their configured cookie behavior and + remain caller-owned, including when ``client`` is also supplied. auth_interceptor: Optional authentication interceptor for secured endpoints. timeout: Request timeout configuration. Can be a float (applied to all timeout components), httpx.Timeout object (for full control), or None (uses 10.0s connect, 60.0s read, @@ -265,15 +266,13 @@ def __init__( super().__init__(id=id, name=name, description=description, **kwargs) self._http_client: httpx.AsyncClient | None = http_client - # By default, leave a caller-supplied HTTP client open; specific paths may override this below. - # every construction path must set this before __aexit__ can run + # Only HTTP clients created by this agent are closed on exit. self._close_http_client = False self._timeout_config = self._create_timeout_config(timeout) bindings = supported_protocol_bindings if supported_protocol_bindings is not None else ["JSONRPC"] if client is not None: self.client = client self._non_streaming_client: Client | None = None - self._close_http_client = True return if agent_card is None: if url is None: diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 7c8dc8d5b57..b7bb617db90 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -501,24 +501,6 @@ async def test_run_streaming_with_message_response(a2a_agent: A2AAgent, mock_a2a assert mock_a2a_client.call_count == 1 -async def test_context_manager_cleanup() -> None: - """Test context manager cleanup of http client.""" - - # Create mock http client that tracks aclose calls - mock_http_client = AsyncMock() - mock_a2a_client = MagicMock() - - agent = A2AAgent(client=cast(Any, mock_a2a_client)) - agent._http_client = mock_http_client - - # Test context manager cleanup - async with agent: - pass - - # Verify aclose was called - mock_http_client.aclose.assert_called_once() - - async def test_context_manager_no_cleanup_when_no_http_client() -> None: """Test context manager when _http_client is None.""" @@ -531,7 +513,8 @@ async def test_context_manager_no_cleanup_when_no_http_client() -> None: pass -async def test_context_manager_does_not_close_caller_supplied_http_client() -> None: +@mark.parametrize("supply_a2a_client", [False, True]) +async def test_context_manager_does_not_close_caller_supplied_http_client(supply_a2a_client: bool) -> None: """A caller-supplied HTTP client must stay open and retain its cookie behavior after __aexit__.""" def handle_request(request: httpx.Request) -> httpx.Response: @@ -541,7 +524,11 @@ def handle_request(request: httpx.Request) -> httpx.Response: async with httpx.AsyncClient( transport=httpx.MockTransport(handle_request), cookies={"caller_cookie": "kept"} ) as http_client: - async with A2AAgent(url="https://a2a.example.test/", http_client=http_client): + async with A2AAgent( + url="https://a2a.example.test/", + client=MagicMock() if supply_a2a_client else None, + http_client=http_client, + ): pass assert not http_client.is_closed From 71894793f5f3bff0c8ab7099217039cbe3135cd1 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Tue, 15 Sep 2026 10:46:48 +0200 Subject: [PATCH 4/4] Python: Clarify AG-UI test types Type conditional exception contexts explicitly and use separate variables for chat and HTTP responses. Preserve the existing test cases and assertions while satisfying Mypy and Zuban. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../ag-ui/tests/ag_ui/test_ag_ui_client.py | 22 ++++++++++++------- .../ag-ui/tests/ag_ui/test_workflow_run.py | 4 +--- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index 65298dde9c6..679cfb11eec 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -5,7 +5,7 @@ import asyncio import json from collections.abc import AsyncGenerator, Awaitable, Mapping, MutableSequence -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext from typing import Any, cast import httpx @@ -114,7 +114,7 @@ async def handle_async_request( http_client = client.http_service.http_client for index, thread_id in enumerate(("thread_a", second_thread_id)): messages = [Message(role="user", contents=["Hello"])] - expected_error = ( + expected_error: AbstractContextManager[object] = ( pytest.raises(httpx.HTTPStatusError, match="503") if index == 0 and first_status == 503 else nullcontext() @@ -220,8 +220,11 @@ async def test_context_manager_closes_owned_http_client(self, raise_in_context: """Owned HTTP clients close on both normal and exceptional context exits.""" client = StubAGUIChatClient(endpoint="https://agui.example.test/", timeout=17.0) http_client = client.http_service.http_client + expected_error: AbstractContextManager[object] = ( + pytest.raises(RuntimeError, match="context failed") if raise_in_context else nullcontext() + ) try: - with pytest.raises(RuntimeError, match="context failed") if raise_in_context else nullcontext(): + with expected_error: async with client: assert not http_client.is_closed assert http_client.timeout == httpx.Timeout(17.0) @@ -261,14 +264,17 @@ async def handler(request: httpx.Request) -> httpx.Response: timeout=17.0, ) as http_client: client = StubAGUIChatClient(endpoint="https://agui.example.test/", http_client=http_client) - with pytest.raises(RuntimeError, match="context failed") if raise_in_context else nullcontext(): + expected_error: AbstractContextManager[object] = ( + pytest.raises(RuntimeError, match="context failed") if raise_in_context else nullcontext() + ) + with expected_error: async with client: assert client.http_service.http_client is http_client - response = await client.get_response( + chat_response = await client.get_response( [Message(role="user", contents=["Hello"])], options={"metadata": {"thread_id": "thread_1"}}, ) - assert response.text == "ok" + assert chat_response.text == "ok" if raise_in_context: raise RuntimeError("context failed") @@ -278,8 +284,8 @@ async def handler(request: httpx.Request) -> httpx.Response: assert http_client.cookies["configured"] == "caller-owned" assert http_client.cookies["response"] == "session-only" assert http_client.timeout == httpx.Timeout(17.0) - response = await http_client.get("https://agui.example.test/health") - assert response.status_code == 200 + http_response = await http_client.get("https://agui.example.test/health") + assert http_response.status_code == 200 assert [request.method for request in requests] == ["POST", "GET"] assert all(request.headers["authorization"] == "Bearer test-token" for request in requests) diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index be6213d2198..d9cb44fe01c 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -481,9 +481,7 @@ async def test_workflow_run_approval_pause_closes_open_real_tool_call() -> None: """ @executor(id="approval_with_streamed_tool") - async def approval_with_streamed_tool( - message: Any, ctx: WorkflowContext[Any, AgentResponseUpdate] - ) -> None: + async def approval_with_streamed_tool(message: Any, ctx: WorkflowContext[Any, AgentResponseUpdate]) -> None: del message function_call = Content.from_function_call( call_id="weather-call",