From 1c4e4d4a60e586b97ec4ab82c358291eaf8f07f2 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 12 May 2026 18:32:05 -0700 Subject: [PATCH] Fix 3 OpenTelemetry tracing bugs in proxy integration (#27757) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Co-authored-by: Claude Opus 4.7 (1M context) --- litellm/integrations/custom_guardrail.py | 29 + litellm/integrations/opentelemetry.py | 136 ++++- .../integrations/test_custom_guardrail.py | 94 +++- .../integrations/test_opentelemetry.py | 513 +++++++++++++++++- 4 files changed, 756 insertions(+), 16 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index a03aef481e..c937ad0a7b 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -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, diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index abce3c3e1c..48d7a07a56 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -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: diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index d09c4ac2c3..a881044dc1 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -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") diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 898f42b45b..3af4a21a60 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -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)}", + )