fix(litellm): copy streaming latency attrs to canonical attempt span

packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py

@@ -292,6 +292,38 @@ def _maybe_warn_retry_attempt_eviction(evicted: int) -> None:
     )
 
 
+def _set_active_retry_attempt_attribute(parent, key: str, value) -> None:
+    """Copy wrapper-owned attrs onto the canonical retry_attempt span.
+
+    LiteLLM streaming latency is measured by FR's safety wrapper while backend
+    dedup makes ``fortifyroot.litellm.attempt_N`` the canonical usage span.
+    Storing these attrs on the attempt span avoids relying on wrapper/attempt
+    spans arriving in the same OTLP batch.
+    """
+    if parent is None:
+        return
+
+    with _FR_RETRY_ATTEMPT_MAP_LOCK:
+        entries = [
+            entry
+            for entry in _FR_RETRY_ATTEMPT_MAP.values()
+            if entry.get("parent") is parent and not entry.get("ended")
+        ]
+        for entry in entries:
+            streaming_attrs = entry.setdefault("streaming_attrs", {})
+            streaming_attrs[key] = value
+            span = entry.get("span")
+            if span is None:
+                continue
+            try:
+                span.set_attribute(key, value)
+            except Exception:
+                logger.debug(
+                    "Failed to set streaming latency on retry_attempt span",
+                    exc_info=True,
+                )
+
+
 def _resolve_routed_provider(kwargs) -> Optional[str]:
     """Best-effort: derive the ROUTED provider (e.g. ``openai``) from
     LiteLLM kwargs for the retry_attempt span's ``gen_ai.system``
@@ -568,6 +600,8 @@ def _finalize_retry_attempt_span(
 
     try:
         if success:
+            for key, value in (entry.get("streaming_attrs") or {}).items():
+                span.set_attribute(key, value)
             response_model = get_object_value(response_obj, "model")
             if response_model:
                 span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_MODEL, str(response_model))
@@ -882,6 +916,7 @@ def _invoke_completion(tracer, wrapped, args, kwargs, *, is_text_completion=Fals
             span_name,
             _set_response_attributes,
             token=span_token,
+            set_canonical_span_attribute=_set_active_retry_attempt_attribute,
         )
 
     context_api.detach(span_token)
@@ -942,6 +977,7 @@ async def _invoke_acompletion(
             span_name,
             _set_response_attributes,
             token=span_token,
+            set_canonical_span_attribute=_set_active_retry_attempt_attribute,
         )
 
     context_api.detach(span_token)

packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py

@@ -38,7 +38,15 @@ def is_async_streaming_response(kwargs, response) -> bool:
     return False
 
 
-def wrap_sync_streaming_response(span, response, request_type, span_name, set_response_attributes, token=None):
+def wrap_sync_streaming_response(
+    span,
+    response,
+    request_type,
+    span_name,
+    set_response_attributes,
+    token=None,
+    set_canonical_span_attribute=None,
+):
     """Wrap a sync streaming response with per-chunk safety and span lifecycle.
 
     ``token`` is the OTel context token established by ``_invoke_completion``
@@ -66,13 +74,15 @@ def wrap_sync_streaming_response(span, response, request_type, span_name, set_re
                     span,
                     FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
                     first_token_at - stream_started_at,
+                    set_canonical_span_attribute,
                 )
             yield chunk
         if first_token_at is not None:
             _set_streaming_latency_attr(
                 span,
                 FR_STREAMING_TIME_TO_GENERATE_MS,
                 time.time() - first_token_at,
+                set_canonical_span_attribute,
             )
         _finalize_streaming_span(span, complete_response, set_response_attributes)
     except Exception as exc:
