Skip to content

fix: 修复 agent 执行中断场景下 span 状态与输出丢失的 4 类问题 #300

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
raychen911 merged 1 commit into from
Aug 19, 2026
+187 −67
Merged

fix: 修复 agent 执行中断场景下 span 状态与输出丢失的 4 类问题 #300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 44 additions & 5 deletions trpc_agent_sdk/agents/_base_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

from __future__ import annotations

import asyncio
import time
from abc import abstractmethod
from functools import partial
Expand Down Expand Up @@ -283,24 +284,57 @@ async def run_async(
# Track all non-partial events for building action trace
non_partial_events = []

# Track accumulated partial text as it streams in (mirrors the
# pattern used by LlmProcessor.call_llm_async's
# _build_interrupted_content). If GeneratorExit/CancelledError
# fires before any non-partial event exists, non_partial_events
# is empty and _build_action_string_from_events([]) would
# otherwise produce "", silently dropping everything that was
# already streamed to the caller.
partial_text_parts: list[str] = []

mono_start = time.monotonic()
t_first_visible: Optional[float] = None
error_type: Optional[str] = None
error_message: Optional[str] = None
interrupted_partial_text: Optional[str] = None

try:
gen_co = run_stream_filters(ctx.agent_context, None, self.filters, handle) # type: ignore
async for event in gen_co:
if t_first_visible is None and event.has_content():
t_first_visible = time.monotonic()
if not event.partial and event.content is not None:
# Collect non-partial events with content for tracing
# This excludes state update events which have content=None
non_partial_events.append(event)
if event.partial:
if event.content and event.content.parts:
for part in event.content.parts:
if part.text:
partial_text_parts.append(part.text)
else:
# A non-partial event finalizes this turn's output;
# any partial text accumulated so far has now been
# superseded by it, so drop it to avoid duplicating
# output already captured in non_partial_events.
partial_text_parts.clear()
if event.content is not None:
# Collect non-partial events with content for tracing
# This excludes state update events which have content=None
non_partial_events.append(event)
yield event # type: ignore
except GeneratorExit:
error_type = "AgentGeneratorExit"
error_message = "Agent execution stopped with GeneratorExit."
interrupted_partial_text = "".join(partial_text_parts)
raise
except asyncio.CancelledError:
# Like GeneratorExit, asyncio.CancelledError subclasses
# BaseException, not Exception, so it is not caught by
# `except Exception` below. External cancellation
# (task.cancel(), asyncio.wait_for() timeout, ASGI
# disconnect) surfaces here as this exception; salvage the
# partial text the same way as the GeneratorExit branch.
error_type = "AgentCancelledError"
error_message = "Agent execution stopped with asyncio.CancelledError."
interrupted_partial_text = "".join(partial_text_parts)
raise
except RunLimitException as ex:
error_type = ex.error_code
Expand All @@ -314,8 +348,13 @@ async def run_async(
# Compute state after agent run
state_end = dict(ctx.session.state)

# Build formatted action string from all non-partial events
# Build formatted action string from all non-partial events.
# Fall back to the accumulated (but never finalized) partial
# text when the run was interrupted before producing any
# non-partial event, so streamed output is not silently lost.
agent_action = _build_action_string_from_events(non_partial_events)
if not agent_action and interrupted_partial_text:
agent_action = f"[INTERRUPTED]\n{interrupted_partial_text}"

# Call trace function with agent execution details
with trace.use_span(agent_span, end_on_exit=False):
Expand Down
168 changes: 114 additions & 54 deletions trpc_agent_sdk/runners.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,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 get_trpc_agent_span_name
from trpc_agent_sdk.telemetry import tracer
from trpc_agent_sdk.telemetry import trace_cancellation
from trpc_agent_sdk.telemetry import trace_runner
Expand Down Expand Up @@ -395,66 +396,103 @@ async def run_async(
# an async generator is cancelled, but try/finally always executes
# even under CancelledError (PEP 492).
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()

session = await self.session_service.get_session(app_name=self.app_name,
user_id=user_id,
session_id=session_id,
agent_context=agent_context)
if not session:
# Create new session if not found - use create_session instead of save_session
session = await self.session_service.create_session(app_name=self.app_name,
user_id=user_id,
session_id=session_id,
agent_context=agent_context)
logger.debug("Created new session: %s", session.id)
else:
logger.debug("Using existing session: %s with %s events", session.id, len(session.events))

# Capture state before runner execution
state_begin = dict(session.state)

session.conversation_count += 1
history_content: Content | None = None
if isinstance(new_message, list):
user_message = new_message[-1]
history_content = new_message[0]
else:
user_message = new_message
invocation_context = self._new_invocation_context(
session,
new_message=user_message,
run_config=run_config,
agent_context=agent_context,
trpc_span_name = get_trpc_agent_span_name()
# Seed the baseline runner.name attribute immediately, before any
# initialization below runs. Session lookup/creation,
# InvocationContext construction, saving the user message, and
# cancel.register_run() can all raise before the main try/finally
# block further down is reached. When that happens, the except
# branch below calls trace_runner(invocation_context=None), which
# (by design) skips runner.name since there is no
# InvocationContext yet - so this is the only place that
# attribute gets set on the init-failure path. All other
# trace_runner() attributes (gen_ai.system, gen_ai.operation.name,
# runner.app_name/user_id/session_id) are unconditional inside
# trace_runner() itself, so seeding them here would just be
# overwritten with the same value and is intentionally omitted.
invocation_span.set_attribute(
f"{trpc_span_name}.runner.name",
f"[trpc-agent]: {self.app_name}/{self.agent.name}",
)
root_agent = self.agent
if run_config.save_history_enabled and history_content:
history_event = Event(content=history_content,
invocation_id=invocation_context.invocation_id,
id=Event.new_id(),
author=root_agent.name)
await self.session_service.append_event(session=session, event=history_event)

if user_message:
await self._append_new_message_to_session(

try:
# Create default agent context if not provided
if agent_context is None:
agent_context = new_agent_context()

session = await self.session_service.get_session(app_name=self.app_name,
user_id=user_id,
session_id=session_id,
agent_context=agent_context)
if not session:
# Create new session if not found - use create_session instead of save_session
session = await self.session_service.create_session(app_name=self.app_name,
user_id=user_id,
session_id=session_id,
agent_context=agent_context)
logger.debug("Created new session: %s", session.id)
else:
logger.debug("Using existing session: %s with %s events", session.id, len(session.events))

# Capture state before runner execution
state_begin = dict(session.state)

session.conversation_count += 1
history_content: Content | None = None
if isinstance(new_message, list):
user_message = new_message[-1]
history_content = new_message[0]
else:
user_message = new_message
invocation_context = self._new_invocation_context(
session,
user_message,
invocation_context,
new_message=user_message,
run_config=run_config,
agent_context=agent_context,
)
root_agent = self.agent
if run_config.save_history_enabled and history_content:
history_event = Event(content=history_content,
invocation_id=invocation_context.invocation_id,
id=Event.new_id(),
author=root_agent.name)
await self.session_service.append_event(session=session, event=history_event)

if user_message:
await self._append_new_message_to_session(
session,
user_message,
invocation_context,
)

invocation_context.agent = self._find_agent_to_run(session, root_agent, run_config)
invocation_context.agent = self._find_agent_to_run(session, root_agent, run_config)

# Register for cancellation tracking
session_key = await cancel.register_run(
app_name=self.app_name,
user_id=user_id,
session_id=session_id,
)
# Register for cancellation tracking
session_key = await cancel.register_run(
app_name=self.app_name,
user_id=user_id,
session_id=session_id,
)

# Store session_key in invocation_context for use by agents
invocation_context.session_key = session_key
# Store session_key in invocation_context for use by agents
invocation_context.session_key = session_key
except Exception as ex:
# Initialization failed before the main try/finally block
# below was reached, so there is no InvocationContext/Session
# to feed into trace_runner() yet. Call trace_runner() with
# invocation_context=None so the span still gets the baseline
# business attributes plus the error status, instead of being
# left as "unknown".
trace_runner(
app_name=self.app_name,
user_id=user_id,
session_id=session_id,
invocation_context=None,
new_message=(new_message[-1] if isinstance(new_message, list) else new_message),
error_type=type(ex).__name__,
error_message=str(ex),
)
raise

# Track the last non-streaming event for tracing
last_non_streaming_event = None
Expand All @@ -464,6 +502,7 @@ async def run_async(
runner_trace_recorded = False
trace_error_type: Optional[str] = None
trace_error_message: Optional[str] = None
trace_partial_text: Optional[str] = None

try:
# Support multiple levels of agent transfers
Expand Down Expand Up @@ -572,6 +611,7 @@ async def run_async(
except GeneratorExit:
trace_error_type = "RunnerGeneratorExit"
trace_error_message = "Runner invocation stopped with GeneratorExit."
trace_partial_text = "".join(temp_text_parts)
raise

except RunLimitException as ex:
Expand Down Expand Up @@ -628,6 +668,25 @@ async def run_async(
branch=invocation_context.branch,
)

except asyncio.CancelledError as ex:
# asyncio.CancelledError subclasses BaseException, not
# Exception, so it is NOT caught by `except Exception` below.
# This is how external cancellation actually happens in
# practice - task.cancel(), asyncio.wait_for() timing out, or
# an ASGI client disconnecting - as opposed to the SDK's
# cooperative cancel_run_async() + raise_if_cancelled()
# checkpoints (handled above via RunCancelledException).
# Without this branch, none of the except clauses above run,
# trace_error_type/trace_error_message stay None, and the
# finally: block's trace_runner() call never marks the span
# as an error nor surfaces the partial output that was
# already streamed to the caller.
trace_error_type = "CancelledError"
trace_error_message = str(ex) or "Runner invocation was cancelled (asyncio.CancelledError)."
trace_partial_text = "".join(temp_text_parts)
logger.info("Run for session %s was cancelled via asyncio.CancelledError", session_id)
raise

except Exception as ex:
trace_error_type = type(ex).__name__
trace_error_message = str(ex)
Expand All @@ -652,6 +711,7 @@ async def run_async(
state_end=state_end,
error_type=trace_error_type,
error_message=trace_error_message,
partial_text=trace_partial_text,
)

# Always cleanup cancellation tracking
Expand Down
37 changes: 29 additions & 8 deletions trpc_agent_sdk/telemetry/_trace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,13 +107,14 @@ def trace_runner(
app_name: str,
user_id: str,
session_id: str,
invocation_context: InvocationContext,
invocation_context: Optional[InvocationContext] = None,
new_message: Optional[Content] = None,
last_event: Optional[Event] = None,
state_begin: Optional[dict[str, Any]] = None,
state_end: Optional[dict[str, Any]] = None,
error_type: Optional[str] = None,
error_message: Optional[str] = None,
partial_text: Optional[str] = None,
):
"""Traces runner execution.

Expand All @@ -125,21 +126,30 @@ def trace_runner(
user_id: The user ID of the session.
session_id: The session ID of the session.
invocation_context: The invocation context for the current agent run.
May be ``None`` when the runner fails before an invocation
context can be constructed (e.g. during initialization). In
that case, attributes that depend on it are skipped.
new_message: The new message that started this invocation.
last_event: The last non-streaming event from the agent execution.
state_begin: The state before the runner execution.
state_end: The state after the runner execution.
error_type: The error type when the runner does not complete normally.
error_message: The error message when the runner does not complete normally.
partial_text: Accumulated partial (streamed but not yet finalized) text.
Used as a fallback for the ``runner.output`` attribute when the
invocation was interrupted (e.g. GeneratorExit or an external
asyncio.CancelledError) before any non-streaming event existed,
so the already-streamed output is not silently lost.
"""
span = trace.get_current_span()
span.set_attribute("gen_ai.system", _trpc_agent_span_name)
span.set_attribute("gen_ai.operation.name", "run_runner")
span.set_attribute(f"{_trpc_agent_span_name}.runner.app_name", app_name)
span.set_attribute(
f"{_trpc_agent_span_name}.runner.name",
f"[trpc-agent]: {app_name}/{invocation_context.agent.name}",
)
if invocation_context is not None:
span.set_attribute(
f"{_trpc_agent_span_name}.runner.name",
f"[trpc-agent]: {app_name}/{invocation_context.agent.name}",
)
span.set_attribute(f"{_trpc_agent_span_name}.runner.user_id", user_id)
span.set_attribute(f"{_trpc_agent_span_name}.runner.session_id", session_id)
input_str = ""
Expand All @@ -149,6 +159,8 @@ def trace_runner(
output_str = ""
if last_event and last_event.content and last_event.content.parts:
output_str = _join_parts_with_thought_tag(last_event.content.parts)
elif partial_text:
output_str = f"[INTERRUPTED]\n{partial_text}"
span.set_attribute(f"{_trpc_agent_span_name}.runner.output", output_str)

# Set state attributes for begin and end
Expand Down Expand Up @@ -485,9 +497,18 @@ def trace_call_llm(
llm_response_json,
)

if error_type:
span.set_status(trace.StatusCode.ERROR, error_message or error_type)
span.set_attribute("error.type", error_type)
# The caller-supplied error_type reflects an exception that propagated out
# of the model call. But the SDK-managed retry layer (retry_model_call)
# can also swallow a raised exception and yield a normal-looking
# LlmResponse with error_code/error_message set instead of re-raising
# (see models/_retry.py:_build_error_response). Fall back to inspecting
# llm_response.error_code so those calls are still marked as errors here.
effective_error_type = error_type or llm_response.error_code
effective_error_message = error_message or llm_response.error_message

if effective_error_type:
span.set_status(trace.StatusCode.ERROR, effective_error_message or effective_error_type)
span.set_attribute("error.type", effective_error_type)

if stream_function_calls_raw:
span.set_attribute(
Expand Down
Loading