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

tests/agents/core/test_llm_processor.py

@@ -242,3 +242,126 @@ 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_accumulated_text_and_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)
+            events = [await anext(stream), await anext(stream)]
+            await stream.aclose()
+            return events
+
+        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"):
+            events = asyncio.run(run())
+
+        assert all(event.partial is True for event in events)
+        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.partial is True
+        assert response.content.role == "model"
+        assert response.content.parts[0].text == "part1part2"
+        assert response.custom_metadata is None
+        assert mock_report.call_args.args[2] is response
+        assert mock_report.call_args.kwargs["error_type"] == "LlmCallGeneratorExit"
+
+    def test_partial_stream_close_keeps_latest_function_call_content(self, invocation_context):
+        m = MockLLMModel(model_name="test-llmproc-model")
+        function_call = Part.from_function_call(name="get_weather_report", args={"city": "Beijing"})
+        m._responses = [
+            LlmResponse(content=Content(parts=[function_call]), 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"), \
+             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.get_function_calls()
+        response = mock_trace.call_args.args[3]
+        assert response.error_code == "LlmCallGeneratorExit"
+        assert response.content.parts[0].function_call.name == "get_weather_report"
+        assert response.content.parts[0].function_call.args == {"city": "Beijing"}
+
+    def test_partial_stream_close_joins_thought_and_visible_text(self, invocation_context):
+        m = MockLLMModel(model_name="test-llmproc-model")
+        thought1 = Part(text="I should call get_")
+        thought1.thought = True
+        thought2 = Part(text="weather_report function with")
+        thought2.thought = True
+        m._responses = [
+            LlmResponse(content=Content(parts=[thought1]), partial=True),
+            LlmResponse(content=Content(parts=[thought2]), partial=True),
+            LlmResponse(content=Content(parts=[Part(text="Let me check.")]), partial=True),
+        ]
+        proc = LlmProcessor(m)
+        request = LlmRequest()
+
+        async def run():
+            stream = proc.call_llm_async(request, invocation_context, stream=True)
+            events = [await anext(stream), await anext(stream), await anext(stream)]
+            await stream.aclose()
+            return events
+
+        with patch("trpc_agent_sdk.agents.core._llm_processor.report_call_llm"), \
+             patch("trpc_agent_sdk.agents.core._llm_processor.trace_call_llm") as mock_trace, \
+             patch("trpc_agent_sdk.agents.core._llm_processor.tracer"):
+            asyncio.run(run())
+
+        response = mock_trace.call_args.args[3]
+        assert response.error_code == "LlmCallGeneratorExit"
+        assert response.content.parts[0].text == "I should call get_weather_report function with"
+        assert response.content.parts[0].thought is True
+        assert response.content.parts[1].text == "Let me check."
+        assert not response.content.parts[1].thought
+
+    def test_stream_exception_traces_accumulated_partial_content(self, invocation_context):
+        m = MockLLMModel(model_name="test-llmproc-model")
+        proc = LlmProcessor(m)
+        request = LlmRequest()
+
+        async def failing_generate(request, stream=False, ctx=None):
+            yield LlmResponse(content=Content(parts=[Part(text="part1")]), partial=True)
+            yield LlmResponse(content=Content(parts=[Part(text="part2")]), partial=True)
+            raise RuntimeError("upstream failed")
+
+        async def run():
+            stream = proc.call_llm_async(request, invocation_context, stream=True)
+            return [await anext(stream), await anext(stream), await anext(stream)]
+
+        with patch.object(m, "generate_async", failing_generate), \
+             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"):
+            events = asyncio.run(run())
+
+        assert events[0].content.parts[0].text == "part1"
+        assert events[1].content.parts[0].text == "part2"
+        assert events[2].is_error()
+        response = mock_trace.call_args.args[3]
+        assert response.error_code == "LLM_CALL_ERROR"
+        assert response.error_message == "upstream failed"
+        assert response.partial is True
+        assert response.content.parts[0].text == "part1part2"
+        assert response.custom_metadata == {"error_type": "RuntimeError"}
+        assert mock_report.call_args.kwargs["error_type"] == "RuntimeError"

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
 # ---------------------------------------------------------------------------
@@ -913,7 +937,7 @@ def test_basic_llm_trace(self, mock_get_span):
         span.set_attribute.assert_any_call("trpc.python.agent.event_id", "e-1")
 
     @patch("trpc_agent_sdk.telemetry._trace.trace.get_current_span")
-    def test_error_response_marks_span_and_records_exception(self, mock_get_span):
+    def test_error_response_sets_status_message_and_keeps_llm_response_output(self, mock_get_span):
         span = _mock_span()
         mock_get_span.return_value = span
         ctx = _make_invocation_context()
@@ -927,19 +951,19 @@ def test_error_response_marks_span_and_records_exception(self, mock_get_span):
 
         trace_call_llm(ctx, event_id="e-1", llm_request=req, llm_response=resp)
 
+        # Output remains the LlmResponse JSON; status carries the error message.
+        span.set_attribute.assert_any_call("trpc.python.agent.llm_response", '{"content": "response"}')
         span.set_status.assert_called_once_with(trace.StatusCode.ERROR, "rate limit exceeded")
         span.set_attribute.assert_any_call("error.type", "RateLimitError")
         span.set_attribute.assert_any_call(
             "trpc.python.agent.llm.error_code",
             "STREAMING_ERROR",
         )
-        span.add_event.assert_called_once_with(
-            "exception",
-            {
-                "exception.type": "RateLimitError",
-                "exception.message": "rate limit exceeded",
-            },
+        span.set_attribute.assert_any_call(
+            "trpc.python.agent.llm.error_message",
+            "rate limit exceeded",
         )
+        span.add_event.assert_not_called()
 
     @patch("trpc_agent_sdk.telemetry._trace.trace.get_current_span")
     def test_with_usage_metadata(self, mock_get_span):