diff --git a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py index ee4b19f3..722c9a10 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py @@ -66,6 +66,7 @@ class _ActiveToolSpan: raw_name: str display_name: str input: Any + dispatch_key: tuple[str, str] tool_use_id: str | None = None parent_tool_use_id: str | None = None handler_active: bool = False @@ -356,10 +357,11 @@ def start_tool_spans(self, message: Any, llm_span_export: str | None) -> None: raw_name=parsed_tool_name.raw_name, display_name=parsed_tool_name.display_name, input=tool_input, + dispatch_key=_make_dispatch_key(parsed_tool_name.display_name, tool_input), tool_use_id=tool_use_id, parent_tool_use_id=message_parent_tool_use_id, ) - dispatch_key = _make_dispatch_key(parsed_tool_name.display_name, tool_input) + dispatch_key = self._active_spans[tool_use_id].dispatch_key self._dispatch_queues.setdefault(dispatch_key, collections.deque()).append(tool_use_id) def finish_tool_spans(self, message: Any) -> None: @@ -458,7 +460,7 @@ def _end_tool_span( self._completed_span_exports[tool_use_id] = active_tool_span.span.export() # Remove from dispatch queue so stale entries don't accumulate. - dispatch_key = _make_dispatch_key(active_tool_span.display_name, active_tool_span.input) + dispatch_key = active_tool_span.dispatch_key queue = self._dispatch_queues.get(dispatch_key) if queue: try: @@ -668,6 +670,7 @@ class _AgentContext: llm_span: Any | None = None llm_parent_export: str | None = None llm_output: list[dict[str, Any]] | None = None + llm_output_logged: bool = False llm_message_id: str | None = None next_llm_start: float | None = None task_span: Any | None = None @@ -985,7 +988,6 @@ def _start_or_merge_llm_span( ctx: _AgentContext, ) -> tuple[dict[str, Any] | None, bool]: """Start one LLM span per provider message ID and merge its snapshots.""" - current_message = _serialize_assistant_message(message) message_id = getattr(message, "message_id", None) if not isinstance(message_id, str): message_id = self._latest_partial_message_id_by_parent.get(getattr(message, "parent_tool_use_id", None)) @@ -1001,15 +1003,18 @@ def _start_or_merge_llm_span( ctx.llm_span and same_provider_message and ctx.llm_parent_export == parent_export - and current_message is not None + and hasattr(message, "content") ): + current_message = _serialize_assistant_message(message) + if current_message is None: + return None, False merged = _merge_assistant_messages( ctx.llm_output[0] if ctx.llm_output else None, current_message, ) if merged is not None: ctx.llm_output = [merged] - ctx.llm_span.log(output=ctx.llm_output) + ctx.llm_output_logged = False self._log_assistant_usage(message, ctx, message_id) return merged, True @@ -1032,6 +1037,8 @@ def _start_or_merge_llm_span( ctx.llm_span = span ctx.llm_parent_export = parent_export ctx.llm_output = [final_content] if final_content is not None else None + # start_span already logged this initial output; later merges mark it dirty. + ctx.llm_output_logged = True ctx.llm_message_id = message_id if isinstance(message_id, str) else None ctx.next_llm_start = None self._log_assistant_usage(message, ctx, ctx.llm_message_id) @@ -1059,9 +1066,13 @@ def _log_assistant_usage( def _mark_output_usage_unknown(self, ctx: _AgentContext) -> None: """Mark missing output only once an LLM span is finalized.""" - if ctx.llm_span is None or ctx.llm_message_id in self._final_output_usage_message_ids: + if ctx.llm_span is None: return - ctx.llm_span.log(metadata={"usage_output_tokens_unknown": True}) + if ctx.llm_output is not None and not ctx.llm_output_logged: + ctx.llm_span.log(output=ctx.llm_output) + ctx.llm_output_logged = True + if ctx.llm_message_id not in self._final_output_usage_message_ids: + ctx.llm_span.log(metadata={"usage_output_tokens_unknown": True}) def _process_task_event(self, message: Any, agent_span_export: str | None) -> None: """Handle TaskStarted / TaskProgress / TaskNotification system messages.""" @@ -1538,11 +1549,16 @@ def _create_llm_span_for_messages( model = getattr(last_message, "model", None) input_messages = _build_llm_input(prompt, conversation_history) - outputs: list[dict[str, Any]] = [] - for msg in messages: - if hasattr(msg, "content"): - content = _serialize_content_blocks(msg.content) - outputs.append({"content": content, "role": "assistant"}) + serialized_messages = [ + { + "content": _serialize_content_blocks(msg.content), + "role": "assistant", + } + if hasattr(msg, "content") + else None + for msg in messages + ] + outputs = [serialized for serialized in serialized_messages if serialized is not None] metadata: dict[str, Any] = {"provider": "anthropic"} if model: @@ -1559,11 +1575,8 @@ def _create_llm_span_for_messages( ) # Return final message content for conversation history and the span - if hasattr(last_message, "content"): - content = _serialize_content_blocks(last_message.content) - return {"content": content, "role": "assistant"}, llm_span - - return None, llm_span + final_content = serialized_messages[-1] if serialized_messages else None + return final_content, llm_span def _serialize_assistant_message(message: Any) -> dict[str, Any] | None: