diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 3d5ac41e86b..1261305bfbc 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -263,6 +263,8 @@ def __init__( super().__init__(id=id, name=name, description=description, **kwargs) self._http_client: httpx.AsyncClient | None = http_client + # every construction path must set this before __aexit__ can run + 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: diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 8ff147f98aa..9c111a2b11e 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -2536,3 +2536,37 @@ async def test_input_required_sets_task_id_instead_of_reference(mock_a2a_client: # endregion + + +async def test_context_manager_with_user_supplied_http_client() -> None: + """A2AAgent(url=..., http_client=mine) must exit cleanly and must not close it. + + The third construction path (no client, caller-provided http_client) used to + leave _close_http_client unset, so __aexit__ raised AttributeError. + """ + mock_http_client = MagicMock() + mock_http_client.aclose = AsyncMock() + + agent = A2AAgent(url="http://agent.example", http_client=mock_http_client) + + async with agent: + pass + + # the client belongs to the caller; we must not close it + mock_http_client.aclose.assert_not_called() + + +async def test_context_manager_closes_self_created_http_client() -> None: + """A2AAgent(url=...) builds its own client and closes it on exit.""" + + with patch("agent_framework_a2a._agent.httpx.AsyncClient") as cls: + mock_http_client = MagicMock() + mock_http_client.aclose = AsyncMock() + cls.return_value = mock_http_client + + agent = A2AAgent(url="http://agent.example") + + async with agent: + pass + + mock_http_client.aclose.assert_called_once()