mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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:
committed by
GitHub
Unverified
parent
f5f909f2e2
commit
2d3ba95036
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user