Bugfix: 修复模型调用失败/重试无trace的问题

tests/agents/core/test_llm_processor.py

@@ -10,6 +10,7 @@
 import asyncio
 from typing import List
 from unittest.mock import Mock
+from unittest.mock import patch
 
 import pytest
 
@@ -23,6 +24,7 @@
 
 
 class _StubAgent(BaseAgent):
+
     async def _run_async_impl(self, ctx):
         yield
 
@@ -53,21 +55,17 @@ def register_test_model():
 @pytest.fixture
 def model():
     m = MockLLMModel(model_name="test-llmproc-model")
-    m._responses = [
-        LlmResponse(
-            content=Content(parts=[Part(text="hello")]),
-            partial=False,
-        )
-    ]
+    m._responses = [LlmResponse(
+        content=Content(parts=[Part(text="hello")]),
+        partial=False,
+    )]
     return m
 
 
 @pytest.fixture
 def invocation_context():
     service = InMemorySessionService()
-    session = asyncio.run(
-        service.create_session(app_name="test", user_id="u1", session_id="s1")
-    )
+    session = asyncio.run(service.create_session(app_name="test", user_id="u1", session_id="s1"))
     agent = _StubAgent(name="test_agent")
     ctx = InvocationContext(
         session_service=service,
@@ -86,6 +84,7 @@ def invocation_context():
 
 
 class TestCreateEventFromResponse:
+
     def test_maps_response_fields(self, model, invocation_context):
         proc = LlmProcessor(model)
         response = LlmResponse(
@@ -121,6 +120,7 @@ def test_preserves_error_fields(self, model, invocation_context):
 
 
 class TestCreateErrorEvent:
+
     def test_creates_error_event(self, model, invocation_context):
         proc = LlmProcessor(model)
         event = proc._create_error_event(invocation_context, "err_code", "err_msg")
@@ -136,6 +136,7 @@ def test_creates_error_event(self, model, invocation_context):
 
 
 class TestProcessPlanningResponse:
+
     def test_no_planner_returns_event_unchanged(self, model, invocation_context):
         proc = LlmProcessor(model)
         event = Event(
@@ -159,6 +160,7 @@ def test_event_without_content_skips_planning(self, model, invocation_context):
 
 
 class TestCallLlmAsync:
+
     def test_yields_events_for_responses(self, model, invocation_context):
         proc = LlmProcessor(model)
         request = LlmRequest()
@@ -210,3 +212,33 @@ async def run():
         assert len(content_events) == 2
         assert content_events[0].partial is True
         assert content_events[1].partial is False
+
+    def test_error_response_is_traced_before_consumer_stops(self, invocation_context):
+        m = MockLLMModel(model_name="test-llmproc-model")
+        m._responses = [
+            LlmResponse(
+                error_code="STREAMING_ERROR",
+                error_message="rate limit exceeded",
+                partial=False,
+            )
+        ]
+        proc = LlmProcessor(m)
+        request = LlmRequest()
+
+        async def run():
+            stream = proc.call_llm_async(request, invocation_context, stream=True)
+            event = await anext(stream)
+            # The downstream LlmAgent returns immediately for an error event,
+            # so tracing and span-context cleanup must be complete at this point.
+            mock_trace.assert_called_once()
+            span_context.__exit__.assert_called_once()
+            await stream.aclose()
+            return event
+
+        with patch("trpc_agent_sdk.agents.core._llm_processor.trace_call_llm") as mock_trace, \
+             patch("trpc_agent_sdk.agents.core._llm_processor.tracer") as mock_tracer:
+            span_context = mock_tracer.start_as_current_span.return_value
+            event = asyncio.run(run())
+
+        assert event.error_code == "STREAMING_ERROR"
+        assert mock_trace.call_args.args[3].error_message == "rate limit exceeded"

tests/models/test_retry.py

@@ -10,8 +10,11 @@
 from collections.abc import AsyncGenerator
 from typing import Optional
 from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
 from unittest.mock import patch
 
+from opentelemetry import trace
+
 from trpc_agent_sdk.configs import ExponentialBackoffConfig
 from trpc_agent_sdk.configs import ModelRetryConfig
 from trpc_agent_sdk.models._llm_response import LlmResponse
@@ -33,7 +36,6 @@ def __init__(self, status_code: int | str, headers: Optional[dict] = None):
             self.response = type("Resp", (), {"headers": headers})()
 
 
-
 class _HeadersError(Exception):
 
     def __init__(self, headers: dict):
@@ -73,13 +75,11 @@ async def _collect(
     *,
     get_retry_info=None,
 ) -> list[LlmResponse]:
-    return [
-        response async for response in retry_model_call(
-            call_model,
-            config,
-            get_retry_info=get_retry_info,
-        )
-    ]
+    return [response async for response in retry_model_call(
+        call_model,
+        config,
+        get_retry_info=get_retry_info,
+    )]
 
 
 class TestRetryHelpers:
@@ -192,12 +192,46 @@ async def call_model() -> AsyncGenerator[LlmResponse, None]:
             yield _content_response("ok")
 
         with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()) as sleep:
-            responses = await _collect(call_model, self._retry_cfg(), get_retry_info=lambda _: ModelRetryInfo(should_retry=True))
+            responses = await _collect(call_model,
+                                       self._retry_cfg(),
+                                       get_retry_info=lambda _: ModelRetryInfo(should_retry=True))
         assert attempts == 2
         assert sleep.await_count == 1
         assert responses[-1].content.parts[0].text == "ok"
         assert all(response.error_code is None for response in responses)
 
+    async def test_retry_records_failed_attempt_span(self):
+        attempts = 0
+        error = _StatusError(429)
+
+        async def call_model() -> AsyncGenerator[LlmResponse, None]:
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise error
+            yield _content_response("ok")
+
+        span = MagicMock()
+        with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()), \
+             patch("trpc_agent_sdk.models._retry._retry_tracer") as retry_tracer:
+            retry_tracer.start_as_current_span.return_value.__enter__.return_value = span
+            responses = await _collect(
+                call_model,
+                self._retry_cfg(),
+                get_retry_info=lambda _: ModelRetryInfo(should_retry=True),
+            )
+
+        assert responses[-1].content.parts[0].text == "ok"
+        retry_tracer.start_as_current_span.assert_called_once_with("model_retry")
+        span.record_exception.assert_called_once_with(error)
+        span.set_status.assert_called_once_with(trace.StatusCode.ERROR, "status 429")
+        span.set_attribute.assert_any_call("gen_ai.operation.name", "model_retry")
+        span.set_attribute.assert_any_call("error.type", "_StatusError")
+        span.set_attribute.assert_any_call("gen_ai.retry.number", 1)
+        span.set_attribute.assert_any_call("gen_ai.retry.max_retries", 2)
+        span.set_attribute.assert_any_call("gen_ai.retry.delay_seconds", 0.0)
+        span.set_attribute.assert_any_call("http.response.status_code", 429)
+
     async def test_exhausts_budget_then_yields_error(self):
         attempts = 0
 
@@ -208,7 +242,9 @@ async def call_model() -> AsyncGenerator[LlmResponse, None]:
             yield
 
         with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()) as sleep:
