diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 7fe4649a65..4f86eb5afc 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -80,6 +80,7 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( # 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 @@ def __init__( 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 diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index f703c92a96..82d3d02875 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -5,6 +5,7 @@ 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: . 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)