refactor: 重写http client传入模型的方式

Author: raychen911

tests/models/test_anthropic_model.py

@@ -710,13 +710,13 @@ async def test_generate_single_error_raises_and_closes_client(self):
         model = AnthropicModel(model_name="claude-3-5-sonnet-20241022", api_key="test-key")
         client = MagicMock()
         client.messages.create = AsyncMock(side_effect=TimeoutError("timeout"))
-        client.close = AsyncMock()
+        model._http_client_provider.close_http_client = AsyncMock()
 
         with patch.object(model, "_create_async_client", return_value=client):
             with pytest.raises(TimeoutError):
                 await model._generate_single({}, LlmRequest(contents=[]))
 
-        client.close.assert_awaited_once()
+        model._http_client_provider.close_http_client.assert_awaited_once_with(client)
 
     @pytest.mark.asyncio
     async def test_generate_async_converts_provider_exception_to_retry_error_response(self):

tests/models/test_anthropic_model_ext.py

@@ -472,11 +472,11 @@ async def test_api_error_raises_and_closes_client(self):
         model = _model()
         mock_client = AsyncMock()
         mock_client.messages.create = AsyncMock(side_effect=RuntimeError("timeout"))
-        mock_client.close = AsyncMock()
+        model._http_client_provider.close_http_client = AsyncMock()
         with patch.object(model, "_create_async_client", return_value=mock_client):
             with pytest.raises(RuntimeError, match="timeout"):
                 await model._generate_single({}, LlmRequest(contents=[]))
-            mock_client.close.assert_awaited_once()
+            model._http_client_provider.close_http_client.assert_awaited_once_with(mock_client)
 
 
 # ---------------------------------------------------------------------------

tests/models/test_openai_model.py

@@ -8,8 +8,6 @@
 from unittest.mock import Mock
 from unittest.mock import patch
 
-import httpx
-import openai
 import pytest
 from trpc_agent_sdk.models import LlmRequest
 from trpc_agent_sdk.models import OpenAIModel
@@ -242,23 +240,26 @@ def test_model_type_is_model(self):
 
         assert model._type == FilterType.MODEL
 
-    def test_create_async_client_uses_custom_http_client_factory(self):
-        """A custom http_client_factory is passed through to AsyncOpenAI."""
+    def test_create_async_client_uses_custom_http_client_provider(self):
+        """A custom http_client_provider_factory is passed through to AsyncOpenAI."""
         shared_http_client = Mock()
-        http_client_factory = Mock(return_value=shared_http_client)
+        http_client_provider = Mock()
+        http_client_provider.create_http_client.return_value = shared_http_client
+        http_client_provider_factory = Mock(return_value=http_client_provider)
         model = OpenAIModel(
             model_name="gpt-4",
             api_key="test_key",
             base_url="https://custom.api.com",
             client_args={"timeout": 30},
-            http_client_factory=http_client_factory,
+            http_client_provider_factory=http_client_provider_factory,
         )
 
         with patch("trpc_agent_sdk.models._openai_model.openai.AsyncOpenAI") as mock_async_openai:
             client = model._create_async_client()
 
         assert client is mock_async_openai.return_value
-        http_client_factory.assert_called_once_with()
+        http_client_provider_factory.assert_called_once_with()
+        http_client_provider.create_http_client.assert_called_once_with()
         mock_async_openai.assert_called_once_with(
             api_key="test_key",
             max_retries=0,
@@ -268,39 +269,129 @@ def test_create_async_client_uses_custom_http_client_factory(self):
             http_client=shared_http_client,
         )
 
-    def test_create_async_client_default_factory_reuses_shared_http_client(self):
-        """Default factory should reuse one shared httpx.AsyncClient across model calls."""
-        from trpc_agent_sdk.models import _openai_model
+    def test_create_async_client_default_factory_reuses_loop_local_http_client(self):
+        """Default provider should reuse the same httpx.AsyncClient within one loop."""
+        from trpc_agent_sdk.models import _httpx_client
 
-        _openai_model._shared_http_client = None
         shared_http_client = Mock()
+        shared_http_client.is_closed = False
         model = OpenAIModel(model_name="gpt-4", api_key="test_key")
 
         try:
-            with patch("trpc_agent_sdk.models._openai_model.httpx.AsyncClient",
+            _httpx_client._shared_http_clients.clear()
+            with patch("trpc_agent_sdk.models._httpx_client.httpx.AsyncClient",
                        return_value=shared_http_client) as mock_httpx_client:
-                with patch("trpc_agent_sdk.models._openai_model.openai.AsyncOpenAI") as mock_async_openai:
-                    model._create_async_client()
-                    model._create_async_client()
+                with patch("trpc_agent_sdk.models._httpx_client._get_loop_key", return_value=1):
+                    with patch("trpc_agent_sdk.models._openai_model.openai.AsyncOpenAI") as mock_async_openai:
+                        model._create_async_client()
+                        model._create_async_client()
         finally:
-            _openai_model._shared_http_client = None
+            _httpx_client._shared_http_clients.clear()
 
-        mock_httpx_client.assert_called_once_with()
+        mock_httpx_client.assert_called_once_with(
+            limits=_httpx_client._DEFAULT_HTTP_CLIENT_LIMITS,
+            timeout=_httpx_client._DEFAULT_HTTP_CLIENT_TIMEOUT,
+            follow_redirects=True,
+        )
         first_call_kwargs = mock_async_openai.call_args_list[0].kwargs
         second_call_kwargs = mock_async_openai.call_args_list[1].kwargs
         assert first_call_kwargs["http_client"] is shared_http_client
         assert second_call_kwargs["http_client"] is shared_http_client
 
+    def test_create_shared_http_client_rebuilds_closed_client(self):
+        """Closed cached clients should be replaced on the next factory call."""
+        from trpc_agent_sdk.models import _httpx_client
+
+        closed_client = Mock()
+        closed_client.is_closed = True
+        fresh_client = Mock()
+        fresh_client.is_closed = False
+
+        try:
+            _httpx_client._shared_http_clients.clear()
+            client_key = (1234, 1)
+            _httpx_client._shared_http_clients[client_key] = closed_client
+            with patch("trpc_agent_sdk.models._httpx_client._get_client_key", return_value=client_key):
+                with patch("trpc_agent_sdk.models._httpx_client.httpx.AsyncClient",
+                           return_value=fresh_client) as mock_httpx_client:
+                    assert _httpx_client._create_shared_http_client() is fresh_client
+        finally:
+            _httpx_client._shared_http_clients.clear()
+
+        mock_httpx_client.assert_called_once()
+
+    def test_create_shared_http_client_does_not_reuse_across_loop_keys(self):
+        """Different event loops should get different default httpx clients."""
+        from trpc_agent_sdk.models import _httpx_client
+
+        first_client = Mock()
+        first_client.is_closed = False
+        second_client = Mock()
+        second_client.is_closed = False
+
+        try:
+            _httpx_client._shared_http_clients.clear()
+            with patch("trpc_agent_sdk.models._httpx_client.httpx.AsyncClient",
+                       side_effect=[first_client, second_client]) as mock_httpx_client:
+                with patch("trpc_agent_sdk.models._httpx_client._get_client_key", side_effect=[(1234, 1), (1234, 2)]):
+                    assert _httpx_client._create_shared_http_client() is first_client
+                    assert _httpx_client._create_shared_http_client() is second_client
+        finally:
+            _httpx_client._shared_http_clients.clear()
+
+        assert mock_httpx_client.call_count == 2
+
+    def test_create_shared_http_client_does_not_reuse_across_process_keys(self):
+        """Different process keys should get different default httpx clients."""
+        from trpc_agent_sdk.models import _httpx_client
+
+        parent_client = Mock()
+        parent_client.is_closed = False
+        child_client = Mock()
+        child_client.is_closed = False
+
+        try:
+            _httpx_client._shared_http_clients.clear()
+            with patch("trpc_agent_sdk.models._httpx_client.httpx.AsyncClient",
+                       side_effect=[parent_client, child_client]) as mock_httpx_client:
+                with patch("trpc_agent_sdk.models._httpx_client._get_client_key",
+                           side_effect=[(1234, 1), (5678, 1)]):
+                    assert _httpx_client._create_shared_http_client() is parent_client
+                    assert _httpx_client._create_shared_http_client() is child_client
+        finally:
+            _httpx_client._shared_http_clients.clear()
+
+        assert mock_httpx_client.call_count == 2
+
+    def test_reset_shared_http_clients_after_fork_clears_cache_and_rebuilds_lock(self):
+        """Fork child reset should drop inherited clients and replace inherited locks."""
+        from trpc_agent_sdk.models import _httpx_client
+
+        inherited_client = Mock()
+        old_lock = _httpx_client._shared_http_clients_lock
+
+        try:
+            _httpx_client._shared_http_clients[(1234, 1)] = inherited_client
+
+            _httpx_client._reset_shared_http_clients_after_fork()
+
+            assert _httpx_client._shared_http_clients == {}
+            assert _httpx_client._shared_http_clients_lock is not old_lock
+        finally:
+            _httpx_client._shared_http_clients.clear()
+
     def test_create_async_client_overwrites_stale_client_args_http_client(self):
-        """Factory owns http_client injection even if client_args already has one."""
+        """Provider owns http_client injection even if client_args already has one."""
         stale_http_client = Mock()
         fresh_http_client = Mock()
-        http_client_factory = Mock(return_value=fresh_http_client)
+        http_client_provider = Mock()
+        http_client_provider.create_http_client.return_value = fresh_http_client
+        http_client_provider_factory = Mock(return_value=http_client_provider)
         model = OpenAIModel(
             model_name="gpt-4",
             api_key="test_key",
             client_args={"http_client": stale_http_client, "timeout": 30},
-            http_client_factory=http_client_factory,
+            http_client_provider_factory=http_client_provider_factory,
         )
 
         with patch("trpc_agent_sdk.models._openai_model.openai.AsyncOpenAI") as mock_async_openai:
@@ -354,15 +445,19 @@ async def test_generate_async_simple_text_response(self):
 
     @pytest.mark.asyncio
     async def test_generate_async_validation_failure(self):
-        """Test generate_async converts validation failures to error responses."""
+        """Test generate_async returns an error response on invalid request."""
         model = OpenAIModel(model_name="gpt-4", api_key="test_key")
 
+        # Empty contents
         request = LlmRequest(contents=[], config=None, tools_dict={})
 
-        responses = [response async for response in model.generate_async(request, stream=False)]
+        responses = []
+        async for response in model.generate_async(request, stream=False):
+            responses.append(response)
+
         assert len(responses) == 1
         assert responses[0].error_code == "API_ERROR"
-        assert "At least one content is required" in (responses[0].error_message or "")
+        assert "At least one content is required" in responses[0].error_message
 
     @pytest.mark.asyncio
     async def test_generate_async_with_config_parameters(self):
@@ -849,47 +944,3 @@ def test_null_prompt_tokens_details_does_not_crash(self):
         }
         meta = OpenAIModel._build_usage_metadata(usage_data)
         assert meta.cache_read_input_tokens is None
