From d18fcb09d6f9f073fb100329775a0eb818d677ef Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 10:10:12 -0700 Subject: [PATCH] fix(otel): detach post-response service spans by request phase, name redis spans by operation (#43237) * fix(otel): detach post-response service spans by request phase, name redis spans by operation Service spans logged from the post-response phase (success callbacks, the response-cache write) now root their own trace linked to the request span even while the server span is still recording, instead of only when they happen to end after it. Redis service spans are named `redis `; the litellm call chain that issued them moves to the `litellm.service.caller` attribute via a typed `ServiceLoggerPayload.caller` field. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): keep the service caller on failure and legacy spans, test the production phase dispatch sites Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): mark anthropic messages stream cache write as post-response phase The /v1/messages streaming cache writer awaits async_add_cache inline instead of going through create_cache_write_task, so its redis span stayed parented under the request trace. Wrap the write in post_response_phase so it detaches like the chat completions write. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): write the Messages stream cache in a background task after handoff Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/_internal_context.py | 16 ++ litellm/_service_logger.py | 8 + litellm/caching/caching_handler.py | 4 +- litellm/caching/redis_cache.py | 123 ++++++---- litellm/integrations/opentelemetry.py | 4 +- litellm/integrations/otel/README.md | 30 ++- litellm/integrations/otel/logger.py | 8 +- litellm/integrations/otel/mappers/genai.py | 1 + litellm/integrations/otel/mappers/legacy.py | 2 + litellm/integrations/otel/model/payloads.py | 2 + litellm/integrations/otel/model/semconv.py | 1 + litellm/integrations/otel/plumbing/context.py | 22 +- litellm/litellm_core_utils/litellm_logging.py | 15 +- .../messages/response_cache.py | 27 ++- litellm/types/services.py | 1 + tests/unit/caching/test_caching_handler.py | 32 +++ .../otel/test_otel_v2_components.py | 6 +- .../integrations/otel/test_otel_v2_logger.py | 220 +++++++++++++++++- .../test_litellm_logging.py | 57 +++++ .../messages/test_response_cache.py | 38 +++ 20 files changed, 522 insertions(+), 95 deletions(-) diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index 8132008731f..389add8ed0f 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -21,6 +21,22 @@ is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", defau # moment they can land on either side of a window boundary and disagree with each other. _billing_time: Final[ContextVar[datetime | None]] = ContextVar("billing_time", default=None) +_post_response: Final[ContextVar[bool]] = ContextVar("post_response", default=False) + + +@contextmanager +def post_response_phase() -> Generator[None]: + """Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns.""" + token: Final = _post_response.set(True) + try: + yield + finally: + _post_response.reset(token) + + +def in_post_response_phase() -> bool: + return _post_response.get() + @contextmanager def pinned_billing_time(moment: datetime) -> Generator[None]: diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 0ccac4b5291..1a5f46e9261 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -159,6 +159,7 @@ class ServiceLogging(CustomLogger): parent_otel_span: Span | None = None, start_time: datetime | float | None = None, end_time: float | datetime | None = None, + caller: str | None = None, ): """ Handles both sync and async monitoring by checking for existing event loop. @@ -172,6 +173,7 @@ class ServiceLogging(CustomLogger): service=service, duration=duration, call_type=call_type, + caller=caller, parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, @@ -187,6 +189,7 @@ class ServiceLogging(CustomLogger): parent_otel_span: Span | None = None, start_time: datetime | float | None = None, end_time: float | datetime | None = None, + caller: str | None = None, ): """ Handles both sync and async monitoring by checking for existing event loop. @@ -200,6 +203,7 @@ class ServiceLogging(CustomLogger): duration=duration, error=error, call_type=call_type, + caller=caller, parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, @@ -215,6 +219,7 @@ class ServiceLogging(CustomLogger): start_time: datetime | float | None = None, end_time: datetime | float | None = None, event_metadata: dict | None = None, + caller: str | None = None, ): """ - For counting if the redis, postgres call is successful @@ -228,6 +233,7 @@ class ServiceLogging(CustomLogger): service=service, duration=duration, call_type=call_type, + caller=caller, event_metadata=event_metadata, ) @@ -313,6 +319,7 @@ class ServiceLogging(CustomLogger): start_time: datetime | float | None = None, end_time: float | datetime | None = None, event_metadata: dict | None = None, + caller: str | None = None, ): """ - For counting if the redis, postgres call is unsuccessful @@ -332,6 +339,7 @@ class ServiceLogging(CustomLogger): service=service, duration=duration, call_type=call_type, + caller=caller, event_metadata=event_metadata, ) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 0887b8bb897..cfc9edd7158 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -24,6 +24,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar from pydantic import BaseModel, ConfigDict, ValidationError import litellm +from litellm._internal_context import post_response_phase from litellm._logging import print_verbose, verbose_logger from litellm.caching import InMemoryCache from litellm.caching.caching import S3Cache @@ -158,7 +159,8 @@ async def _complete_cache_write_despite_cancellation(write_factory: Callable[[], def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "asyncio.Task[None]": - task: Final = asyncio.create_task(_complete_cache_write_despite_cancellation(write_factory)) + with post_response_phase(): + task: Final = asyncio.create_task(_complete_cache_write_despite_cancellation(write_factory)) _PENDING_CACHE_WRITES.add(task) task.add_done_callback(_PENDING_CACHE_WRITES.discard) return task diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 0b56c28f9b1..29e390b1d9a 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -839,7 +839,8 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"set_cache <- {_get_call_stack_info()}", + call_type="set_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, ) @@ -860,7 +861,8 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"increment_cache <- {_get_call_stack_info()}", + call_type="increment_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, ) @@ -874,7 +876,8 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"increment_cache_ttl <- {_get_call_stack_info()}", + call_type="increment_cache_ttl", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, ) @@ -887,7 +890,8 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"increment_cache_expire <- {_get_call_stack_info()}", + call_type="increment_cache_expire", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, ) @@ -963,7 +967,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_scan_iter <- {_get_call_stack_info()}", + call_type="async_scan_iter", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, ) @@ -979,7 +984,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_scan_iter <- {_get_call_stack_info()}", + call_type="async_scan_iter", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, ) @@ -1100,7 +1106,8 @@ class RedisCache(BaseCache): start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), - call_type=f"async_set_cache <- {_get_call_stack_info()}", + call_type="async_set_cache", + caller=_get_call_stack_info(), ) ) log_redis_failure( @@ -1129,7 +1136,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_set_cache <- {_get_call_stack_info()}", + call_type="async_set_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1145,7 +1153,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_set_cache <- {_get_call_stack_info()}", + call_type="async_set_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1213,7 +1222,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_set_cache_pipeline <- {_get_call_stack_info()}", + call_type="async_set_cache_pipeline", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1229,7 +1239,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_set_cache_pipeline <- {_get_call_stack_info()}", + call_type="async_set_cache_pipeline", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1263,7 +1274,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=time.time() - start_time, - call_type=f"async_set_cache_pipeline_with_ttls <- {_get_call_stack_info()}", + call_type="async_set_cache_pipeline_with_ttls", + caller=_get_call_stack_info(), start_time=start_time, end_time=time.time(), ) @@ -1274,7 +1286,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=time.time() - start_time, error=e, - call_type=f"async_set_cache_pipeline_with_ttls <- {_get_call_stack_info()}", + call_type="async_set_cache_pipeline_with_ttls", + caller=_get_call_stack_info(), start_time=start_time, end_time=time.time(), ) @@ -1322,7 +1335,8 @@ class RedisCache(BaseCache): start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), - call_type=f"async_set_cache_sadd <- {_get_call_stack_info()}", + call_type="async_set_cache_sadd", + caller=_get_call_stack_info(), ) ) # NON blocking - notify users Redis is throwing an exception @@ -1342,7 +1356,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_set_cache_sadd <- {_get_call_stack_info()}", + call_type="async_set_cache_sadd", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1356,7 +1371,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_set_cache_sadd <- {_get_call_stack_info()}", + call_type="async_set_cache_sadd", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1427,7 +1443,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_increment <- {_get_call_stack_info()}", + call_type="async_increment", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -1443,7 +1460,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_increment <- {_get_call_stack_info()}", + call_type="async_increment", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -1531,7 +1549,8 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"get_cache <- {_get_call_stack_info()}", + call_type="get_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -1590,7 +1609,8 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"batch_get_cache <- {_get_call_stack_info()}", + call_type="batch_get_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -1614,7 +1634,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=failed_at - start_time, error=e, - call_type=f"batch_get_cache <- {_get_call_stack_info()}", + call_type="batch_get_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=failed_at, parent_otel_span=parent_otel_span, @@ -1643,7 +1664,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_get_cache <- {_get_call_stack_info()}", + call_type="async_get_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -1659,7 +1681,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_get_cache <- {_get_call_stack_info()}", + call_type="async_get_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -1704,7 +1727,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_batch_get_cache <- {_get_call_stack_info()}", + call_type="async_batch_get_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -1732,7 +1756,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_batch_get_cache <- {_get_call_stack_info()}", + call_type="async_batch_get_cache", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -1757,7 +1782,8 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"sync_ping <- {_get_call_stack_info()}", + call_type="sync_ping", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, ) @@ -1771,7 +1797,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"sync_ping <- {_get_call_stack_info()}", + call_type="sync_ping", + caller=_get_call_stack_info(), ) verbose_logger.error("LiteLLM Redis Cache PING: - Got exception from REDIS : %s", e) raise e @@ -1789,7 +1816,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_ping <- {_get_call_stack_info()}", + call_type="async_ping", + caller=_get_call_stack_info(), ) ) return response @@ -1803,7 +1831,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_ping <- {_get_call_stack_info()}", + call_type="async_ping", + caller=_get_call_stack_info(), ) ) verbose_logger.error("LiteLLM Redis Cache PING: - Got exception from REDIS : %s", e) @@ -1955,7 +1984,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_increment_pipeline <- {_get_call_stack_info()}", + call_type="async_increment_pipeline", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1971,7 +2001,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_increment_pipeline <- {_get_call_stack_info()}", + call_type="async_increment_pipeline", + caller=_get_call_stack_info(), start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -2049,7 +2080,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_rpush <- {_get_call_stack_info()}", + call_type="async_rpush", + caller=_get_call_stack_info(), ) ) return response @@ -2063,7 +2095,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_rpush <- {_get_call_stack_info()}", + call_type="async_rpush", + caller=_get_call_stack_info(), ) ) log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH: - Got exception from REDIS", e) @@ -2096,7 +2129,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=time.time() - start_time, - call_type=f"async_rpush_and_trim <- {_get_call_stack_info()}", + call_type="async_rpush_and_trim", + caller=_get_call_stack_info(), ) ) return int(results[0]) @@ -2106,7 +2140,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=time.time() - start_time, error=e, - call_type=f"async_rpush_and_trim <- {_get_call_stack_info()}", + call_type="async_rpush_and_trim", + caller=_get_call_stack_info(), ) ) log_redis_failure( @@ -2163,7 +2198,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}", + call_type="async_rpush_pipeline", + caller=_get_call_stack_info(), ) ) return results @@ -2176,7 +2212,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}", + call_type="async_rpush_pipeline", + caller=_get_call_stack_info(), ) ) log_redis_failure( @@ -2230,7 +2267,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_lpop <- {_get_call_stack_info()}", + call_type="async_lpop", + caller=_get_call_stack_info(), ) ) @@ -2256,7 +2294,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_lpop <- {_get_call_stack_info()}", + call_type="async_lpop", + caller=_get_call_stack_info(), ) ) log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache LPOP: - Got exception from REDIS", e) @@ -2354,7 +2393,8 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}", + call_type="async_lpop_pipeline", + caller=_get_call_stack_info(), ) ) return results @@ -2367,7 +2407,8 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}", + call_type="async_lpop_pipeline", + caller=_get_call_stack_info(), ) ) log_redis_failure( diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index c1531f4e4ae..948e3113337 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -25,7 +25,7 @@ from litellm.integrations.otel.mappers.utils import drop_none from litellm.integrations.otel.model.baggage import promoted_metadata from litellm.integrations.otel.model.db_endpoint import db_span_attributes from litellm.integrations.otel.model.metadata import flatten_metadata -from litellm.integrations.otel.model.semconv import Metric +from litellm.integrations.otel.model.semconv import LiteLLM, Metric from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -784,6 +784,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) for key, value in attributes.items(): self.safe_set_attribute(span=span, key=key, value=value) + if payload.caller is not None: + self.safe_set_attribute(span=span, key=LiteLLM.SERVICE_CALLER, value=payload.caller) return span async def async_service_success_hook( diff --git a/litellm/integrations/otel/README.md b/litellm/integrations/otel/README.md index d8dfabe23d6..1b97e159105 100644 --- a/litellm/integrations/otel/README.md +++ b/litellm/integrations/otel/README.md @@ -60,7 +60,10 @@ traceable units of work: instead (see below). Spans are named `"{service} {call_type}"` (e.g. `"redis set"`) so repeated calls -to one service stay distinguishable. Like every other span they parent to the +to one service stay distinguishable. `call_type` is the operation only; the +litellm call chain that issued it (`async_set_cache <- async_add_cache`) travels +as `ServiceLoggerPayload.caller` and lands on the `litellm.service.caller` +attribute, so one operation is one span name. Like every other span they parent to the **ambient** context, falling back to the threaded `litellm_parent_otel_span` only when ambient has no live span; a background job with neither starts its own root trace. @@ -69,16 +72,21 @@ trace. and the spend-counter increment all run after the response is on the wire, so they add nothing to the request's latency. Parenting them under the (already ended) server span stretched the request trace past the request itself, which is what a -viewer shows as trace duration. `context.resolve_service_span_context` compares -the call's end time with the resolved parent's end time: a call that finished -after its parent ended starts a **new root trace** carrying a **span link** back -to the request span (the `FollowsFrom` relationship of OpenTracing; the default -`:link` propagation style of the OTel Ruby ActiveJob and Sidekiq -instrumentations). Identity Baggage still rides along, so the detached span keeps -its team / key / user attributes. Only an SDK span that has really ended detaches: -a sampled-out or remote `NonRecordingSpan` is never recording but is still the -right parent. A call that ended before the server span did stays a child even when -its `asyncio.create_task`-dispatched hook runs after the response. +viewer shows as trace duration. `context.resolve_service_span_context` detaches +a call in two cases: it was logged from the post-response phase +(`litellm._internal_context.post_response_phase`, entered by the success +handlers and by the response-cache write task, inherited by every task spawned +inside), or it finished after the resolved parent ended. Either way it starts a +**new root trace** carrying a **span link** back to the request span (the +`FollowsFrom` relationship of OpenTracing; the default `:link` propagation style +of the OTel Ruby ActiveJob and Sidekiq instrumentations). The phase check matters +for streaming: the stream-finished callbacks run before the ASGI server span +closes, so by end time alone the cache write would look like request latency. +Identity Baggage still rides along, so the detached span keeps its team / key / +user attributes. Only an SDK span detaches: a sampled-out or remote +`NonRecordingSpan` is never recording but is still the right parent. A call that +ended before the server span did stays a child even when its +`asyncio.create_task`-dispatched hook runs after the response. Caller-supplied `event_metadata` is **sanitized** before it reaches a span (primitives only, no live objects, no secrets/headers, bounded) — see diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 0466e00a959..e21711c2708 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -3,6 +3,7 @@ from collections import OrderedDict from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager +from dataclasses import replace from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast @@ -661,12 +662,7 @@ class OpenTelemetryV2(CustomLogger): if error_override is None and start_time is None and end_time is None and parent_otel_span is None: return None if error_override is not None and data.error is None: - data = ServiceSpanData( - service_name=data.service_name, - call_type=data.call_type, - error=SpanError(message=error_override), - event_metadata=data.event_metadata, - ) + data = replace(data, error=SpanError(message=error_override)) # Parent like every other span: ambient context first (so identity Baggage # rides along and the call nests under whatever request phase is active — # e.g. a DB lookup under the live ``auth`` span), falling back to the diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 1a9b897ca28..e37da8908e4 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -148,6 +148,7 @@ class GenAIMapper: _SERVICE_ATTRS: dict[str, Callable[[ServiceSpanData], AttrValue | None]] = { LiteLLM.SERVICE_NAME: lambda d: d.service_name, LiteLLM.SERVICE_CALL_TYPE: lambda d: d.call_type, + LiteLLM.SERVICE_CALLER: lambda d: d.caller, } def __init__(self, tool_attr_budget: int = MAX_TOOL_DEFINITION_ATTRS_PER_SPAN) -> None: diff --git a/litellm/integrations/otel/mappers/legacy.py b/litellm/integrations/otel/mappers/legacy.py index d25c25cd127..df15fe86a94 100644 --- a/litellm/integrations/otel/mappers/legacy.py +++ b/litellm/integrations/otel/mappers/legacy.py @@ -37,6 +37,7 @@ _LEGACY_PRESENCE_PENALTY: Final = "llm.presence_penalty" _LEGACY_STOP_SEQUENCES: Final = "llm.chat.stop_sequences" _LEGACY_SERVICE: Final = "service" _LEGACY_CALL_TYPE: Final = "call_type" +_LEGACY_CALLER: Final = "caller" _LEGACY_ERROR: Final = Error.MESSAGE_LEGACY @@ -66,6 +67,7 @@ class LegacyMapper: _SERVICE_ATTRS: dict[str, Callable[[ServiceSpanData], AttrValue | None]] = { _LEGACY_SERVICE: lambda d: d.service_name, _LEGACY_CALL_TYPE: lambda d: d.call_type, + _LEGACY_CALLER: lambda d: d.caller, _LEGACY_ERROR: lambda d: d.error.message if d.error is not None and d.error.message else None, } diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index ea4ded90480..7e47abfb20d 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -309,6 +309,7 @@ class GuardrailSpanData: class ServiceSpanData: service_name: str call_type: str | None = None + caller: str | None = None error: SpanError | None = None # Caller-supplied attributes to stamp on the service span, passed through # from ``async_service_*_hook(event_metadata=...)``. The mapper owns how @@ -330,6 +331,7 @@ class ServiceSpanData: return cls( service_name=payload.service.value, call_type=payload.call_type, + caller=payload.caller, error=SpanError(message=payload.error) if payload.error else None, event_metadata=sanitize_event_metadata(event_metadata), ) diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index f552ba37655..19b319009e8 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -326,6 +326,7 @@ class LiteLLM: GUARDRAIL_COST_IN_SPEND: Final = "litellm.guardrail.cost_in_spend" SERVICE_NAME: Final = "litellm.service.name" SERVICE_CALL_TYPE: Final = "litellm.service.call_type" + SERVICE_CALLER: Final = "litellm.service.caller" PREPROCESSING_MS: Final = "litellm.preprocessing.duration_ms" # The logical name of the MCP server a tool call was routed to. There is no # semconv key for an MCP server's *name* (the convention uses ``server.address`` diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 9de5c1ac1cb..f5f221cf278 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -21,6 +21,7 @@ from opentelemetry.trace.propagation.tracecontext import ( TraceContextTextMapPropagator, ) +from litellm._internal_context import in_post_response_phase from litellm.integrations.otel.model.semconv import HTTP if TYPE_CHECKING: @@ -231,21 +232,28 @@ def resolve_service_span_context( ) -> tuple[Context, tuple[Link, ...]]: """Parent context + links for a service/DB span that ended at ``end_time_ns``. - A call that finished after its parent ended (post-response spend tracking) - starts its own root trace with a span link back to the parent instead of - stretching the parent's trace. Baggage stays on the returned context. + Work the caller did not wait for starts its own root trace with a span link + back to the parent instead of stretching the parent's trace: anything logged + from the post-response phase (success callbacks, the response-cache write, + see :func:`litellm._internal_context.post_response_phase`), whether or not + the server span has closed yet, and anything that finished after its parent + ended. Baggage stays on the returned context. """ ctx: Final = resolve_parent_context(threaded) parent: Final = get_current_span(ctx) - if not _ended_before(parent, end_time_ns): + if not _is_post_response(parent, end_time_ns): return ctx, () return set_span_in_context(INVALID_SPAN, ctx), (Link(parent.get_span_context()),) -def _ended_before(span: Span, end_time_ns: int | None) -> bool: - if not isinstance(span, ReadableSpan) or span.end_time is None: +def _is_post_response(parent: Span, end_time_ns: int | None) -> bool: + if not isinstance(parent, ReadableSpan): return False - return end_time_ns is None or end_time_ns > span.end_time + if in_post_response_phase(): + return True + if parent.end_time is None: + return False + return end_time_ns is None or end_time_ns > parent.end_time def resolve_request_span_context() -> Context: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e955c0157c6..f6211869913 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -21,6 +21,7 @@ from pydantic import BaseModel, JsonValue import litellm from litellm import _custom_logger_compatible_callbacks_literal +from litellm._internal_context import post_response_phase from litellm._logging import ( _is_debugging_on, _redact_string, @@ -2739,9 +2740,10 @@ class Logging(LiteLLMLoggingBaseClass): """Restores trace_id/session_id contextvars once this attempt's own success logging (including any nested calls its callbacks trigger) is fully done.""" try: - return self._success_handler_body( - result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs - ) + with post_response_phase(): + return self._success_handler_body( + result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs + ) finally: self._restore_correlation_context() @@ -3177,9 +3179,10 @@ class Logging(LiteLLMLoggingBaseClass): """Restores trace_id/session_id contextvars once this attempt's own success logging (including any nested calls its callbacks trigger) is fully done.""" try: - return await self._async_success_handler_body( - result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs - ) + with post_response_phase(): + return await self._async_success_handler_body( + result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs + ) finally: self._restore_correlation_context() diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py index dc2d4408c20..b60458f8401 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final import litellm from litellm._logging import verbose_logger +from litellm.caching.caching_handler import create_cache_write_task from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( AnthropicMessagesStreamingResponse, BaseAnthropicMessagesStreamingIterator, @@ -57,7 +58,7 @@ class AnthropicMessagesStreamCacheWriter: try: chunk: Final = await self.stream.__anext__() except StopAsyncIteration: - await self._persist() + self._persist() raise self.collected_chunks.append(chunk.encode("utf-8") if isinstance(chunk, str) else chunk) return chunk @@ -65,8 +66,9 @@ class AnthropicMessagesStreamCacheWriter: async def aclose(self) -> None: await aclose_if_supported(self.stream) - async def _persist(self) -> None: - if self.persisted or litellm.cache is None: + def _persist(self) -> None: + cache: Final = litellm.cache + if self.persisted or cache is None: return collected_stream: Final = b"".join(self.collected_chunks) if not _is_message_stop_chunk(collected_stream) or _is_provider_error_chunk(collected_stream): @@ -88,14 +90,19 @@ class AnthropicMessagesStreamCacheWriter: try: events: Final = _split_sse_events(collected_stream.decode("utf-8")) - cached_payload: Final = {CACHED_STREAM_EVENTS_KEY: events} - await litellm.cache.async_add_cache( - cached_payload, - dynamic_cache_object=self.caching_handler.dual_cache, - **request_kwargs, - ) - except Exception as e: # noqa: BLE001 # a cache write must never surface as a client-visible stream error + except UnicodeDecodeError as e: verbose_logger.exception("Anthropic Messages stream cache write failed: %s", e) + return + cached_payload: Final = {CACHED_STREAM_EVENTS_KEY: events} + dual_cache: Final = self.caching_handler.dual_cache + + async def _write() -> None: + try: + await cache.async_add_cache(cached_payload, dynamic_cache_object=dual_cache, **request_kwargs) + except Exception as e: # noqa: BLE001 # a cache write must never surface as a client-visible stream error + verbose_logger.exception("Anthropic Messages stream cache write failed: %s", e) + + create_cache_write_task(_write) class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterator): diff --git a/litellm/types/services.py b/litellm/types/services.py index c558f6fb9d2..b8c4265b6be 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -100,6 +100,7 @@ class ServiceLoggerPayload(BaseModel): service: ServiceTypes = Field(description="who is this for? - postgres/redis") duration: float = Field(description="How long did the request take?") call_type: str = Field(description="The call of the service, being made") + caller: str | None = Field(None, description="The litellm call chain that made the service call, innermost first") event_metadata: dict | None = Field(description="The metadata logged during service success/failure") def to_json(self, **kwargs): diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 425d657312a..6cf8e901cd7 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -43,6 +43,7 @@ import json import httpx import respx from fastapi.testclient import TestClient +from litellm._internal_context import in_post_response_phase from litellm.caching.caching_handler import _PENDING_CACHE_WRITES @@ -2073,6 +2074,37 @@ def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatc assert len(writes) == 1 +def test_async_cache_write_runs_in_the_post_response_phase_without_leaking_it(monkeypatch): + """The response-cache write happens after the response is handed to the caller, so the + service spans it logs must detach from the request trace even while the server span is + still open. The marker must stay inside the write task and not leak into the request.""" + import litellm + + phases = [] + + class _PhaseRecordingCache: + supported_call_types = ["acompletion"] + cache = None + + async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs): + phases.append(in_post_response_phase()) + + async def acompletion(**kwargs): + return None + + handler = LLMCachingHandler(original_function=acompletion, request_kwargs={}, start_time=datetime.now()) + monkeypatch.setattr(litellm, "cache", _PhaseRecordingCache()) + + async def _request(): + await handler.async_set_cache(result=litellm.ModelResponse(), original_function=acompletion, kwargs={}) + leaked = in_post_response_phase() + await asyncio.gather(*_PENDING_CACHE_WRITES) + return leaked + + assert asyncio.run(_request()) is False, "the phase must not leak into the request task" + assert phases == [True], "async_add_cache must observe the post-response phase" + + @pytest.mark.asyncio async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monkeypatch): """The spend log for a cache hit must reuse the key the lookup already computed instead of hashing again.""" diff --git a/tests/unit/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py index fd10210c5ba..fb7be0dda14 100644 --- a/tests/unit/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -144,16 +144,19 @@ def test_service_span_data_from_payload(): class _Payload: service = _Service() call_type = "async_set_cache" + caller = "async_set_cache <- async_add_cache" error = None data = ServiceSpanData.from_payload(_Payload()) assert data.service_name == "redis" assert data.call_type == "async_set_cache" + assert data.caller == "async_set_cache <- async_add_cache" assert data.error is None class _FailPayload: service = _Service() call_type = "async_set_cache" + caller = None error = "boom" failed = ServiceSpanData.from_payload(_FailPayload()) @@ -445,10 +448,11 @@ def test_legacy_mapper_all_request_params(): def test_legacy_mapper_covers_service_with_v1_bare_keys(): """Service spans dual-emit V1's bare ``service``/``call_type``/``error`` keys.""" attrs = LegacyMapper().map( - ServiceSpanData("redis", call_type="set", event_metadata={"k": "v"}), + ServiceSpanData("redis", call_type="set", caller="set <- add", event_metadata={"k": "v"}), ) assert attrs["service"] == "redis" assert attrs["call_type"] == "set" + assert attrs["caller"] == "set <- add" assert attrs["k"] == "v" # event_metadata is stamped bare (V1 behavior) diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index d478c670e58..62bf75bd083 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -9,8 +9,8 @@ hooks, proxy SERVER span lifecycle (start + setters), parent-context resolution import asyncio import contextlib import os -from unittest.mock import patch from datetime import datetime, timedelta, timezone +from unittest.mock import patch import pytest @@ -23,20 +23,13 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E4 from opentelemetry.trace import SpanKind # noqa: E402 from opentelemetry.trace.status import StatusCode # noqa: E402 +from litellm._internal_context import in_post_response_phase, post_response_phase # noqa: E402 from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY # noqa: E402 from litellm.integrations.otel import ( # noqa: E402 GenAI, LiteLLM, OpenTelemetryV2Config, ) -from litellm.integrations.otel.plumbing import providers # noqa: E402 -from litellm.integrations.otel.plumbing.context import ( # noqa: E402 - reset_mcp_message_trace_carrier, - reset_mcp_message_transport_span, - set_mcp_message_trace_carrier, - set_mcp_message_transport_span, - set_request_root_span, -) from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402 from litellm.integrations.otel.model.config import ExporterSpec # noqa: E402 from litellm.integrations.otel.model.spans import ( # noqa: E402 @@ -44,6 +37,14 @@ from litellm.integrations.otel.model.spans import ( # noqa: E402 SpanRole, ) from litellm.integrations.otel.model.utils import to_ns, to_seconds # noqa: E402 +from litellm.integrations.otel.plumbing import providers # noqa: E402 +from litellm.integrations.otel.plumbing.context import ( # noqa: E402 + reset_mcp_message_trace_carrier, + reset_mcp_message_transport_span, + set_mcp_message_trace_carrier, + set_mcp_message_transport_span, + set_request_root_span, +) # --------------------------------------------------------------------------- # # Fixtures @@ -1772,9 +1773,10 @@ class _Service: class _ServicePayload: - def __init__(self, service="redis", call_type="set", error=None): + def __init__(self, service="redis", call_type="set", error=None, caller=None): self.service = _Service(service) self.call_type = call_type + self.caller = caller self.error = error @@ -1785,6 +1787,54 @@ def _service_parent(logger): ) +async def _redis_get_through_service_logger(logger): + """Drive a real ``RedisCache.async_get_cache`` (client doubled at the edge) through the real + ``ServiceLogging`` into ``logger``, the way the proxy's cache reads reach OTel.""" + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm._service_logger import ServiceLogging + from litellm.caching.redis_cache import RedisCache + + async_client = MagicMock() + async_client.get = AsyncMock(return_value=None) + async_client.ping = AsyncMock(return_value=True) + with ( + patch("litellm._redis.get_redis_client", return_value=MagicMock()), + patch("litellm._redis.get_redis_connection_pool", return_value=MagicMock()), + patch("litellm._redis.get_redis_async_client", return_value=async_client), + patch.object(litellm, "service_callback", [logger]), + patch.object( + litellm, + "in_memory_llm_clients_cache", + MagicMock(get_cache=MagicMock(return_value=None)), + ), + ): + cache = RedisCache( + host="127.0.0.1", port=6379, service_logger_obj=ServiceLogging() + ) + await cache.async_get_cache("otel-naming-key") + await asyncio.gather( + *(t for t in asyncio.all_tasks() if t is not asyncio.current_task()) + ) + + +def test_redis_service_span_is_named_by_operation_and_keeps_the_caller_chain_as_an_attribute(): + """``redis async_get_cache``, not ``redis async_get_cache <- caller <- caller``: the stack + walk that used to be spliced into the span name rides on ``litellm.service.caller`` instead, + so one operation is one span name and ``db.operation.name`` is the bare operation.""" + logger, exporter = _logger() + asyncio.run(_redis_get_through_service_logger(logger)) + (span,) = [s for s in exporter.get_finished_spans() if s.name.startswith("redis")] + assert span.name == "redis async_get_cache" + assert span.attributes[LiteLLM.SERVICE_CALL_TYPE] == "async_get_cache" + assert span.attributes["db.operation.name"] == "async_get_cache" + callers = span.attributes[LiteLLM.SERVICE_CALLER].split(" <- ") + assert callers[0] == "_redis_get_through_service_logger" and len(callers) == 2, ( + callers + ) + + def test_async_service_success_hook_emits_service_span(): logger, exporter = _logger() parent = _service_parent(logger) @@ -1853,7 +1903,7 @@ def test_async_service_failure_hook_marks_error_status(): try: asyncio.run( logger.async_service_failure_hook( - payload=_ServicePayload("postgres", "query"), + payload=_ServicePayload("postgres", "query", caller="query <- get_user_object"), error="boom", parent_otel_span=parent, ) @@ -1868,6 +1918,7 @@ def test_async_service_failure_hook_marks_error_status(): # Without an explicit error_type from the payload, V2 stamps the fallback. assert span.attributes["error.type"] == "error" assert span.attributes[LiteLLM.SERVICE_NAME] == "postgres" + assert span.attributes[LiteLLM.SERVICE_CALLER] == "query <- get_user_object" def test_async_service_failure_hook_preserves_payload_error_over_override(): @@ -2086,6 +2137,153 @@ def test_service_call_under_a_remote_parent_is_never_detached(): assert list(span.links) == [] +def _service_hook_from_post_response_task( + logger, payload, *, parent, ambient, end_time +): + """Log ``payload`` the way the proxy's post-response tail does: the hook runs on a + task spawned from inside ``post_response_phase`` while the server span is still open.""" + + async def _dispatch(): + with post_response_phase(): + task = asyncio.create_task( + logger.async_service_success_hook( + payload=payload, + parent_otel_span=parent, + start_time=end_time - 0.4, + end_time=end_time, + ) + ) + assert not in_post_response_phase(), ( + "the phase must not leak into the request task" + ) + await task + + if ambient is None: + asyncio.run(_dispatch()) + return + with trace.use_span(ambient, end_on_exit=False): + asyncio.run(_dispatch()) + + +@pytest.mark.parametrize("parent_source", ["ambient", "threaded"]) +def test_service_call_from_the_post_response_phase_detaches_before_the_server_span_ends( + parent_source, +): + """The streaming tail: the response-cache write and the success callbacks run + after the client has the whole response but before the ASGI server span closes, + so the call ends before its parent does. Timing alone would keep it a child; + being dispatched from the post-response phase is what detaches it, with a link.""" + logger, exporter = _logger() + server = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + assert server.is_recording() + try: + _service_hook_from_post_response_task( + logger, + _ServicePayload("redis", "async_set_cache"), + parent=server if parent_source == "threaded" else None, + ambient=server if parent_source == "ambient" else None, + end_time=_REQUEST_END - 0.1, + ) + finally: + server.end(end_time=to_ns(_REQUEST_END)) + span = {s.name: s for s in exporter.get_finished_spans()}["redis async_set_cache"] + request_ctx = server.get_span_context() + assert span.end_time < server.end_time + assert span.parent is None + assert span.context.trace_id != request_ctx.trace_id + assert [(link.context.trace_id, link.context.span_id) for link in span.links] == [ + (request_ctx.trace_id, request_ctx.span_id) + ] + + +def test_service_call_from_the_post_response_phase_under_a_remote_parent_is_never_detached(): + from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags + + logger, exporter = _logger() + remote = NonRecordingSpan( + SpanContext( + trace_id=0xABC, + span_id=0x123, + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + ) + ) + _service_hook_from_post_response_task( + logger, + _ServicePayload("redis", "get"), + parent=remote, + ambient=None, + end_time=_REQUEST_END, + ) + span = {s.name: s for s in exporter.get_finished_spans()}["redis get"] + assert span.parent.span_id == 0x123 + assert span.context.trace_id == 0xABC + assert list(span.links) == [] + + +def test_redis_write_from_a_success_callback_detaches_while_the_server_span_is_still_open(): + """The production dispatch path: ``Logging.async_success_handler`` runs the + success callbacks, one of which writes to redis and logs the service span + through the OTel logger. With the server span still recording (the streaming + tail), the redis span must still root its own trace linked to the request.""" + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import ModelResponse + + logger, exporter = _logger() + + class _RedisWritingCallback(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await logger.async_service_success_hook( + payload=_ServicePayload("redis", "async_increment", caller="async_increment_cache <- async_log_success_event"), + parent_otel_span=None, + start_time=_REQUEST_END - 0.5, + end_time=_REQUEST_END - 0.1, + ) + + async def _request(): + logging_obj = Logging( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(timezone.utc), + litellm_call_id="call-1", + function_id="fn-1", + dynamic_async_success_callbacks=[_RedisWritingCallback()], + ) + logging_obj.update_environment_variables( + model="gpt-4o", + user="u", + optional_params={}, + litellm_params={"metadata": {}, "acompletion": True}, + custom_llm_provider="openai", + ) + await logging_obj.async_success_handler( + result=ModelResponse(model="gpt-4o", choices=[{"message": {"role": "assistant", "content": "ok"}}]), + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + ) + + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + try: + with trace.use_span(server, end_on_exit=False): + asyncio.run(_request()) + finally: + server.end(end_time=to_ns(_REQUEST_END)) + span = {s.name: s for s in exporter.get_finished_spans()}["redis async_increment"] + request_ctx = server.get_span_context() + assert span.end_time < server.end_time + assert span.parent is None + assert span.context.trace_id != request_ctx.trace_id + assert [(link.context.trace_id, link.context.span_id) for link in span.links] == [ + (request_ctx.trace_id, request_ctx.span_id) + ] + assert span.attributes[LiteLLM.SERVICE_CALLER] == "async_increment_cache <- async_log_success_event" + + # --------------------------------------------------------------------------- # # Proxy SERVER span lifecycle # --------------------------------------------------------------------------- # diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index cb1e281e356..f211505d06d 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -19,6 +19,7 @@ from openai import AsyncOpenAI from openai._legacy_response import HttpxBinaryResponseContent import litellm +from litellm._internal_context import in_post_response_phase from litellm._logging import session_id_var, trace_id_var from litellm.constants import REDACTED_BY_LITELLM, SENTRY_PII_DENYLIST from litellm.cost_calculator import ocr_batch_cost @@ -1970,6 +1971,62 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call dummy_logger.log_stream_event.assert_not_called() +class _PhaseRecordingLogger(CustomLogger): + """Records whether each success callback ran inside the post-response phase.""" + + def __init__(self) -> None: + super().__init__() + self.phases: list[bool] = [] + + def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.phases.append(in_post_response_phase()) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.phases.append(in_post_response_phase()) + + +def _success_response() -> ModelResponse: + return ModelResponse( + id="resp-123", + model="gpt-4o-mini", + choices=[{"message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop", "index": 0}], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + + +def test_success_handler_runs_sync_callbacks_in_the_post_response_phase(logging_obj): + """Service spans logged by success callbacks must detach from the request trace even + while the server span is still open, so the callbacks run inside the phase marker.""" + logging_obj.stream = False + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + recorder = _PhaseRecordingLogger() + + with patch.object(logging_obj, "get_combined_callback_list", return_value=[recorder]): + logging_obj.success_handler(result=_success_response()) + + assert recorder.phases == [True], "log_success_event must observe the post-response phase" + assert in_post_response_phase() is False, "the phase must end with the handler" + + +@pytest.mark.asyncio +async def test_async_success_handler_runs_async_callbacks_in_the_post_response_phase(logging_obj): + logging_obj.stream = False + logging_obj.model_call_details["litellm_params"] = {"acompletion": True} + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + recorder = _PhaseRecordingLogger() + + with patch.object(logging_obj, "get_combined_callback_list", return_value=[recorder]): + await logging_obj.async_success_handler( + result=_success_response(), + start_time=datetime.datetime.now(datetime.timezone.utc), + end_time=datetime.datetime.now(datetime.timezone.utc), + ) + + assert recorder.phases == [True], "async_log_success_event must observe the post-response phase" + assert in_post_response_phase() is False, "the phase must not leak into the request task" + + def test_is_sync_litellm_request(): assert LitellmLogging._is_sync_litellm_request({}) is True assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False diff --git a/tests/unit/llms/anthropic/experimental_pass_through/messages/test_response_cache.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_response_cache.py index 22d14614108..aecc84cfcaa 100644 --- a/tests/unit/llms/anthropic/experimental_pass_through/messages/test_response_cache.py +++ b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_response_cache.py @@ -7,6 +7,7 @@ import pytest import datetime import litellm +from litellm._internal_context import in_post_response_phase from litellm.caching.caching import Cache, LiteLLMCacheType from litellm.caching.caching_handler import LLMCachingHandler from litellm.llms.anthropic.experimental_pass_through.messages import handler @@ -130,6 +131,7 @@ async def test_streaming_request_is_replayed_from_cache(local_cache, request_kwa monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + await asyncio.sleep(0) second_stream = await litellm.anthropic_messages(**request_kwargs, stream=True) second = await _collect(second_stream) @@ -181,6 +183,7 @@ async def test_multibyte_utf8_split_across_chunks_streams_and_caches(local_cache monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + await asyncio.sleep(0) second = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) assert len(fake_handler.calls) == 1 @@ -198,6 +201,7 @@ async def test_message_stop_split_across_chunks_still_caches(local_cache, reques monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + await asyncio.sleep(0) second = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) assert len(fake_handler.calls) == 1 @@ -278,6 +282,40 @@ class _HeldBackStream: raise StopAsyncIteration +@pytest.mark.asyncio +async def test_stream_cache_write_runs_in_post_response_phase(request_kwargs, monkeypatch): + """Every event, message_stop included, is already with the client when the stream write + runs, so it must not hold the stream open and the redis span it logs must detach from the + request trace like the chat completions write does. The marker must not leak into the consumer.""" + phases: list[bool] = [] + write_started = asyncio.Event() + release_write = asyncio.Event() + + class _PhaseRecordingCache: + supported_call_types = ["anthropic_messages"] + cache = None + + async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs): + phases.append(in_post_response_phase()) + write_started.set() + await release_write.wait() + + monkeypatch.setattr(litellm, "cache", _PhaseRecordingCache()) + caching_handler = LLMCachingHandler( + original_function=handler.anthropic_messages, + request_kwargs=dict(request_kwargs), + start_time=datetime.datetime.now(), + ) + writer = AnthropicMessagesStreamCacheWriter(stream=_byte_stream(STREAM_EVENTS), caching_handler=caching_handler) + + collected = await asyncio.wait_for(_collect(writer), timeout=1) + assert collected == STREAM_EVENTS, "the stream must close without waiting for the write" + assert in_post_response_phase() is False, "the phase must not leak into the stream consumer" + await asyncio.wait_for(write_started.wait(), timeout=1) + release_write.set() + assert phases == [True], "async_add_cache must observe the post-response phase" + + def test_cache_writer_forwards_has_buffered_provider_output(request_kwargs): caching_handler = LLMCachingHandler( original_function=handler.anthropic_messages,