Bugfix: 修复llm调用抛出GeneratorExit时无trace上报的问题

tests/agents/core/test_llm_processor.py

@@ -242,3 +242,35 @@ async def run():
 
         assert event.error_code == "STREAMING_ERROR"
         assert mock_trace.call_args.args[3].error_message == "rate limit exceeded"
+
+    def test_partial_stream_close_traces_consumer_error(self, invocation_context):
+        m = MockLLMModel(model_name="test-llmproc-model")
+        m._responses = [
+            LlmResponse(content=Content(parts=[Part(text="part1")]), partial=True),
+            LlmResponse(content=Content(parts=[Part(text="part2")]), partial=True),
+        ]
+        proc = LlmProcessor(m)
+        request = LlmRequest()
+
+        async def run():
+            stream = proc.call_llm_async(request, invocation_context, stream=True)
+            event = await anext(stream)
+            await stream.aclose()
+            return event
+
+        with patch("trpc_agent_sdk.agents.core._llm_processor.report_call_llm") as mock_report, \
+             patch("trpc_agent_sdk.agents.core._llm_processor.trace_call_llm") as mock_trace, \
+             patch("trpc_agent_sdk.agents.core._llm_processor.tracer"):
+            event = asyncio.run(run())
+
+        assert event.partial is True
+        mock_trace.assert_called_once()
+        assert mock_trace.call_args.args[2] is request
+        response = mock_trace.call_args.args[3]
+        assert response.error_code == "LlmCallGeneratorExit"
+        assert response.error_message == "LLM call stopped with GeneratorExit."
+        assert response.interrupted is True
+        assert response.content is None
+        assert response.custom_metadata is None
+        assert mock_report.call_args.args[2] is response
+        assert mock_report.call_args.kwargs["error_type"] == "LlmCallGeneratorExit"

tests/agents/core/test_tools_processor.py

@@ -33,11 +33,13 @@ def _compat_get_skill_processor_parameters(agent_context):
 
 
 class _StubAgent(BaseAgent):
+
     async def _run_async_impl(self, ctx):
         yield
 
 
 class MockLLMModel(LLMModel):
+
     @classmethod
     def supported_models(cls) -> List[str]:
         return [r"test-tools-proc-.*"]
@@ -65,9 +67,7 @@ def sample_tool(name: str, value: str) -> dict:
 @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 +86,7 @@ def invocation_context():
 
 
 class TestToolsProcessorInit:
+
     def test_stores_tools(self):
         tool = FunctionTool(sample_tool)
         proc = ToolsProcessor([tool])
@@ -102,6 +103,7 @@ def test_empty_tools(self):
 
 
 class TestFindTool:
+
     def test_finds_matching_tool(self):
         tool = FunctionTool(sample_tool)
         proc = ToolsProcessor([tool])
@@ -131,6 +133,7 @@ async def run():
 
 
 class TestFindToolPublic:
+
     def test_resolves_and_finds(self, invocation_context):
         tool = FunctionTool(sample_tool)
         proc = ToolsProcessor([tool])
@@ -150,6 +153,7 @@ async def run():
 
 
 class TestExecuteToolsSequential:
+
     def test_single_tool_call(self, invocation_context):
         tool = FunctionTool(sample_tool)
         proc = ToolsProcessor([tool])
@@ -198,6 +202,7 @@ async def run():
 
 
 class TestMergeParallelFunctionResponseEvents:
