From b9c13ad067a412b9ad3640b9e6fa3d3605858b62 Mon Sep 17 00:00:00 2001 From: weimch Date: Tue, 4 Aug 2026 12:07:37 +0800 Subject: [PATCH] =?UTF-8?q?Bugfix:=20=E4=BF=AE=E5=A4=8Dllm=E8=B0=83?= =?UTF-8?q?=E7=94=A8=E6=8A=9B=E5=87=BAGeneratorExit=E6=97=B6=E6=97=A0trace?= =?UTF-8?q?=E4=B8=8A=E6=8A=A5=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 原因：服务部署时，如果客户端直接断开连接，则运行中的Agent会收到协程被取消的异常，之前异常时没有上报，现在修复这个问题 --- tests/agents/core/test_llm_processor.py | 123 +++++++++++++++++++ tests/agents/core/test_tools_processor.py | 22 ++-- tests/agents/test_base_agent.py | 34 ++++- tests/telemetry/test_trace.py | 38 ++++-- tests/test_runner.py | 41 +++++++ trpc_agent_sdk/agents/_base_agent.py | 33 +++-- trpc_agent_sdk/agents/core/_llm_processor.py | 87 +++++++++---- trpc_agent_sdk/runners.py | 73 ++++++----- trpc_agent_sdk/telemetry/__init__.py | 2 + trpc_agent_sdk/telemetry/_trace.py | 28 +++-- 10 files changed, 388 insertions(+), 93 deletions(-) diff --git a/tests/agents/core/test_llm_processor.py b/tests/agents/core/test_llm_processor.py index 1b3762616..e69170163 100644 --- a/tests/agents/core/test_llm_processor.py +++ b/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" diff --git a/tests/agents/core/test_tools_processor.py b/tests/agents/core/test_tools_processor.py index 8dd9255f2..ce35acec6 100644 --- a/tests/agents/core/test_tools_processor.py +++ b/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() diff --git a/tests/agents/test_base_agent.py b/tests/agents/test_base_agent.py index cf976c01f..a716111b0 100644 --- a/tests/agents/test_base_agent.py +++ b/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" diff --git a/tests/telemetry/test_trace.py b/tests/telemetry/test_trace.py index 51bbd4dfe..6b8453b43 100644 --- a/tests/telemetry/test_trace.py +++ b/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): diff --git a/tests/test_runner.py b/tests/test_runner.py index a9a1a1a03..3756a962b 100644 --- a/tests/test_runner.py +++ b/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.""" diff --git a/trpc_agent_sdk/agents/_base_agent.py b/trpc_agent_sdk/agents/_base_agent.py index 90ecef95d..ff39076c6 100644 --- a/trpc_agent_sdk/agents/_base_agent.py +++ b/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,10 +262,6 @@ 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. @@ -267,7 +269,7 @@ async def run_async( # 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 diff --git a/trpc_agent_sdk/agents/core/_llm_processor.py b/trpc_agent_sdk/agents/core/_llm_processor.py index b0ee51c8e..ecbfb87ac 100644 --- a/trpc_agent_sdk/agents/core/_llm_processor.py +++ b/trpc_agent_sdk/agents/core/_llm_processor.py @@ -18,6 +18,8 @@ from typing import AsyncGenerator from typing import Optional +from opentelemetry import trace + from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.events import Event from trpc_agent_sdk.log import logger @@ -28,6 +30,8 @@ from trpc_agent_sdk.telemetry import report_call_llm from trpc_agent_sdk.telemetry import trace_call_llm from trpc_agent_sdk.telemetry import tracer +from trpc_agent_sdk.types import Content +from trpc_agent_sdk.types import Part class LlmProcessor: @@ -83,9 +87,12 @@ async def call_llm_async(self, # Step 2: Call the model and process responses with telemetry tracing. terminal_event: Optional[Event] = None - with tracer.start_as_current_span('call_llm'): + with tracer.start_as_current_span('call_llm') as call_llm_span: event_id = Event.new_id() final_llm_response = None + latest_llm_response = None + partial_thought_parts: list[str] = [] + partial_text_parts: list[str] = [] aggregated_raw_function_calls: list[dict] = [] aggregated_event_function_calls: list[dict] = [] instruction = getattr(context.agent, 'instruction', None) @@ -100,17 +107,44 @@ def _append_function_calls(target: list[dict], calls: list) -> None: "args": getattr(call, "args", None), }) + def _build_interrupted_content() -> Optional[Content]: + """Join streamed partial deltas for the interrupted call_llm trace.""" + parts: list[Part] = [] + thought_text = "".join(partial_thought_parts) + visible_text = "".join(partial_text_parts) + if thought_text: + thought_part = Part(text=thought_text) + thought_part.thought = True + parts.append(thought_part) + if visible_text: + parts.append(Part(text=visible_text)) + if latest_llm_response is not None and latest_llm_response.content is not None: + for part in latest_llm_response.content.parts or []: + if part.function_call: + parts.append(part) + if parts: + return Content(role="model", parts=parts) + if latest_llm_response is not None and latest_llm_response.content is not None: + return latest_llm_response.content + return None + t_start = time.monotonic() t_first_token: Optional[float] = None metrics_error_type: Optional[str] = None try: async for llm_response in self.model.generate_async(request, stream=stream, ctx=context): + latest_llm_response = llm_response if t_first_token is None and llm_response.has_content(): t_first_token = time.monotonic() # Collect raw model-level function calls from every chunk. raw_calls = [] if llm_response.content and llm_response.content.parts: for part in llm_response.content.parts: + if llm_response.partial and part.text: + if part.thought: + partial_thought_parts.append(part.text) + else: + partial_text_parts.append(part.text) if part.function_call: raw_calls.append(part.function_call) _append_function_calls(aggregated_raw_function_calls, raw_calls) @@ -130,40 +164,45 @@ def _append_function_calls(target: list[dict], calls: list) -> None: if not llm_response.partial: final_llm_response = llm_response - # Trace before yielding because consumers stop - # immediately after receiving an error event. - trace_call_llm( - context, - event_id, - request, - llm_response, - instruction_metadata=instruction_metadata, - stream_function_calls_raw=aggregated_raw_function_calls, - stream_function_calls_post_planner=aggregated_event_function_calls, - ) terminal_event = event # Finish the model stream and exit the span context # before exposing the terminal event downstream. continue yield event + except GeneratorExit: + metrics_error_type = "LlmCallGeneratorExit" + final_llm_response = LlmResponse( + content=_build_interrupted_content(), + partial=True, + error_code=metrics_error_type, + error_message="LLM call stopped with GeneratorExit.", + interrupted=True, + ) + raise except Exception as ex: metrics_error_type = type(ex).__name__ - trace_call_llm( - context, - event_id, - request, - LlmResponse( - error_code="LLM_CALL_ERROR", - error_message=str(ex), - custom_metadata={"error_type": type(ex).__name__}, - ), - instruction_metadata=instruction_metadata, - stream_function_calls_raw=aggregated_raw_function_calls, - stream_function_calls_post_planner=aggregated_event_function_calls, + final_llm_response = LlmResponse( + content=_build_interrupted_content(), + partial=True, + error_code="LLM_CALL_ERROR", + error_message=str(ex), + custom_metadata={"error_type": type(ex).__name__}, ) raise finally: + response_for_trace = final_llm_response or latest_llm_response or LlmResponse() + with trace.use_span(call_llm_span, end_on_exit=False): + trace_call_llm( + context, + event_id, + request, + response_for_trace, + instruction_metadata=instruction_metadata, + stream_function_calls_raw=aggregated_raw_function_calls, + stream_function_calls_post_planner=aggregated_event_function_calls, + ) + duration_s = time.monotonic() - t_start ttft_s = (t_first_token - t_start) if t_first_token is not None else duration_s report_call_llm( diff --git a/trpc_agent_sdk/runners.py b/trpc_agent_sdk/runners.py index 88e2f7e90..7ee3f00b4 100644 --- a/trpc_agent_sdk/runners.py +++ b/trpc_agent_sdk/runners.py @@ -19,6 +19,8 @@ from typing import Callable from typing import Optional +from opentelemetry import trace + from trpc_agent_sdk import cancel from trpc_agent_sdk.agents import BaseAgent from trpc_agent_sdk.artifacts import BaseArtifactService @@ -34,6 +36,7 @@ from trpc_agent_sdk.memory import BaseMemoryService from trpc_agent_sdk.sessions import BaseSessionService from trpc_agent_sdk.sessions import Session +from trpc_agent_sdk.telemetry import mark_span_error from trpc_agent_sdk.telemetry import tracer from trpc_agent_sdk.telemetry import trace_cancellation from trpc_agent_sdk.telemetry import trace_runner @@ -387,7 +390,7 @@ async def run_async( # 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("invocation"): + with tracer.start_as_current_span("invocation") as invocation_span: # Create default agent context if not provided if agent_context is None: agent_context = new_agent_context() @@ -454,6 +457,7 @@ async def run_async( # Track accumulated partial text for cancellation handling temp_text_parts: list[str] = [] + runner_trace_recorded = False try: # Support multiple levels of agent transfers @@ -559,24 +563,13 @@ async def run_async( # background worker to avoid blocking request completion. await self._schedule_post_turn_processing(invocation_context=invocation_context, ) - # Compute state after runner execution - state_end = dict(session.state) - if (last_non_streaming_event and last_non_streaming_event.actions - and last_non_streaming_event.actions.state_delta): - state_end.update(last_non_streaming_event.actions.state_delta) - - # Call trace function with runner execution details - trace_runner( - app_name=self.app_name, - user_id=user_id, - session_id=session_id, - invocation_context=invocation_context, - new_message=user_message, - last_event=last_non_streaming_event, - state_begin=state_begin, - state_end=state_end, + except GeneratorExit: + mark_span_error( + invocation_span, + error_type="RunnerGeneratorExit", + description="Runner invocation stopped with GeneratorExit.", ) - + raise except RunCancelledException as ex: logger.info("Run for session %s was cancelled", session_id) logger.debug("Cancellation details: %s", ex, exc_info=True) @@ -597,18 +590,20 @@ async def run_async( await self.session_service.update_session(session=session) # Trace the cancellation event - trace_cancellation( - app_name=self.app_name, - user_id=user_id, - session_id=session_id, - invocation_context=invocation_context, - reason=str(ex), - new_message=user_message, - last_event=last_non_streaming_event, - partial_text=temp_text, - state_begin=state_begin, - state_partial=state_partial, - ) + with trace.use_span(invocation_span, end_on_exit=False): + trace_cancellation( + app_name=self.app_name, + user_id=user_id, + session_id=session_id, + invocation_context=invocation_context, + reason=str(ex), + new_message=user_message, + last_event=last_non_streaming_event, + partial_text=temp_text, + state_begin=state_begin, + state_partial=state_partial, + ) + runner_trace_recorded = True # Yield cancellation event to notify client yield AgentCancelledEvent( @@ -619,6 +614,24 @@ async def run_async( ) finally: + if not runner_trace_recorded: + state_end = dict(session.state) + if (last_non_streaming_event and last_non_streaming_event.actions + and last_non_streaming_event.actions.state_delta): + state_end.update(last_non_streaming_event.actions.state_delta) + + with trace.use_span(invocation_span, end_on_exit=False): + trace_runner( + app_name=self.app_name, + user_id=user_id, + session_id=session_id, + invocation_context=invocation_context, + new_message=user_message, + last_event=last_non_streaming_event, + state_begin=state_begin, + state_end=state_end, + ) + # Always cleanup cancellation tracking await cancel.cleanup_run( app_name=self.app_name, diff --git a/trpc_agent_sdk/telemetry/__init__.py b/trpc_agent_sdk/telemetry/__init__.py index 03bacbd44..e50d18a5e 100644 --- a/trpc_agent_sdk/telemetry/__init__.py +++ b/trpc_agent_sdk/telemetry/__init__.py @@ -11,6 +11,7 @@ from ._metrics import report_execute_tool from ._metrics import report_invoke_agent from ._trace import get_trpc_agent_span_name +from ._trace import mark_span_error from ._trace import set_trpc_agent_span_name from ._trace import trace_agent from ._trace import trace_call_llm @@ -26,6 +27,7 @@ "report_call_llm", "report_execute_tool", "report_invoke_agent", + "mark_span_error", "trace_agent", "trace_call_llm", "trace_cancellation", diff --git a/trpc_agent_sdk/telemetry/_trace.py b/trpc_agent_sdk/telemetry/_trace.py index 6c18a59ea..2a1af505b 100644 --- a/trpc_agent_sdk/telemetry/_trace.py +++ b/trpc_agent_sdk/telemetry/_trace.py @@ -62,6 +62,21 @@ def get_trpc_agent_span_name() -> str: return _trpc_agent_span_name +def mark_span_error(span: trace.Span, error_type: str, description: str) -> None: + """Mark a span as failed with an operation-specific error. + + The caller supplies an error type and description that identify the failed + operation. + + Args: + span: The failed operation's span. + error_type: The operation-specific error type. + description: The human-readable error description. + """ + span.set_status(trace.StatusCode.ERROR, description) + span.set_attribute("error.type", error_type) + + def _join_parts_with_thought_tag(parts) -> str: """Join part texts, wrapping thought parts in tags. @@ -457,6 +472,9 @@ def trace_call_llm( llm_response_json, ) + # call_llm observation output is always the LlmResponse JSON above. On + # error, also set ERROR status (with description) and error attributes, but + # skip exception events so exporters keep using llm_response as output. error_code = getattr(llm_response, "error_code", None) if error_code: error_message = getattr(llm_response, "error_message", None) @@ -464,19 +482,11 @@ def trace_call_llm( error_type = custom_metadata.get("error_type") if isinstance(custom_metadata, dict) else None error_type = str(error_type or error_code) status_description = str(error_message or error_code) - - span.set_status(trace.StatusCode.ERROR, status_description) - span.set_attribute("error.type", error_type) + mark_span_error(span, error_type, status_description) span.set_attribute(f"{_trpc_agent_span_name}.llm.error_code", str(error_code)) if error_message: span.set_attribute(f"{_trpc_agent_span_name}.llm.error_message", str(error_message)) - exception_attributes = { - "exception.type": error_type, - "exception.message": status_description, - } - span.add_event("exception", exception_attributes) - if stream_function_calls_raw: span.set_attribute( f"{_trpc_agent_span_name}.stream_function_calls.raw",