-
-
-class _RetryTestError(Exception):
-
-    def __init__(self, status_code=None, headers=None):
-        super().__init__(f"status {status_code}" if status_code is not None else "retry test")
-        if status_code is not None:
-            self.status_code = status_code
-        if headers is not None:
-            self.response = type("Resp", (), {"headers": headers})()
-
-
-class TestOpenAIModelRetryHooks:
-
-    def _model(self):
-        return OpenAIModel(model_name="gpt-4", api_key="test_key")
-
-    def test_x_should_retry_header_has_priority(self):
-        model = self._model()
-        assert model._get_model_retry_info(_RetryTestError(400, {"x-should-retry": "true"})).should_retry is True
-        assert model._get_model_retry_info(_RetryTestError(500, {"x-should-retry": "false"})).should_retry is False
-
-    @pytest.mark.parametrize("status_code", [408, 409, 429, 500, 503])
-    def test_retryable_status_codes(self, status_code):
-        assert self._model()._get_model_retry_info(_RetryTestError(status_code)).should_retry is True
-
-    @pytest.mark.parametrize("status_code", [400, 401, 403, 404, 499])
-    def test_non_retryable_status_codes(self, status_code):
-        assert self._model()._get_model_retry_info(_RetryTestError(status_code)).should_retry is False
-
-    def test_timeout_exception_retried(self):
-        request = httpx.Request("GET", "https://example.com")
-        assert self._model()._get_model_retry_info(httpx.TimeoutException("timeout", request=request)).should_retry is True
-
-    def test_openai_error_not_retried_without_status_decision(self):
-        assert self._model()._get_model_retry_info(openai.OpenAIError("boom")).should_retry is False
-
-    def test_other_exception_retried(self):
-        assert self._model()._get_model_retry_info(ValueError("boom")).should_retry is True
-
-    def test_retry_after_extracted_from_headers(self):
-        info = self._model()._get_model_retry_info(_RetryTestError(429, {"retry-after": "3"}))
-        assert info.should_retry is True
-        assert info.retry_after == 3.0