+
     def test_single_event_returns_as_is(self):
         proc = ToolsProcessor([])
         event = Event(
@@ -253,11 +258,10 @@ def test_merged_actions(self):
 
 
 class TestToolsProcessorErrorEvent:
+
     def test_creates_error_with_function_response(self, invocation_context):
         proc = ToolsProcessor([])
-        event = proc._create_error_event(
-            invocation_context, "test_error", "Something failed", "call-1", "my_tool"
-        )
+        event = proc._create_error_event(invocation_context, "test_error", "Something failed", "call-1", "my_tool")
         assert event.error_code == "test_error"
         assert event.error_message == "Something failed"
         assert event.content is not None
@@ -273,7 +277,6 @@ def test_error_event_without_tool_info(self, invocation_context):
 # _update_streaming_tool_names
 # ---------------------------------------------------------------------------
 
-
 # ---------------------------------------------------------------------------
 # execute_tools_async - progress-streaming tool path
 # ---------------------------------------------------------------------------
@@ -335,6 +338,7 @@ async def run():
         assert fr.response == {"status": "done", "url": "https://x", "steps": 2}
 
     def test_streaming_tool_error_yields_error_event(self, invocation_context):
+
         async def boom(query: str):
             yield {"status": "started"}
             raise RuntimeError("kaboom")
@@ -392,8 +396,7 @@ async def run():
 
             # The streaming call yields partials AND its own final event.
             stream_partials = [
-                ev for ev in events
-                if ev.partial and (ev.custom_metadata or {}).get("tool_call_id") == "c-stream"
+                ev for ev in events if ev.partial and (ev.custom_metadata or {}).get("tool_call_id") == "c-stream"
             ]
             stream_finals = [
                 ev for ev in events if ev.partial is not True and ev.content and any(
@@ -426,6 +429,7 @@ async def run():
 
 
 class TestUpdateStreamingToolNames:
+
     def test_no_streaming_tools(self):
         proc = ToolsProcessor([])
         request = LlmRequest()

tests/agents/test_base_agent.py

@@ -34,6 +34,7 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event,
 
 
 class MockLLMModel(LLMModel):
+
     @classmethod
     def supported_models(cls) -> List[str]:
         return [r"test-base-.*"]
@@ -56,9 +57,7 @@ def register_test_model():
 @pytest.fixture
 def invocation_context():
     service = InMemorySessionService()
-    session = asyncio.run(
-        service.create_session(app_name="test_app", user_id="user-1", session_id="s-1")
-    )
+    session = asyncio.run(service.create_session(app_name="test_app", user_id="user-1", session_id="s-1"))
     agent = ConcreteAgent(name="test_agent")
     return InvocationContext(
         session_service=service,
@@ -75,6 +74,7 @@ def invocation_context():
 
 
 class TestBuildActionStringFromEvents:
+
     def test_empty_events(self):
         assert _build_action_string_from_events([]) == ""
 
@@ -166,6 +166,7 @@ def test_multiple_events_joined_by_double_newline(self):
 
 
 class TestCreateInvocationContext:
+
     def test_same_agent_keeps_branch(self, invocation_context):
         agent = invocation_context.agent
         invocation_context.branch = "existing_branch"
@@ -195,6 +196,7 @@ def test_no_branch_initializes_with_name(self, invocation_context):
 
 
 class TestBaseAgentModelPostInit:
+
     def test_invalid_filter_name_raises(self):
         with pytest.raises(ValueError, match="not found"):
             ConcreteAgent(name="bad_agent", filters_name=["nonexistent_filter"])
@@ -208,6 +210,7 @@ def test_callback_filter_appended(self):
 
 
 class TestBaseAgentGetSubagents:
+
     def test_returns_sub_agents_list(self):
         child = ConcreteAgent(name="child")
         parent = ConcreteAgent(name="parent", sub_agents=[child])
@@ -216,3 +219,28 @@ def test_returns_sub_agents_list(self):
     def test_empty_sub_agents(self):
         agent = ConcreteAgent(name="solo")
         assert agent.get_subagents() == []
+
+
+class TestBaseAgentTracing:
+
+    def test_closing_stream_marks_agent_span_interrupted(self, invocation_context):
+        agent = invocation_context.agent
+
+        async def run():
+            stream = agent.run_async(invocation_context)
+            await anext(stream)
+            await stream.aclose()
+
+        with patch("trpc_agent_sdk.agents._base_agent.mark_span_error") as mock_span_error, \
+             patch("trpc_agent_sdk.agents._base_agent.report_invoke_agent") as mock_report, \
+             patch("trpc_agent_sdk.agents._base_agent.trace_agent"), \
+             patch("trpc_agent_sdk.agents._base_agent.tracer") as mock_tracer:
+            asyncio.run(run())
+
+        agent_span = mock_tracer.start_as_current_span.return_value.__enter__.return_value
+        mock_span_error.assert_called_once_with(
+            agent_span,
+            error_type="AgentGeneratorExit",
+            description="Agent execution stopped with GeneratorExit.",
+        )
+        assert mock_report.call_args.kwargs["error_type"] == "AgentGeneratorExit"

tests/telemetry/test_trace.py

@@ -29,6 +29,7 @@
     _build_llm_request_for_trace,
     _safe_json_serialize,
     get_trpc_agent_span_name,
+    mark_span_error,
     set_trpc_agent_span_name,
     trace_agent,
     trace_call_llm,
@@ -179,6 +180,29 @@ def test_serialize_empty_dict(self):
         assert json.loads(_safe_json_serialize({})) == {}
 
 
+# ---------------------------------------------------------------------------
+# Tests: mark_span_error
+# ---------------------------------------------------------------------------
+
+
+class TestMarkSpanError:
+
+    def test_marks_interruption_with_operation_specific_error(self):
+        span = _mock_span()
+
+        mark_span_error(
+            span,
+            error_type="RunnerGeneratorExit",
+            description="Runner invocation stopped with GeneratorExit.",
+        )
+
+        span.set_status.assert_called_once_with(
+            trace.StatusCode.ERROR,
+            "Runner invocation stopped with GeneratorExit.",
+        )
+        span.set_attribute.assert_called_once_with("error.type", "RunnerGeneratorExit")
+
+
 # ---------------------------------------------------------------------------
 # Tests: trace_runner
 # ---------------------------------------------------------------------------

tests/test_runner.py

@@ -273,6 +273,47 @@ async def mock_agent_run(ctx):
         assert events[0].partial is True
         assert events[1].partial is False
 
+    @pytest.mark.asyncio
+    async def test_closing_stream_marks_invocation_span_interrupted(self, runner, mock_session_service, mock_agent,
+                                                                    mock_session):
+        mock_session_service.get_session.return_value = mock_session
+
+        async def mock_agent_run(ctx):
+            yield Event(
+                invocation_id=ctx.invocation_id,
+                author="test_agent",
+                content=Content(parts=[Part(text="Partial")]),
+                partial=True,
+            )
+            yield Event(
+                invocation_id=ctx.invocation_id,
+                author="test_agent",
+                content=Content(parts=[Part(text="Complete")]),
+                partial=False,
+            )
+
+        mock_agent.run_async = mock_agent_run
+
+        with patch("trpc_agent_sdk.runners.mark_span_error") as mock_span_error, \
+             patch("trpc_agent_sdk.runners.trace_runner"), \
+             patch("trpc_agent_sdk.runners.tracer") as mock_tracer:
+            stream = runner.run_async(
+                user_id="test_user",
+                session_id="test_session",
+                new_message=Content(parts=[Part(text="Hello")]),
+                run_config=RunConfig(streaming=True),
+            )
+            event = await anext(stream)
+            await stream.aclose()
+
+        assert event.partial is True
+        invocation_span = mock_tracer.start_as_current_span.return_value.__enter__.return_value
+        mock_span_error.assert_called_once_with(
+            invocation_span,
+            error_type="RunnerGeneratorExit",
+            description="Runner invocation stopped with GeneratorExit.",
+        )
+
     @pytest.mark.asyncio
     async def test_run_async_non_streaming_mode(self, runner, mock_session_service, mock_agent, mock_session):
         """Test non-streaming mode only yields complete events."""

trpc_agent_sdk/agents/_base_agent.py

@@ -33,6 +33,8 @@
 from typing import final
 from typing_extensions import override
 
+from opentelemetry import trace
+
 from trpc_agent_sdk.abc import AgentABC
 from trpc_agent_sdk.abc import FilterType
 from trpc_agent_sdk.code_executors import BaseCodeExecutor
@@ -43,6 +45,10 @@
 from trpc_agent_sdk.events import Event
 from trpc_agent_sdk.filter import get_filter
 from trpc_agent_sdk.filter import run_stream_filters
+from trpc_agent_sdk.telemetry import mark_span_error
+from trpc_agent_sdk.telemetry import report_invoke_agent
+from trpc_agent_sdk.telemetry import tracer
+from trpc_agent_sdk.telemetry import trace_agent
 
 from ._callback import AgentCallback
 from ._callback import AgentCallbackFilter
@@ -256,18 +262,14 @@ async def run_async(
                 - State changes
                 - Actions
         """
-        from trpc_agent_sdk.telemetry import report_invoke_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,
         # etc.) can correctly resolve their parent.
         # We use start_span + attach/detach rather than start_as_current_span
         # because __aexit__ of the context manager is not guaranteed to run when
         # an async generator is cancelled, but try/finally always executes
         # even under CancelledError (PEP 492).
-        with tracer.start_as_current_span(f"agent_run [{self.name}]"):
+        with tracer.start_as_current_span(f"agent_run [{self.name}]") as agent_span:
             ctx = self._create_invocation_context(parent_context)
             if ctx.agent_context is None:
                 ctx.agent_context = create_agent_context()
@@ -294,6 +296,14 @@ async def run_async(
                         # This excludes state update events which have content=None
                         non_partial_events.append(event)
                     yield event  # type: ignore
+            except GeneratorExit:
+                metrics_error_type = "AgentGeneratorExit"
+                mark_span_error(
+                    agent_span,
+                    error_type=metrics_error_type,
+                    description="Agent execution stopped with GeneratorExit.",
+                )
+                raise
             except Exception as ex:
                 metrics_error_type = type(ex).__name__
                 raise
@@ -305,12 +315,13 @@ async def run_async(
                 agent_action = _build_action_string_from_events(non_partial_events)
 
                 # Call trace function with agent execution details
-                trace_agent(
-                    invocation_context=ctx,
-                    agent_action=agent_action,
-                    state_begin=state_begin,
-                    state_end=state_end,
-                )
+                with trace.use_span(agent_span, end_on_exit=False):
+                    trace_agent(
+                        invocation_context=ctx,
+                        agent_action=agent_action,
+                        state_begin=state_begin,
+                        state_end=state_end,
+                    )
 
                 duration_s = time.monotonic() - mono_start
                 ttft_s = (t_first_visible - mono_start) if t_first_visible is not None else duration_s