better output marker logic

Author: AbhiPrasad

py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py

@@ -263,6 +263,10 @@ async def calculator_handler(args):
         assert len(ordered_llm_spans) == len(expected_partial_usage)
         for llm_span, usage in zip(ordered_llm_spans, expected_partial_usage, strict=True):
             assert llm_span["metrics"] | _metrics_from_exact_anthropic_usage(usage) == llm_span["metrics"]
+        llm_spans_with_output = [span for span in llm_spans if "completion_tokens" in span["metrics"]]
+        assert llm_spans_with_output
+        for llm_span in llm_spans_with_output:
+            assert "usage_output_tokens_unknown" not in llm_span.get("metadata", {})
     elif _sdk_version_at_least("0.1.11"):
         expected_usage_by_message_id: dict[str, dict[str, Any]] = {}
         for message in received_messages:

py/src/braintrust/integrations/claude_agent_sdk/tracing.py

@@ -744,6 +744,7 @@ def cleanup(self) -> None:
         """End all open LLM spans, TASK spans, and TOOL spans; clear thread-local."""
         for ctx in self._contexts.values():
             if ctx.llm_span:
+                self._mark_output_usage_unknown(ctx)
                 ctx.llm_span.end()
                 ctx.llm_span = None
             if ctx.task_span:
@@ -884,6 +885,8 @@ def _handle_stream_event(self, message: Any) -> None:
 
     def _handle_result(self, message: Any) -> None:
         self._active_key = None
+        for ctx in self._contexts.values():
+            self._mark_output_usage_unknown(ctx)
         result_value = getattr(message, "result", None)
         if result_value is not None:
             self._result_output = result_value
@@ -1014,6 +1017,7 @@ def _start_or_merge_llm_span(
         first_token_time = time.time()
 
         if ctx.llm_span:
+            self._mark_output_usage_unknown(ctx)
             ctx.llm_span.end(end_time=resolved_start)
 
         final_content, span = _create_llm_span_for_messages(
@@ -1051,10 +1055,14 @@ def _log_assistant_usage(
         has_final_output = message_id in self._final_output_usage_message_ids if message_id else False
         metrics, _ = extract_anthropic_usage(usage, include_output=has_final_output)
         _, metadata = extract_anthropic_usage(raw_message_usage, include_output=False)
-        if not has_final_output and "prompt_tokens" in metrics:
-            metadata["usage_output_tokens_unknown"] = True
         ctx.llm_span.log(metrics=metrics or None, metadata=metadata or None)
 
+    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:
+            return
+        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."""
         task_id = _msg_field(message, "task_id")