@@ -93,6 +103,7 @@ async def wrap_async_streaming_response(
     span_name,
     set_response_attributes,
     token=None,
+    set_canonical_span_attribute=None,
 ):
     """Wrap an async streaming response with per-chunk safety and span lifecycle.
 
@@ -119,13 +130,15 @@ async def wrap_async_streaming_response(
                     span,
                     FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
                     first_token_at - stream_started_at,
+                    set_canonical_span_attribute,
                 )
             yield chunk
         if first_token_at is not None:
             _set_streaming_latency_attr(
                 span,
                 FR_STREAMING_TIME_TO_GENERATE_MS,
                 time.time() - first_token_at,
+                set_canonical_span_attribute,
             )
         _finalize_streaming_span(span, complete_response, set_response_attributes)
     except Exception as exc:
@@ -248,9 +261,17 @@ async def _aclose_streaming_response(response) -> None:
         await aclose()
 
 
-def _set_streaming_latency_attr(span, key: str, seconds: float) -> None:
+def _set_streaming_latency_attr(
+    span,
+    key: str,
+    seconds: float,
+    set_canonical_span_attribute=None,
+) -> None:
+    value = int(round(seconds * 1000))
     if span.is_recording():
-        span.set_attribute(key, int(round(seconds * 1000)))
+        span.set_attribute(key, value)
+    if set_canonical_span_attribute is not None:
+        set_canonical_span_attribute(span, key, value)
 
 
 def _chunk_has_output_text(chunk) -> bool:

packages/opentelemetry-instrumentation-litellm/tests/test_retry_attempt_emission.py

@@ -38,6 +38,11 @@
     _resolve_routed_provider,
     _start_retry_attempt_span,
     _finalize_retry_attempt_span,
+    _set_active_retry_attempt_attribute,
+)
+from opentelemetry.instrumentation.litellm.streaming_safety import (
+    FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
+    FR_STREAMING_TIME_TO_GENERATE_MS,
 )
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
@@ -237,6 +242,12 @@ def test_single_attempt_emits_one_retry_attempt_under_parent(fresh_tracer):
             "messages": [{"role": "user", "content": "hello masked [EMAIL]"}],
         }
         _start_retry_attempt_span(kwargs)
+        _set_active_retry_attempt_attribute(
+            parent, FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, 123
+        )
+        _set_active_retry_attempt_attribute(
+            parent, FR_STREAMING_TIME_TO_GENERATE_MS, 456
+        )
 
         # Mock response object with usage/model.
         class MockUsage:
@@ -272,6 +283,8 @@ class MockResponse:
     assert retry_span.attributes.get("gen_ai.response.id") == "resp-abc"
     assert retry_span.attributes.get("gen_ai.usage.input_tokens") == 10
     assert retry_span.attributes.get("gen_ai.usage.output_tokens") == 5
+    assert retry_span.attributes.get(FR_STREAMING_TIME_TO_FIRST_TOKEN_MS) == 123
+    assert retry_span.attributes.get(FR_STREAMING_TIME_TO_GENERATE_MS) == 456
     assert retry_span.attributes.get("gen_ai.prompt.0.role") == "user"
     assert retry_span.attributes.get("gen_ai.prompt.0.content") == "hello masked [EMAIL]"
 

packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py

@@ -769,6 +769,7 @@ def test_streaming_chunk_output_detection_handles_delta_message_and_text():
 def test_sync_streaming_wrapper_sets_streaming_latency_attrs():
     exporter, tracer = _test_tracer()
     span = tracer.start_span("litellm.completion")
+    canonical_attrs = []
     chunks = [
         SimpleNamespace(
             choices=[
@@ -793,18 +794,24 @@ def test_sync_streaming_wrapper_sets_streaming_latency_attrs():
                 "chat",
                 "litellm.completion",
                 lambda *_: None,
+                set_canonical_span_attribute=lambda *args: canonical_attrs.append(args),
             )
         )
 
     attrs = exporter.get_finished_spans()[0].attributes
     assert attrs[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 250
     assert attrs[FR_STREAMING_TIME_TO_GENERATE_MS] == 650
+    assert canonical_attrs == [
+        (span, FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, 250),
+        (span, FR_STREAMING_TIME_TO_GENERATE_MS, 650),
+    ]
 
 
 @pytest.mark.asyncio
 async def test_async_streaming_wrapper_sets_streaming_latency_attrs():
     exporter, tracer = _test_tracer()
     span = tracer.start_span("litellm.completion")
+    canonical_attrs = []
 
     async def _response():
         yield SimpleNamespace(
@@ -823,13 +830,18 @@ async def _response():
                 "chat",
                 "litellm.completion",
                 lambda *_: None,
+                set_canonical_span_attribute=lambda *args: canonical_attrs.append(args),
             )
         ]
 
     assert len(yielded) == 1
     attrs = exporter.get_finished_spans()[0].attributes
     assert attrs[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 125
     assert attrs[FR_STREAMING_TIME_TO_GENERATE_MS] == 375
+    assert canonical_attrs == [
+        (span, FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, 125),
+        (span, FR_STREAMING_TIME_TO_GENERATE_MS, 375),
+    ]
 
 
 def test_streaming_wrapper_leaves_latency_attrs_unset_for_empty_stream():