diff --git a/trpc_agent_sdk/agents/_base_agent.py b/trpc_agent_sdk/agents/_base_agent.py index 322268d3..473bc730 100644 --- a/trpc_agent_sdk/agents/_base_agent.py +++ b/trpc_agent_sdk/agents/_base_agent.py @@ -21,6 +21,7 @@ from __future__ import annotations +import asyncio import time from abc import abstractmethod from functools import partial @@ -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 @@ -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): diff --git a/trpc_agent_sdk/runners.py b/trpc_agent_sdk/runners.py index 7c26a230..da83db37 100644 --- a/trpc_agent_sdk/runners.py +++ b/trpc_agent_sdk/runners.py @@ -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 @@ -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 @@ -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 @@ -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: @@ -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) @@ -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 diff --git a/trpc_agent_sdk/telemetry/_trace.py b/trpc_agent_sdk/telemetry/_trace.py index 9596c908..074f6c0f 100644 --- a/trpc_agent_sdk/telemetry/_trace.py +++ b/trpc_agent_sdk/telemetry/_trace.py @@ -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. @@ -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 = "" @@ -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 @@ -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(