trpc_agent_sdk/agents/_base_agent.py

@@ -257,8 +257,8 @@ async def run_async(
                 - Actions
         """
         from trpc_agent_sdk.telemetry import report_invoke_agent
-        from trpc_agent_sdk.telemetry._trace import tracer
-        from trpc_agent_sdk.telemetry._trace import trace_agent
+        from trpc_agent_sdk.telemetry import tracer
+        from trpc_agent_sdk.telemetry import trace_agent
 
         # Manually propagate span context using attach/detach instead of
         # start_as_current_span. This ensures child spans (call_llm, execute_tool,

trpc_agent_sdk/models/__init__.py

@@ -45,6 +45,10 @@
 from ._openai_model import OpenAIModel
 from ._openai_model import ToolCall
 from ._openai_model import ToolKey
+from ._httpx_client import close_shared_http_clients
+from ._httpx_client import temporary_http_client_provider_factory
+from ._httpx_client import shared_http_client_provider_factory
+from ._httpx_client import HttpClientProviderFactory
 from ._registry import ModelRegistry
 from ._registry import register_model
 
@@ -85,6 +89,10 @@
     "OpenAIModel",
     "ToolCall",
     "ToolKey",
+    "close_shared_http_clients",
+    "temporary_http_client_provider_factory",
+    "shared_http_client_provider_factory",
+    "HttpClientProviderFactory",
     "ModelRegistry",
     "register_model",
 ]

trpc_agent_sdk/models/_anthropic_model.py

@@ -38,6 +38,9 @@
 from ._llm_model import LLMModel
 from ._llm_request import LlmRequest
 from ._llm_response import LlmResponse
+from ._httpx_client import BaseHttpClientProvider
+from ._httpx_client import HttpClientProviderFactory
+from ._httpx_client import shared_http_client_provider_factory
 from ._registry import register_model
 
 _EPHEMERAL = "ephemeral"
@@ -154,29 +157,36 @@ def __init__(
         model_name: str,
         filters_name: Optional[list[str]] = None,
         generate_content_config: Optional[GenerateContentConfig] = None,
+        http_client_provider_factory: HttpClientProviderFactory = shared_http_client_provider_factory,
         **kwargs,
     ):
         super().__init__(model_name, filters_name, **kwargs)
 
         # Extract Anthropic-specific config
         self.client_args = kwargs.get(const.CLIENT_ARGS, {})
-
+        # Allow callers to inject a tuned httpx client so the underlying openai.AsyncOpenAI honors connection-pool
+        # settings such as keepalive_expiry / max_keepalive_connections (avoids reusing stale
+        # keep-alive sockets that gateways close earlier than httpx).
+        http_client_provider_factory = http_client_provider_factory or shared_http_client_provider_factory
+        self._http_client_provider: BaseHttpClientProvider = http_client_provider_factory()
         # Default generation config that can be overridden per request
         self.generate_content_config = generate_content_config
 
-    def _create_async_client(self):
+    def _create_async_client(self) -> AsyncAnthropic:
         """Create a new async client instance."""
 
         # Disable httpx logging to prevent HTTP request logs
         import logging
 
         logging.getLogger("httpx").setLevel(logging.WARNING)
+        client_args = self.client_args.copy()
+        client_args['http_client'] = self._http_client_provider.create_http_client()
 
         return AsyncAnthropic(
             api_key=self._api_key,
             max_retries=0,  # disable retries
             base_url=self._base_url if self._base_url else None,
-            **self.client_args,
+            **client_args,
         )
 
     def is_retriable_status_code(self, status_code: int) -> Optional[bool]:
@@ -555,7 +565,7 @@ async def _generate_single(
             response = await client.messages.create(**api_params)
             return self._message_to_llm_response(response)
         finally:
-            await client.close()
+            await self._http_client_provider.close_http_client(client)
 
     async def _generate_stream(
         self,
@@ -683,7 +693,7 @@ async def _generate_stream(
             logger.error("Error in streaming response", exc_info=True)
             raise
         finally:
-            await client.close()
+            await self._http_client_provider.close_http_client(client)
 
     def _apply_prompt_cache(self, api_params: Dict[str, Any], ctx: InvocationContext | None) -> None:
         """Inject Anthropic native cache_control breakpoints (opt-in, no-op when disabled)."""