From 93e71f67b93c2fd3df868e0e65bf3a4bab8a5b9c Mon Sep 17 00:00:00 2001
From: Varun Sripad <varunsripad@Varuns-MacBook-Air.local>
Date: Wed, 28 Jan 2026 16:29:43 -0600
Subject: [PATCH] fix(langfuse): set LANGFUSE_TRACING_ENVIRONMENT on all span
 types

Fixes #19926: Environment not populated on all Langfuse spans

When setting LANGFUSE_TRACING_ENVIRONMENT, the environment attribute was
only being set on generation spans, not on other span types (service,
guardrail, management endpoint spans) in the trace.

Changes:
- Added set_langfuse_environment_on_span() helper in LangfuseOtelLogger
- Call this helper from all span creation points in OpenTelemetry:
  - async_service_success_hook
  - async_service_failure_hook
  - guardrail span creation (_maybe_log_guardrails)
  - async_management_endpoint_success_hook
  - async_management_endpoint_failure_hook
---
 .../integrations/langfuse/langfuse_otel.py    | 21 +++++++++++
 litellm/integrations/opentelemetry.py         | 37 +++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py
index 08493a0e8ec..5321961aa84 100644
--- a/litellm/integrations/langfuse/langfuse_otel.py
+++ b/litellm/integrations/langfuse/langfuse_otel.py
@@ -40,6 +40,27 @@ class LangfuseOtelLogger(OpenTelemetry):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    @staticmethod
+    def set_langfuse_environment_on_span(span: Span):
+        """
+        Sets the LANGFUSE_TRACING_ENVIRONMENT attribute on any span.
+        
+        This method should be called for all span types (service, guardrail, 
+        management endpoint, etc.) to ensure the environment is populated
+        on all spans in a trace, not just generation spans.
+        
+        Fixes: https://github.com/BerriAI/litellm/issues/19926
+        """
+        from litellm.integrations.arize._utils import safe_set_attribute
+        
+        langfuse_environment = os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
+        if langfuse_environment:
+            safe_set_attribute(
+                span,
+                LangfuseSpanAttributes.LANGFUSE_ENVIRONMENT.value,
+                langfuse_environment,
+            )
+
     @staticmethod
     def set_langfuse_otel_attributes(span: Span, kwargs, response_obj):
         """
diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py
index 997dd044a65..9b3a9e2ec11 100644
--- a/litellm/integrations/opentelemetry.py
+++ b/litellm/integrations/opentelemetry.py
@@ -435,6 +435,13 @@ class OpenTelemetry(CustomLogger):
                 key="service",
                 value=payload.service.value,
             )
+            
+            # Set Langfuse environment on service spans (fixes #19926)
+            if self.callback_name == "langfuse_otel":
+                from litellm.integrations.langfuse.langfuse_otel import (
+                    LangfuseOtelLogger,
+                )
+                LangfuseOtelLogger.set_langfuse_environment_on_span(service_logging_span)
 
             if event_metadata:
                 for key, value in event_metadata.items():
@@ -514,6 +521,13 @@ class OpenTelemetry(CustomLogger):
                         value=value,
                     )
 
+            # Set Langfuse environment on service failure spans (fixes #19926)
+            if self.callback_name == "langfuse_otel":
+                from litellm.integrations.langfuse.langfuse_otel import (
+                    LangfuseOtelLogger,
+                )
+                LangfuseOtelLogger.set_langfuse_environment_on_span(service_logging_span)
+
             service_logging_span.set_status(Status(StatusCode.ERROR))
             service_logging_span.end(end_time=_end_time_ns)
 
@@ -1156,8 +1170,16 @@ class OpenTelemetry(CustomLogger):
                 value=guardrail_information.get("guardrail_response"),
             )
 
+            # Set Langfuse environment on guardrail spans (fixes #19926)
+            if self.callback_name == "langfuse_otel":
+                from litellm.integrations.langfuse.langfuse_otel import (
+                    LangfuseOtelLogger,
+                )
+                LangfuseOtelLogger.set_langfuse_environment_on_span(guardrail_span)
+
             guardrail_span.end(end_time=self._to_ns(end_time_datetime))
 
+
     def _handle_failure(self, kwargs, response_obj, start_time, end_time):
         from opentelemetry.trace import Status, StatusCode
 
@@ -2234,6 +2256,13 @@ class OpenTelemetry(CustomLogger):
                         value=value,
                     )
 
+            # Set Langfuse environment on management endpoint spans (fixes #19926)
+            if self.callback_name == "langfuse_otel":
+                from litellm.integrations.langfuse.langfuse_otel import (
+                    LangfuseOtelLogger,
+                )
+                LangfuseOtelLogger.set_langfuse_environment_on_span(management_endpoint_span)
+
             management_endpoint_span.set_status(Status(StatusCode.OK))
             management_endpoint_span.end(end_time=_end_time_ns)
 
@@ -2284,6 +2313,14 @@ class OpenTelemetry(CustomLogger):
                 key="exception",
                 value=str(_exception),
             )
+
+            # Set Langfuse environment on management endpoint failure spans (fixes #19926)
+            if self.callback_name == "langfuse_otel":
+                from litellm.integrations.langfuse.langfuse_otel import (
+                    LangfuseOtelLogger,
+                )
+                LangfuseOtelLogger.set_langfuse_environment_on_span(management_endpoint_span)
+
             management_endpoint_span.set_status(Status(StatusCode.ERROR))
             management_endpoint_span.end(end_time=_end_time_ns)