fix: Retain the consumers and transport registry when recreating the ClientFactory in remote_a2a_agent.py

PiperOrigin-RevId: 807762203
This commit is contained in:
Google Team Member
2025-09-16 10:58:46 -07:00
committed by Copybara-Service
parent f7bd3c111c
commit 6bd33e1be3
2 changed files with 50 additions and 1 deletions
+5 -1
View File
@@ -188,12 +188,16 @@ class RemoteA2aAgent(BaseAgent):
)
self._httpx_client_needs_cleanup = True
if self._a2a_client_factory:
registry = self._a2a_client_factory._registry
self._a2a_client_factory = A2AClientFactory(
config=dataclasses.replace(
self._a2a_client_factory._config,
httpx_client=self._httpx_client,
)
),
consumers=self._a2a_client_factory._consumers,
)
for label, generator in registry.items():
self._a2a_client_factory.register(label, generator)
if not self._a2a_client_factory:
client_config = A2AClientConfig(
httpx_client=self._httpx_client,
@@ -34,6 +34,7 @@ pytestmark = pytest.mark.skipif(
# Import dependencies with version checking
try:
from a2a.client.client import ClientConfig
from a2a.client.client import Consumer
from a2a.client.client_factory import ClientFactory
from a2a.types import AgentCapabilities
from a2a.types import AgentCard
@@ -246,6 +247,50 @@ class TestRemoteA2aAgentResolution:
assert client == existing_client
assert agent._httpx_client_needs_cleanup is False
@pytest.mark.asyncio
async def test_ensure_httpx_client_updates_factory_with_new_client(self):
"""Test that _ensure_httpx_client updates factory with new client."""
agent = RemoteA2aAgent(
name="test_agent",
agent_card=create_test_agent_card(),
a2a_client_factory=ClientFactory(
ClientConfig(httpx_client=None),
),
)
assert agent._a2a_client_factory._config.httpx_client is None
client = await agent._ensure_httpx_client()
assert client is not None
assert agent._httpx_client == client
assert agent._httpx_client_needs_cleanup is True
assert agent._a2a_client_factory._config.httpx_client == client
@pytest.mark.asyncio
async def test_ensure_httpx_client_reregisters_transports_with_new_client(
self,
):
"""Test that _ensure_httpx_client registers transports with new client."""
factory = ClientFactory(
ClientConfig(httpx_client=None),
)
factory.register("transport_label", lambda: "test")
agent = RemoteA2aAgent(
name="test_agent",
agent_card=create_test_agent_card(),
a2a_client_factory=factory,
)
assert agent._a2a_client_factory._config.httpx_client is None
assert "transport_label" in agent._a2a_client_factory._registry
client = await agent._ensure_httpx_client()
assert client is not None
assert agent._httpx_client == client
assert agent._httpx_client_needs_cleanup is True
assert agent._a2a_client_factory._config.httpx_client == client
assert "transport_label" in agent._a2a_client_factory._registry
@pytest.mark.asyncio
async def test_resolve_agent_card_from_url_success(self):
"""Test successful agent card resolution from URL."""