Fix 3 OpenTelemetry tracing bugs in proxy integration (#27757)

1. Missing litellm_request child span when proxy parent in metadata:
   _get_span_context now returns (ctx, None) for the metadata-injected
   proxy parent so the primary span is always emitted as a child of ctx.
   Proxy span lifecycle managed by new _end_proxy_span_from_kwargs.

2. open_telemetry_logger overwrite by later handlers:
   _init_otel_logger_on_litellm_proxy now uses first-registered-wins —
   only assigns proxy_server.open_telemetry_logger when currently None.

3. Duplicate litellm_request success spans in streaming paths:
   Added _mark_success_span_once with per-handler dedupe key stored in
   kwargs metadata, suppressing the second span when both sync and async
   success callbacks fire for the same request.

Co-authored-by: Yassin Kortam <yassinkortam@g.ucla.edu>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yassin Kortam
2026-05-12 18:32:05 -07:00
committed by GitHub
co-authored by Yassin Kortam Claude Opus 4.7
parent 8f25942ecf
commit 1c4e4d4a60
4 changed files with 756 additions and 16 deletions
+29
View File
@@ -888,6 +888,15 @@ def log_guardrail_information(func):
- pre_call
- during_call
- post_call
Some guardrails (e.g. ``block_code_execution``) call
``add_standard_logging_guardrail_information_to_request_data`` directly
from inside the wrapped function so they can record a richer payload
(structured detections, tracing detail) than this decorator's
"allow"/"mask"/raw-response default. To avoid double-recording in that
case (which would emit two spans, two Datadog records, two spend-log
entries, etc.), snapshot the entry count before invocation: if the
wrapped function already appended its own entry, skip the auto-record.
"""
import functools
import inspect
@@ -907,6 +916,16 @@ def log_guardrail_information(func):
return GuardrailEventHooks.post_call
return None
def _count_recorded_guardrail_entries(request_data: dict) -> int:
total = 0
for container_key in ("metadata", "litellm_metadata"):
container = request_data.get(container_key)
if isinstance(container, dict):
entries = container.get("standard_logging_guardrail_information")
if isinstance(entries, list):
total += len(entries)
return total
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
start_time = datetime.now() # Move start_time inside the wrapper
@@ -919,8 +938,11 @@ def log_guardrail_information(func):
if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
original_inputs = kwargs.get("inputs")
entries_before = _count_recorded_guardrail_entries(request_data)
try:
response = await func(*args, **kwargs)
if _count_recorded_guardrail_entries(request_data) > entries_before:
return response
return self._process_response(
response=response,
request_data=request_data,
@@ -931,6 +953,8 @@ def log_guardrail_information(func):
original_inputs=original_inputs,
)
except Exception as e:
if _count_recorded_guardrail_entries(request_data) > entries_before:
raise
return self._process_error(
e=e,
request_data=request_data,
@@ -952,8 +976,11 @@ def log_guardrail_information(func):
if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
original_inputs = kwargs.get("inputs")
entries_before = _count_recorded_guardrail_entries(request_data)
try:
response = func(*args, **kwargs)
if _count_recorded_guardrail_entries(request_data) > entries_before:
return response
return self._process_response(
response=response,
request_data=request_data,
@@ -962,6 +989,8 @@ def log_guardrail_information(func):
original_inputs=original_inputs,
)
except Exception as e:
if _count_recorded_guardrail_entries(request_data) > entries_before:
raise
return self._process_error(
e=e,
request_data=request_data,
+133 -3
View File
@@ -237,7 +237,14 @@ class OpenTelemetry(CustomLogger):
not isinstance(cb, OpenTelemetry) for cb in litellm.service_callback
):
litellm.service_callback.append(self)
setattr(proxy_server, "open_telemetry_logger", self)
# avoid proxy logger ownership being overwritten by later
# handlers. Multiple integrations (default OTEL, Langfuse OTEL,
# Arize OTEL, etc.) may initialize in sequence; without this guard,
# the last one silently replaces the first and breaks expected
# routing for proxy_server.open_telemetry_logger consumers.
# Behavior: first-registered wins.
if getattr(proxy_server, "open_telemetry_logger", None) is None:
setattr(proxy_server, "open_telemetry_logger", self)
def _get_or_create_provider(
self,
@@ -794,12 +801,100 @@ class OpenTelemetry(CustomLogger):
# End of Team/Key Based Logging Control Flow
#########################################################
def _emit_once(self, kwargs: dict, *scope: object) -> bool:
"""Return True the first time this handler is asked to emit a span
for the given (handler, scope) on this kwargs; False on repeats.
Used to suppress duplicate span emission for two distinct patterns:
1. **Handler-level dual-fire**: streaming code paths trigger both
the sync and async callback for one request, so ``_handle_success``
/ ``_handle_failure`` would otherwise produce two
``litellm_request`` spans. Scope: ``("success",)`` / ``("failure",)``.
2. **Payload-driven multi-entrypoint emission**: a span loop that
reads entries from ``standard_logging_payload`` (currently only
guardrails) is invoked from multiple lifecycle points
(post-call hooks, success callback, failure callback). The list
can be re-read with mutated entries between calls, so dedupe
must be at entry granularity. Scope: the entry's stable identity.
``scope`` parts can be any hashable identity. The marker is stored
in ``kwargs["litellm_params"]["metadata"]["_otel_internal"]`` so it
is request-local (kwargs is shared across the sync/async callbacks
and lifecycle hooks for one request).
"""
litellm_params = kwargs.get("litellm_params")
if not isinstance(litellm_params, dict):
litellm_params = {}
kwargs["litellm_params"] = litellm_params
_metadata = litellm_params.get("metadata")
if not isinstance(_metadata, dict):
_metadata = {}
litellm_params["metadata"] = _metadata
_otel_internal = _metadata.get("_otel_internal")
if not isinstance(_otel_internal, dict):
_otel_internal = {}
_metadata["_otel_internal"] = _otel_internal
spans_logged = _otel_internal.get("spans_logged")
if not isinstance(spans_logged, dict):
spans_logged = {}
_otel_internal["spans_logged"] = spans_logged
dedupe_key = (self.__class__.__name__, id(self), *scope)
if spans_logged.get(dedupe_key) is True:
return False
spans_logged[dedupe_key] = True
return True
def _end_proxy_span_from_kwargs(self, kwargs: dict, end_time) -> None:
"""Close the proxy-level parent span if it is still recording.
This helper retrieves the proxy span directly from kwargs metadata
and closes it after all child spans have been recorded.
Only called from the success path. The failure path deliberately
leaves the proxy span open so ``async_post_call_failure_hook`` can
append the ``"Failed Proxy Server Request"`` child span before
closing it.
Only spans named ``LITELLM_PROXY_REQUEST_SPAN_NAME`` are closed —
externally provided spans must not be closed by LiteLLM.
"""
litellm_params = kwargs.get("litellm_params", {}) or {}
_metadata = litellm_params.get("metadata", {}) or {}
proxy_span = _metadata.get("litellm_parent_otel_span", None)
if (
proxy_span is not None
and getattr(proxy_span, "name", None) == LITELLM_PROXY_REQUEST_SPAN_NAME
and hasattr(proxy_span, "is_recording")
and proxy_span.is_recording()
):
proxy_span.end(end_time=self._to_ns(end_time))
def _handle_success(self, kwargs, response_obj, start_time, end_time):
"""Create the litellm_request span then close the proxy span."""
verbose_logger.debug(
"OpenTelemetry Logger: Logging kwargs: %s, OTEL config settings=%s",
kwargs,
self.config,
)
# sync + async success handlers can both fire for one
# request (notably in streaming code paths). Guard against duplicate
# span writes — but still close the proxy span on the skip path so
# the trace doesn't leak an open root span.
if not self._emit_once(kwargs, "success"):
verbose_logger.debug(
"OpenTelemetry: skipping duplicate success span for handler=%s",
self.__class__.__name__,
)
self._end_proxy_span_from_kwargs(kwargs, end_time)
return
ctx, parent_span = self._get_span_context(kwargs)
if self.config.ignore_context_propagation:
@@ -859,7 +954,7 @@ class OpenTelemetry(CustomLogger):
# 6. Do NOT end parent span - it should be managed by its creator
# External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM
# However, proxy-created spans should be closed here
# However, proxy-created spans should be closed here.
if (
parent_span is not None
and hasattr(parent_span, "name")
@@ -867,6 +962,11 @@ class OpenTelemetry(CustomLogger):
):
parent_span.end(end_time=self._to_ns(end_time))
# close the proxy span explicitly from kwargs metadata
# after all child spans (litellm_request, guardrail, raw_request)
# have been fully recorded and exported.
self._end_proxy_span_from_kwargs(kwargs, end_time)
def _start_primary_span(
self,
kwargs,
@@ -1296,6 +1396,21 @@ class OpenTelemetry(CustomLogger):
for guardrail_information in guardrail_information_list:
start_time_float = guardrail_information.get("start_time")
end_time_float = guardrail_information.get("end_time")
# ``_create_guardrail_span`` is called from three lifecycle
# points (``async_post_call_success_hook``, ``_handle_success``,
# ``_handle_failure``) and re-reads the (mutating) entry list
# each time. Dedupe at entry granularity so a single real
# guardrail invocation produces exactly one span per handler.
if not self._emit_once(
kwargs,
"guardrail",
guardrail_information.get("guardrail_name"),
start_time_float,
guardrail_information.get("guardrail_mode"),
):
continue
start_time_datetime = datetime.now()
if start_time_float is not None:
start_time_datetime = datetime.fromtimestamp(start_time_float)
@@ -1349,6 +1464,21 @@ class OpenTelemetry(CustomLogger):
kwargs,
self.config,
)
# sync + async failure handlers can both fire for one
# request (notably in streaming code paths), producing two
# semantically identical ERROR spans. Unlike the success path, the
# proxy span is intentionally left open here so that
# ``async_post_call_failure_hook`` can append the
# "Failed Proxy Server Request" child span before closing it —
# there is no proxy-span side-effect to preserve on the skip path.
if not self._emit_once(kwargs, "failure"):
verbose_logger.debug(
"OpenTelemetry: skipping duplicate failure span for handler=%s",
self.__class__.__name__,
)
return
_parent_context, parent_otel_span = self._get_span_context(kwargs)
if self.config.ignore_context_propagation:
@@ -2188,7 +2318,7 @@ class OpenTelemetry(CustomLogger):
verbose_logger.debug(
"OpenTelemetry: Using explicit parent span from metadata"
)
return trace.set_span_in_context(parent_otel_span), parent_otel_span
return trace.set_span_in_context(parent_otel_span), None
# Priority 2: HTTP traceparent header
if traceparent is not None:
@@ -929,6 +929,91 @@ class TestEventTypeLogging:
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
@pytest.mark.asyncio
async def test_log_guardrail_information_skips_auto_record_if_function_already_recorded(
self,
):
"""When a wrapped guardrail function records its own entry directly
(e.g. block_code_execution.apply_guardrail records a rich
``[detections...]`` payload), the decorator must NOT also append its
own ``"allow"``/raw-response entry otherwise every backend
(OTEL spans, Datadog, Langfuse, spend logs) double-records one
logical guardrail invocation."""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="block-code",
event_hook=GuardrailEventHooks.pre_call,
)
@log_guardrail_information
async def apply_guardrail(self, inputs, request_data, **kwargs):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=[{"action_taken": "block"}],
request_data=request_data,
guardrail_status="success",
event_type=GuardrailEventHooks.pre_call,
)
return inputs
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.apply_guardrail(
inputs={"texts": ["x"]}, request_data=request_data
)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1, (
f"Decorator must not double-record when the wrapped function "
f"already appended its own entry; got {len(logged_info)} entries"
)
assert logged_info[0]["guardrail_response"] == [{"action_taken": "block"}]
@pytest.mark.asyncio
async def test_log_guardrail_information_skips_auto_record_on_exception_if_function_already_recorded(
self,
):
"""Same as above on the failure path: if the wrapped function
appended an entry in its ``finally`` block before re-raising, the
decorator must just re-raise without auto-recording on top."""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="block-code",
event_hook=GuardrailEventHooks.pre_call,
)
@log_guardrail_information
async def apply_guardrail(self, inputs, request_data, **kwargs):
try:
raise ValueError("blocked")
finally:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=[{"action_taken": "block"}],
request_data=request_data,
guardrail_status="guardrail_intervened",
event_type=GuardrailEventHooks.pre_call,
)
guardrail = TestGuardrail()
request_data = {"metadata": {}}
with pytest.raises(ValueError, match="blocked"):
await guardrail.apply_guardrail(
inputs={"texts": ["x"]}, request_data=request_data
)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_status"] == "guardrail_intervened"
def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
self,
):
@@ -1086,9 +1171,12 @@ class TestCustomGuardrailSpendLogMatchRedaction:
][0]["match"]
== "[REDACTED]"
)
assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
"match"
] == "GG"
assert (
raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
"match"
]
== "GG"
)
def test_add_standard_logging_redacts_regex_field(self):
cg = CustomGuardrail(guardrail_name="test-rail")
@@ -1982,11 +1982,13 @@ class TestOpenTelemetryExternalSpan(unittest.TestCase):
- raw_gen_ai_request spans are children of litellm_request spans
- Correct hierarchy: external_parent litellm_request raw_gen_ai_request
"""
import copy
# Initialize OpenTelemetry
otel = OpenTelemetry(tracer_provider=self.tracer_provider)
# Load test data
kwargs, response_obj = self._create_test_kwargs_and_response()
kwargs1, response_obj = self._create_test_kwargs_and_response()
kwargs2 = copy.deepcopy(kwargs1)
# Create external parent span using our test TracerProvider
tracer = self.tracer_provider.get_tracer(__name__)
@@ -1999,7 +2001,7 @@ class TestOpenTelemetryExternalSpan(unittest.TestCase):
# First completion call
start_time = datetime.utcnow()
end_time = start_time + timedelta(seconds=1)
otel._handle_success(kwargs, response_obj, start_time, end_time)
otel._handle_success(kwargs1, response_obj, start_time, end_time)
# Verify parent span is still recording
self.assertTrue(
@@ -2010,7 +2012,7 @@ class TestOpenTelemetryExternalSpan(unittest.TestCase):
# Second completion call
start_time2 = end_time
end_time2 = start_time2 + timedelta(seconds=1)
otel._handle_success(kwargs, response_obj, start_time2, end_time2)
otel._handle_success(kwargs2, response_obj, start_time2, end_time2)
# Verify parent span is still recording
self.assertTrue(
@@ -3161,7 +3163,6 @@ class TestResponseIdFallback(unittest.TestCase):
mock_span.set_attribute.assert_any_call("litellm.call_id", call_id)
class TestOpenTelemetryResponsesAPI(unittest.TestCase):
"""
Tests for Responses API (/v1/responses) OTel span attributes.
@@ -3374,7 +3375,9 @@ class TestOpenTelemetryResponsesAPI(unittest.TestCase):
# No output messages should be set since the text is empty
raw = self._get_attr(mock_span, "gen_ai.output.messages")
self.assertIsNone(raw, "Empty output text should not produce gen_ai.output.messages")
self.assertIsNone(
raw, "Empty output text should not produce gen_ai.output.messages"
)
def test_choices_still_work(self):
"""Existing choices-based responses must still work (no regression)."""
@@ -3616,7 +3619,6 @@ class TestTransformResponsesAPIOutput(unittest.TestCase):
result = otel._transform_responses_api_output_to_otel(output)
self.assertEqual(result[0]["role"], "assistant")
def test_pydantic_like_objects_accepted(self):
"""Items with .get() but not isinstance(dict) should be accepted."""
@@ -3738,12 +3740,16 @@ class TestResponsesAPIToolCallSpanAttributes(unittest.TestCase):
],
}
otel.set_attributes(span=mock_span, kwargs=self._base_kwargs(), response_obj=response_obj)
otel.set_attributes(
span=mock_span, kwargs=self._base_kwargs(), response_obj=response_obj
)
# Verify per-tool-call attributes were set (same format as choices branch)
attr_names = [call[0][0] for call in mock_span.set_attribute.call_args_list]
tool_call_attrs = [a for a in attr_names if "function_call" in a]
self.assertTrue(len(tool_call_attrs) > 0, "Per-tool-call span attributes should be emitted")
self.assertTrue(
len(tool_call_attrs) > 0, "Per-tool-call span attributes should be emitted"
)
# Verify the name attribute specifically
mock_span.set_attribute.assert_any_call(
@@ -3778,7 +3784,9 @@ class TestResponsesAPIToolCallSpanAttributes(unittest.TestCase):
],
}
otel.set_attributes(span=mock_span, kwargs=self._base_kwargs(), response_obj=response_obj)
otel.set_attributes(
span=mock_span, kwargs=self._base_kwargs(), response_obj=response_obj
)
mock_span.set_attribute.assert_any_call(
"gen_ai.completion.0.function_call.name", "get_weather"
@@ -3786,3 +3794,488 @@ class TestResponsesAPIToolCallSpanAttributes(unittest.TestCase):
mock_span.set_attribute.assert_any_call(
"gen_ai.completion.1.function_call.name", "get_time"
)
class TestOpenTelemetryProxyParentSpanChildEmission(unittest.TestCase):
"""When metadata includes litellm_parent_otel_span (the proxy
span), the primary litellm_request span must still be created as a child
so the trace hierarchy is complete."""
def _build_kwargs(self, parent_span):
return {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"optional_params": {},
"litellm_params": {
"custom_llm_provider": "openai",
"metadata": {"litellm_parent_otel_span": parent_span},
},
"standard_logging_object": {
"id": "test-id",
"call_type": "completion",
"metadata": {},
"hidden_params": {},
},
}
def test_get_span_context_returns_none_parent_for_metadata_span(self):
"""_get_span_context Priority 1 must return (ctx, None) — never the
parent span object so callers always create litellm_request as a
child of ctx."""
tracer_provider = TracerProvider()
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
parent_span = otel.tracer.start_span("some_external_parent")
kwargs = self._build_kwargs(parent_span)
ctx, returned_parent = otel._get_span_context(kwargs)
self.assertIsNotNone(ctx, "ctx should carry the parent for child spans")
self.assertIsNone(
returned_parent,
"parent_span return slot must be None so callers create litellm_request",
)
parent_span.end()
def test_litellm_request_emitted_as_child_of_proxy_parent_span(self):
"""End-to-end: proxy span in metadata should yield exactly one
litellm_request span parented to it, with no extra root span."""
from litellm.integrations.opentelemetry import (
LITELLM_PROXY_REQUEST_SPAN_NAME,
LITELLM_REQUEST_SPAN_NAME,
)
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
proxy_span = otel.tracer.start_span(LITELLM_PROXY_REQUEST_SPAN_NAME)
kwargs = self._build_kwargs(proxy_span)
start = datetime.utcnow()
end = start + timedelta(seconds=1)
otel._handle_success(kwargs, response_obj=None, start_time=start, end_time=end)
spans = span_exporter.get_finished_spans()
litellm_spans = [s for s in spans if s.name == LITELLM_REQUEST_SPAN_NAME]
proxy_spans = [s for s in spans if s.name == LITELLM_PROXY_REQUEST_SPAN_NAME]
self.assertEqual(
len(litellm_spans), 1, "Exactly one litellm_request span must be emitted"
)
self.assertEqual(
len(proxy_spans), 1, "Proxy span should be closed exactly once"
)
litellm_span = litellm_spans[0]
self.assertIsNotNone(
litellm_span.parent, "litellm_request must have a parent (not root)"
)
self.assertEqual(
litellm_span.parent.span_id,
proxy_spans[0].context.span_id,
"litellm_request must be a child of the proxy span",
)
def test_end_proxy_span_from_kwargs_closes_recording_proxy_span(self):
from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
proxy_span = otel.tracer.start_span(LITELLM_PROXY_REQUEST_SPAN_NAME)
self.assertTrue(proxy_span.is_recording())
kwargs = {
"litellm_params": {
"metadata": {"litellm_parent_otel_span": proxy_span},
}
}
otel._end_proxy_span_from_kwargs(kwargs, end_time=datetime.utcnow())
self.assertFalse(
proxy_span.is_recording(), "Proxy span should be closed by helper"
)
def test_end_proxy_span_from_kwargs_does_not_close_external_span(self):
"""Spans not named LITELLM_PROXY_REQUEST_SPAN_NAME must not be closed —
they may belong to external owners (Langfuse SDK, user code, etc.)."""
tracer_provider = TracerProvider()
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
external = otel.tracer.start_span("external_caller_span")
kwargs = {
"litellm_params": {
"metadata": {"litellm_parent_otel_span": external},
}
}
otel._end_proxy_span_from_kwargs(kwargs, end_time=datetime.utcnow())
self.assertTrue(
external.is_recording(),
"External (non-proxy) parent span must not be closed by LiteLLM",
)
external.end()
class TestOpenTelemetryProxyLoggerFirstRegisteredWins(unittest.TestCase):
"""open_telemetry_logger ownership must not be silently
overwritten by later handlers. First-registered wins."""
def _install_fake_proxy_server(self):
"""Install a stub ``litellm.proxy.proxy_server`` so the test does
not depend on optional proxy dependencies (websockets, etc.).
Returns (fake_module, cleanup_fn)."""
import importlib
import types
proxy_pkg_name = "litellm.proxy"
proxy_server_name = "litellm.proxy.proxy_server"
previous_pkg = sys.modules.get(proxy_pkg_name)
previous_mod = sys.modules.get(proxy_server_name)
# Ensure litellm.proxy package object exists
if previous_pkg is None:
try:
pkg = importlib.import_module(proxy_pkg_name)
except Exception:
pkg = types.ModuleType(proxy_pkg_name)
sys.modules[proxy_pkg_name] = pkg
else:
pkg = previous_pkg
fake = types.ModuleType(proxy_server_name)
fake.open_telemetry_logger = None
sys.modules[proxy_server_name] = fake
setattr(pkg, "proxy_server", fake)
def cleanup():
if previous_mod is not None:
sys.modules[proxy_server_name] = previous_mod
setattr(pkg, "proxy_server", previous_mod)
else:
sys.modules.pop(proxy_server_name, None)
if hasattr(pkg, "proxy_server"):
try:
delattr(pkg, "proxy_server")
except AttributeError:
pass
if previous_pkg is None and proxy_pkg_name in sys.modules:
if sys.modules[proxy_pkg_name] is pkg:
# Leave it in place — removing it would break later imports
pass
return fake, cleanup
def test_first_registered_handler_keeps_ownership(self):
fake_proxy_server, cleanup = self._install_fake_proxy_server()
try:
first = OpenTelemetry()
self.assertIs(
fake_proxy_server.open_telemetry_logger,
first,
"First registered handler must own the proxy logger slot",
)
second = OpenTelemetry()
self.assertIs(
fake_proxy_server.open_telemetry_logger,
first,
"Second handler must NOT overwrite the first-registered logger",
)
self.assertIsNot(
fake_proxy_server.open_telemetry_logger,
second,
"Proxy logger must remain pointed at the first handler",
)
finally:
cleanup()
def test_assignment_happens_when_slot_is_unset(self):
fake_proxy_server, cleanup = self._install_fake_proxy_server()
try:
handler = OpenTelemetry()
self.assertIs(fake_proxy_server.open_telemetry_logger, handler)
finally:
cleanup()
def test_existing_non_none_logger_is_preserved(self):
"""If ``proxy_server.open_telemetry_logger`` is already set to any
non-None value, a new handler must not overwrite it even if the
existing value is not an OpenTelemetry instance."""
fake_proxy_server, cleanup = self._install_fake_proxy_server()
try:
sentinel = object()
fake_proxy_server.open_telemetry_logger = sentinel
OpenTelemetry()
self.assertIs(
fake_proxy_server.open_telemetry_logger,
sentinel,
"Existing non-None logger must not be overwritten",
)
finally:
cleanup()
class TestOpenTelemetrySpanDedupe(unittest.TestCase):
"""``_emit_once`` is a per-request, per-handler idempotency guard that
prevents duplicate span emission across two distinct dual-fire patterns:
1. Handler-level: streaming triggers both sync and async success/failure
callbacks for one request the second call would otherwise produce a
duplicate ``litellm_request`` span.
2. Payload-driven entry-level: ``_create_guardrail_span`` is invoked
from three lifecycle points (post-call hook, success, failure) and
re-reads a mutating list the same logical guardrail invocation
would otherwise be emitted up to three times.
"""
def _build_kwargs(self, *, exception: bool = False):
kwargs = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"optional_params": {},
"litellm_params": {
"custom_llm_provider": "openai",
"metadata": {},
},
"standard_logging_object": {
"id": "test-id",
"call_type": "completion",
"metadata": {},
"hidden_params": {},
},
}
if exception:
kwargs["exception"] = Exception("test error")
return kwargs
def test_emit_once_first_call_returns_true_then_false(self):
otel = OpenTelemetry()
kwargs = self._build_kwargs()
self.assertTrue(otel._emit_once(kwargs, "success"))
self.assertFalse(
otel._emit_once(kwargs, "success"),
"Repeat call for same handler+scope+kwargs must be deduped",
)
def test_emit_once_distinct_scopes_dont_collide(self):
"""Different scopes on the same handler+kwargs must each emit once."""
otel = OpenTelemetry()
kwargs = self._build_kwargs()
self.assertTrue(otel._emit_once(kwargs, "success"))
self.assertTrue(
otel._emit_once(kwargs, "failure"),
"Failure scope must be independent of success scope",
)
self.assertTrue(
otel._emit_once(kwargs, "guardrail", "block-code", 1.0, "pre_call"),
"Guardrail entry scope must be independent of success/failure scopes",
)
self.assertFalse(otel._emit_once(kwargs, "success"))
self.assertFalse(otel._emit_once(kwargs, "failure"))
self.assertFalse(
otel._emit_once(kwargs, "guardrail", "block-code", 1.0, "pre_call")
)
def test_emit_once_separate_handlers_each_emit(self):
"""Two distinct handler instances must each emit exactly once for the
same scope."""
otel_a = OpenTelemetry()
otel_b = OpenTelemetry()
kwargs = self._build_kwargs()
self.assertTrue(otel_a._emit_once(kwargs, "success"))
self.assertTrue(
otel_b._emit_once(kwargs, "success"),
"Different handler instance must not share the first handler's marker",
)
self.assertFalse(otel_a._emit_once(kwargs, "success"))
self.assertFalse(otel_b._emit_once(kwargs, "success"))
def test_emit_once_handles_missing_metadata(self):
otel = OpenTelemetry()
kwargs = {"litellm_params": {}}
self.assertTrue(otel._emit_once(kwargs, "success"))
self.assertFalse(otel._emit_once(kwargs, "success"))
def test_emit_once_handles_missing_litellm_params(self):
otel = OpenTelemetry()
kwargs = {}
self.assertTrue(otel._emit_once(kwargs, "success"))
self.assertFalse(otel._emit_once(kwargs, "success"))
def test_handle_success_emits_single_litellm_request_span_on_double_call(self):
"""Sync + async callback paths firing for the same kwargs must
result in exactly one litellm_request span."""
from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
kwargs = self._build_kwargs()
start = datetime.utcnow()
end = start + timedelta(seconds=1)
otel._handle_success(kwargs, response_obj=None, start_time=start, end_time=end)
otel._handle_success(kwargs, response_obj=None, start_time=start, end_time=end)
spans = span_exporter.get_finished_spans()
litellm_spans = [s for s in spans if s.name == LITELLM_REQUEST_SPAN_NAME]
self.assertEqual(
len(litellm_spans),
1,
f"Exactly one litellm_request span expected, got {len(litellm_spans)}",
)
def test_handle_success_dedupe_skip_still_closes_proxy_span(self):
"""When the success path is short-circuited as a duplicate, the
proxy span must still be closed so traces don't leak."""
from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
proxy_span = otel.tracer.start_span(LITELLM_PROXY_REQUEST_SPAN_NAME)
kwargs = self._build_kwargs()
kwargs["litellm_params"]["metadata"]["litellm_parent_otel_span"] = proxy_span
otel._emit_once(kwargs, "success") # pre-mark to force dedupe-skip branch
self.assertTrue(proxy_span.is_recording())
start = datetime.utcnow()
end = start + timedelta(seconds=1)
otel._handle_success(kwargs, response_obj=None, start_time=start, end_time=end)
self.assertFalse(
proxy_span.is_recording(),
"Dedupe-skip path must still close the proxy span via _end_proxy_span_from_kwargs",
)
def test_handle_failure_emits_single_error_span_on_double_call(self):
"""Sync + async failure callback paths firing for the same kwargs
must result in exactly one ERROR litellm_request span."""
from opentelemetry.trace import StatusCode
from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
kwargs = self._build_kwargs(exception=True)
start = datetime.utcnow()
end = start + timedelta(seconds=1)
otel._handle_failure(kwargs, response_obj=None, start_time=start, end_time=end)
otel._handle_failure(kwargs, response_obj=None, start_time=start, end_time=end)
spans = span_exporter.get_finished_spans()
litellm_spans = [s for s in spans if s.name == LITELLM_REQUEST_SPAN_NAME]
self.assertEqual(
len(litellm_spans),
1,
f"Exactly one litellm_request ERROR span expected, got {len(litellm_spans)}",
)
self.assertEqual(litellm_spans[0].status.status_code, StatusCode.ERROR)
def test_create_guardrail_span_dedupes_across_lifecycle_entrypoints(self):
"""``_create_guardrail_span`` is called from post-call-success hook,
``_handle_success``, and ``_handle_failure``. A single guardrail
invocation (identified by ``(name, start_time, mode)``) must produce
exactly one span per handler even when the underlying entry is
mutated between calls (e.g. proxy enriches ``guardrail_response``)."""
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
kwargs = self._build_kwargs()
guardrail_entry = {
"guardrail_name": "block-code",
"guardrail_mode": "pre_call",
"guardrail_response": "allow",
"start_time": 1.0,
"end_time": 2.0,
}
kwargs["standard_logging_object"]["guardrail_information"] = [guardrail_entry]
otel._create_guardrail_span(kwargs=kwargs, context=None)
# Mutate the entry between calls — proxy enriches the response.
guardrail_entry["guardrail_response"] = [
{"type": "code_block", "action_taken": "block"}
]
guardrail_entry["end_time"] = 3.0
otel._create_guardrail_span(kwargs=kwargs, context=None)
otel._create_guardrail_span(kwargs=kwargs, context=None)
guardrail_spans = [
s for s in span_exporter.get_finished_spans() if s.name == "guardrail"
]
self.assertEqual(
len(guardrail_spans),
1,
f"Exactly one guardrail span expected per logical invocation, got {len(guardrail_spans)}",
)
def test_create_guardrail_span_emits_distinct_entries(self):
"""Two real guardrail invocations (different ``start_time``) must
each emit a span entry-level dedupe must not collapse them."""
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
kwargs = self._build_kwargs()
kwargs["standard_logging_object"]["guardrail_information"] = [
{
"guardrail_name": "block-code",
"guardrail_mode": "pre_call",
"guardrail_response": "allow",
"start_time": 1.0,
"end_time": 2.0,
},
{
"guardrail_name": "block-code",
"guardrail_mode": "post_call",
"guardrail_response": "allow",
"start_time": 5.0,
"end_time": 6.0,
},
]
otel._create_guardrail_span(kwargs=kwargs, context=None)
otel._create_guardrail_span(kwargs=kwargs, context=None)
guardrail_spans = [
s for s in span_exporter.get_finished_spans() if s.name == "guardrail"
]
self.assertEqual(
len(guardrail_spans),
2,
f"Two distinct guardrail invocations expected, got {len(guardrail_spans)}",
)