Python: Add configurable timeout support to A2AAgent

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,13 @@ 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.
+                httpx.Timeout object (for full control), or None (use defaults).
             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 +113,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 +141,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

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: <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)