Python: Add configurable timeout support to A2AAgent (#2432)

* a2a timeout config

* added default timeout info

---------

Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
This commit is contained in:
Giles Odigwe
2025-12-05 14:21:36 -08:00
committed by GitHub
Unverified
parent f5f909f2e2
commit 2d3ba95036
2 changed files with 81 additions and 7 deletions
@@ -80,6 +80,7 @@ class A2AAgent(BaseAgent):
client: Client | None = None,
http_client: httpx.AsyncClient | None = None,
auth_interceptor: AuthInterceptor | None = None,
timeout: float | httpx.Timeout | None = None,
**kwargs: Any,
) -> None:
"""Initialize the A2AAgent.
@@ -93,10 +94,14 @@ class A2AAgent(BaseAgent):
client: The A2A client for the agent.
http_client: Optional httpx.AsyncClient to use.
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,
10.0s write, 5.0s pool - optimized for A2A operations).
kwargs: any additional properties, passed to BaseAgent.
"""
super().__init__(id=id, name=name, description=description, **kwargs)
self._http_client: httpx.AsyncClient | None = http_client
self._timeout_config = self._create_timeout_config(timeout)
if client is not None:
self.client = client
self._close_http_client = True
@@ -109,14 +114,8 @@ class A2AAgent(BaseAgent):
# Create or use provided httpx client
if http_client is None:
timeout = httpx.Timeout(
connect=10.0, # 10 seconds to establish connection
read=60.0, # 60 seconds to read response (A2A operations can take time)
write=10.0, # 10 seconds to send request
pool=5.0, # 5 seconds to get connection from pool
)
headers = prepend_agent_framework_to_user_agent()
http_client = httpx.AsyncClient(timeout=timeout, headers=headers)
http_client = httpx.AsyncClient(timeout=self._timeout_config, headers=headers)
self._http_client = http_client # Store for cleanup
self._close_http_client = True
@@ -143,6 +142,32 @@ class A2AAgent(BaseAgent):
f"Fallback error: {fallback_error}"
) from transport_error
def _create_timeout_config(self, timeout: float | httpx.Timeout | None) -> httpx.Timeout:
"""Create httpx.Timeout configuration from user input.
Args:
timeout: User-provided timeout configuration
Returns:
Configured httpx.Timeout object
"""
if timeout is None:
# Default timeout configuration (preserving original values)
return httpx.Timeout(
connect=10.0, # 10 seconds to establish connection
read=60.0, # 60 seconds to read response (A2A operations can take time)
write=10.0, # 10 seconds to send request
pool=5.0, # 5 seconds to get connection from pool
)
if isinstance(timeout, float):
# Simple timeout
return httpx.Timeout(timeout)
if isinstance(timeout, httpx.Timeout):
# Full timeout configuration provided by user
return timeout
msg = f"Invalid timeout type: {type(timeout)}. Expected float, httpx.Timeout, or None."
raise TypeError(msg)
async def __aenter__(self) -> "A2AAgent":
"""Async context manager entry."""
return self
@@ -5,6 +5,7 @@ from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4
import httpx
from a2a.types import (
AgentCard,
Artifact,
@@ -554,3 +555,51 @@ def test_transport_negotiation_both_fail() -> None:
name="test-agent",
agent_card=mock_agent_card,
)
def test_create_timeout_config_httpx_timeout() -> None:
"""Test _create_timeout_config with httpx.Timeout object returns it unchanged."""
agent = A2AAgent(name="Test Agent", client=MockA2AClient(), http_client=None)
custom_timeout = httpx.Timeout(connect=15.0, read=180.0, write=20.0, pool=8.0)
timeout_config = agent._create_timeout_config(custom_timeout)
assert timeout_config is custom_timeout # Same object reference
assert timeout_config.connect == 15.0
assert timeout_config.read == 180.0
assert timeout_config.write == 20.0
assert timeout_config.pool == 8.0
def test_create_timeout_config_invalid_type() -> None:
"""Test _create_timeout_config with invalid type raises TypeError."""
agent = A2AAgent(name="Test Agent", client=MockA2AClient(), http_client=None)
with raises(TypeError, match="Invalid timeout type: <class 'str'>. Expected float, httpx.Timeout, or None."):
agent._create_timeout_config("invalid")
def test_a2a_agent_initialization_with_timeout_parameter() -> None:
"""Test A2AAgent initialization with timeout parameter."""
# Test with URL to trigger httpx client creation
with (
patch("agent_framework_a2a._agent.httpx.AsyncClient") as mock_async_client,
patch("agent_framework_a2a._agent.ClientFactory") as mock_factory,
):
# Mock the factory and client creation
mock_client_instance = MagicMock()
mock_factory.return_value.create.return_value = mock_client_instance
# Create agent with custom timeout
A2AAgent(name="Test Agent", url="https://test-agent.example.com", timeout=120.0)
# Verify httpx.AsyncClient was called with the configured timeout
mock_async_client.assert_called_once()
call_args = mock_async_client.call_args
# Check that timeout parameter was passed
assert "timeout" in call_args.kwargs
timeout_arg = call_args.kwargs["timeout"]
# Verify it's an httpx.Timeout object with our custom timeout applied to all components
assert isinstance(timeout_arg, httpx.Timeout)