-            responses = await _collect(call_model, self._retry_cfg(num_retries=2), get_retry_info=lambda _: ModelRetryInfo(should_retry=True))
+            responses = await _collect(call_model,
+                                       self._retry_cfg(num_retries=2),
+                                       get_retry_info=lambda _: ModelRetryInfo(should_retry=True))
         assert attempts == 3
         assert sleep.await_count == 2
         assert responses[-1].error_code == "API_ERROR"
@@ -224,7 +260,9 @@ async def call_model() -> AsyncGenerator[LlmResponse, None]:
             yield
 
         with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()) as sleep:
-            responses = await _collect(call_model, self._retry_cfg(), get_retry_info=lambda _: ModelRetryInfo(should_retry=False))
+            responses = await _collect(call_model,
+                                       self._retry_cfg(),
+                                       get_retry_info=lambda _: ModelRetryInfo(should_retry=False))
         assert attempts == 1
         assert sleep.await_count == 0
         assert responses[-1].error_code == "API_ERROR"
@@ -258,7 +296,9 @@ async def call_model() -> AsyncGenerator[LlmResponse, None]:
             raise _StatusError(429)
 
         with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()) as sleep:
-            responses = await _collect(call_model, self._retry_cfg(), get_retry_info=lambda _: ModelRetryInfo(should_retry=True))
+            responses = await _collect(call_model,
+                                       self._retry_cfg(),
+                                       get_retry_info=lambda _: ModelRetryInfo(should_retry=True))
         assert attempts == 1
         assert sleep.await_count == 0
         assert responses[0].content.parts[0].text == "partial"
@@ -279,6 +319,8 @@ async def second_attempt() -> AsyncGenerator[LlmResponse, None]:
 
         attempts = iter([first_attempt, second_attempt])
         with patch("trpc_agent_sdk.models._retry.asyncio.sleep", new=AsyncMock()):
-            responses = await _collect(lambda: next(attempts)(), self._retry_cfg(), get_retry_info=lambda _: ModelRetryInfo(should_retry=True))
+            responses = await _collect(lambda: next(attempts)(),
+                                       self._retry_cfg(),
+                                       get_retry_info=lambda _: ModelRetryInfo(should_retry=True))
         assert closed_attempts == ["first"]
         assert responses[-1].content.parts[0].text == "ok"