diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a4de1be1842..ca7877134db 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -66,6 +66,7 @@ _BASE64_INLINE_PATTERN: Final = re.compile( class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class # Class variables or attributes server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset() + is_internal_accounting_callback: ClassVar[bool] = False enforces_request_content: bool = False """ diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 141162bdfa5..8165f1eeeff 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -10,6 +10,7 @@ worker thread, off any event loop — and caches it for the process lifetime. """ from collections.abc import Sequence +from dataclasses import dataclass, field from typing import Final import httpx @@ -80,6 +81,16 @@ def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, ) + from requests.auth import AuthBase + + @dataclass(frozen=True, slots=True) + class _AgentOpsBearerAuth(AuthBase): + token: str = field(repr=False) + + def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest: + authenticated: Final = request.copy() + authenticated.headers["Authorization"] = f"Bearer {self.token}" + return authenticated class _LazyAuthAgentOpsExporter(OTLPSpanExporter): """OTLP/HTTP exporter that mints the AgentOps JWT on its first export. @@ -103,8 +114,8 @@ def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: return try: token: Final = _fetch_agentops_jwt(self._agentops_api_key).get("token") - if token: - self._agentops_session.headers["Authorization"] = f"Bearer {token}" + if isinstance(token, str) and token: + self._agentops_session.auth = _AgentOpsBearerAuth(token) except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9e0461e951d..0ea96b33987 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,6 +19,7 @@ from functools import lru_cache from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast +from anyio import fail_after from httpx import Response, Timeout from pydantic import BaseModel, JsonValue, TypeAdapter @@ -114,9 +115,11 @@ from litellm.litellm_core_utils.logging_utils import ( truncate_base64_in_messages, truncate_base64_in_messages_async, ) +from litellm.litellm_core_utils.logging_worker import optional_callback_budget from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.ptu_pricing import is_spilled_over_ptu_request from litellm.litellm_core_utils.redact_messages import ( + perform_redaction, # pyright: ignore[reportUnknownVariableType] # legacy sanitizer has untyped response variants redact_message_input_output_from_custom_logger, redact_message_input_output_from_logging, redacted_model_call_details_for_custom_logger, @@ -3704,6 +3707,106 @@ class Logging(LiteLLMLoggingBaseClass): e, ) + async def _run_async_success_hooks( + self, callbacks: Sequence[str | Callable[..., object] | CustomLogger], result: object + ) -> object: + for callback in callbacks: + try: + if isinstance(callback, CustomGuardrail): + from litellm.types.guardrails import GuardrailEventHooks + + if ( + callback.should_run_guardrail( + data=self.model_call_details, + event_type=GuardrailEventHooks.logging_only, + ) + is not True + ): + continue + + self.shadow_eval_request_snapshot = None + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + elif isinstance(callback, CustomLogger): + result = redact_message_input_output_from_custom_logger( + result=result, litellm_logging_obj=self, custom_logger=callback + ) + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s", + traceback.format_exc(), + ) + self._handle_callback_failure(callback=callback) + + return result + + async def _prepare_accounting_success_hooks( + self, + callbacks: Sequence[str | Callable[..., object] | CustomLogger], + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> object: + from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE + + try: + with fail_after(optional_callback_budget(LOGGING_WORKER_MAX_TIME_PER_COROUTINE, fraction=0.5)): + return await self._run_async_success_hooks(callbacks, result) + except TimeoutError: + from litellm.constants import X_LITELLM_DISABLE_CALLBACKS + from litellm.proxy.spend_tracking.spend_event import ( + SpendEventBuildError, + build_spend_event_payload, + spend_event_callback_args, + ) + + self.shadow_eval_request_snapshot = None + self.model_call_details = redacted_model_call_details_for_custom_logger( + self.model_call_details, CustomLogger(message_logging=False) + ) + perform_redaction(self.model_call_details, None) + assembled_stream: Final = "async_complete_streaming_response" in self.model_call_details + accounting_response: Final[object] = ( + self.model_call_details["async_complete_streaming_response"] if assembled_stream else result + ) + event: Final = build_spend_event_payload( + self.model_call_details, accounting_response, start_time, end_time, store_bodies=False + ) + if isinstance(event, SpendEventBuildError): + verbose_logger.error("Cannot prepare body-free spend accounting after a logging hook timeout") + return event + args: Final = spend_event_callback_args(event) + mapping_adapter: Final = TypeAdapter(Mapping[str, object]) + empty_headers: Final = MappingProxyType[str, object]({}) + params: Final = mapping_adapter.validate_python(self.model_call_details["litellm_params"]) + request_data: Final = params.get("proxy_server_request") + request: Final = ( + mapping_adapter.validate_python(request_data) if request_data is not None else empty_headers + ) + header_data: Final = request.get("headers") + headers: Final = mapping_adapter.validate_python(header_data) if header_data is not None else empty_headers + self.model_call_details = { + **args.kwargs, + "litellm_params": { + **args.kwargs["litellm_params"], + "no-log": params.get("no-log", False), + "proxy_server_request": { + "headers": {X_LITELLM_DISABLE_CALLBACKS: headers.get(X_LITELLM_DISABLE_CALLBACKS)} + }, + }, + **({"async_complete_streaming_response": args.response_obj} if assembled_stream else {}), + } + verbose_logger.warning("Logging hooks timed out, recording spend without request or response bodies") + return args.response_obj + async def async_success_handler( self, result: object = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) @@ -3921,7 +4024,7 @@ class Logging(LiteLLMLoggingBaseClass): self._get_completed_success_callbacks(asynchronous=True) if not self.stream else callbacks ) - result = _REDACTED_RESULT.validate_python( + redacted_result: Final = _REDACTED_RESULT.validate_python( redact_message_input_output_from_logging( model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=result, @@ -3931,45 +4034,35 @@ class Logging(LiteLLMLoggingBaseClass): ## LOGGING HOOK ## hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks))) - for callback in hook_callbacks: - try: - if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks + accounting_callbacks: Final = tuple( + { + type(callback): callback + for callback in delivery_callbacks + if isinstance(callback, CustomLogger) and callback.is_internal_accounting_callback + }.values() + ) + result = ( + await self._prepare_accounting_success_hooks(hook_callbacks, redacted_result, start_time, end_time) + if accounting_callbacks + else await self._run_async_success_hooks(hook_callbacks, redacted_result) + ) + if accounting_callbacks: + from litellm.proxy.spend_tracking.spend_event import SpendEventBuildError - if ( - callback.should_run_guardrail( - data=self.model_call_details, - event_type=GuardrailEventHooks.logging_only, - ) - is not True - ): - continue - - self.shadow_eval_request_snapshot = None - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - elif isinstance(callback, CustomLogger): - result = redact_message_input_output_from_custom_logger( - result=result, litellm_logging_obj=self, custom_logger=callback - ) - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) - verbose_logger.error( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s", - traceback.format_exc(), - ) - self._handle_callback_failure(callback=callback) + if isinstance(result, SpendEventBuildError): + return + ordered_callbacks: Final = ( + *accounting_callbacks, + *( + callback + for callback in delivery_callbacks + if not (isinstance(callback, CustomLogger) and callback.is_internal_accounting_callback) + ), + ) self.mark_logging_complete(event_type="async_success") - for callback in delivery_callbacks: + for callback in ordered_callbacks: # check if callback can run for this request litellm_params = self.model_call_details.get("litellm_params", {}) should_run = self.should_run_callback( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c11b59a809a..8b222b73d64 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -5,7 +5,7 @@ from contextlib import contextmanager from contextvars import ContextVar from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast +from typing import TYPE_CHECKING, ClassVar, Final, Protocol, TypeAlias, cast from pydantic import TypeAdapter, ValidationError @@ -337,6 +337,8 @@ def _attempt_spend_row_kwargs( class _ProxyDBLogger(CustomLogger): + is_internal_accounting_callback: ClassVar[bool] = True + def __init__( self, spend_event_producer: SpendEventProducer | None = None, diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py index 031e84896b2..e28954e4f1d 100644 --- a/litellm/proxy/spend_tracking/spend_event.py +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -342,10 +342,10 @@ def _function_tool(name: str) -> _FunctionTool: return tool -def build_spend_event( +def build_spend_event_payload( raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool -) -> bytes | SpendEventBuildError: - """Validate the success callback's kwargs and serialize the event once, as a single JSON line.""" +) -> SpendEvent | SpendEventBuildError: + """Validate and project the identity, usage, cost, and optional bodies of a completed call.""" try: kwargs: Final = _SUCCESS_KWARGS.validate_python(raw_kwargs) except ValidationError as error: @@ -379,6 +379,16 @@ def build_spend_event( realtime_tool_calls=kwargs.get("realtime_tool_calls"), request_tool_names=_request_tool_names(kwargs), ) + return event + + +def build_spend_event( + raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool +) -> bytes | SpendEventBuildError: + """Validate the success callback's kwargs and serialize the event once, as a single JSON line.""" + event: Final = build_spend_event_payload(raw_kwargs, response_obj, start_time, end_time, store_bodies) + if isinstance(event, SpendEventBuildError): + return event return event.model_dump_json(fallback=_json_fallback).encode() + b"\n" diff --git a/tests/unit/integrations/otel/test_otel_v2_presets.py b/tests/unit/integrations/otel/test_otel_v2_presets.py index 570c27a24e1..c5b436259a3 100644 --- a/tests/unit/integrations/otel/test_otel_v2_presets.py +++ b/tests/unit/integrations/otel/test_otel_v2_presets.py @@ -3,18 +3,18 @@ event loop: the preset does no network I/O, and a custom exporter mints the JWT lazily on its first export (in the BatchSpanProcessor worker thread).""" import json +from typing import Final import httpx import pytest -import requests_mock import respx +from requests_mock import Mocker from opentelemetry.sdk.trace.export import SpanExportResult from litellm.integrations.otel.plumbing import providers from litellm.integrations.otel.model.config import ExporterSpec from litellm.integrations.otel.presets import agentops as agentops_mod from litellm.integrations.otel.presets.agentops import ( - _AGENTOPS_AUTH_ENDPOINT, _AGENTOPS_ENDPOINT, _AGENTOPS_EXPORTER_KIND, _build_agentops_exporter, @@ -81,44 +81,61 @@ def test_dynamic_cred_presets_tag_exporter_with_matching_owner(monkeypatch): ) -def _agentops_exporter(): - return _build_agentops_exporter( - ExporterSpec( - kind=_AGENTOPS_EXPORTER_KIND, - endpoint=_AGENTOPS_ENDPOINT, - options={"api_key": "ak"}, - ) +@pytest.mark.parametrize("global_headers", (None, "Authorization=Basic%20shared")) +def test_agentops_exporter_mints_jwt_lazily( + requests_mock: Mocker, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + global_headers: str | None, +) -> None: + monkeypatch.delenv("OTEL_EXPORTER_OTLP_HEADERS", raising=False) + monkeypatch.delenv("OTEL_EXPORTER_OTLP_TRACES_HEADERS", raising=False) + if global_headers is not None: + monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", global_headers) + auth: Final = respx_mock.post(agentops_mod._AGENTOPS_AUTH_ENDPOINT).respond(200, json={"token": "jwt-xyz"}) + collector: Final = requests_mock.post(_AGENTOPS_ENDPOINT, status_code=200) + exporter: Final = _build_agentops_exporter( + ExporterSpec(kind=_AGENTOPS_EXPORTER_KIND, endpoint=_AGENTOPS_ENDPOINT, options={"api_key": "ak"}) ) + try: + assert auth.call_count == 0 + assert exporter.export([]) is SpanExportResult.SUCCESS + assert exporter.export([]) is SpanExportResult.SUCCESS + assert auth.call_count == 1 + assert [json.loads(call.request.content) for call in auth.calls] == [{"api_key": "ak"}] + assert [request.headers.get("Authorization") for request in collector.request_history] == [ + "Bearer jwt-xyz", + "Bearer jwt-xyz", + ] + finally: + exporter.shutdown() -def test_agentops_exporter_mints_the_jwt_once_on_first_export_and_sends_it_with_every_batch(): - pytest.importorskip("opentelemetry.exporter.otlp.proto.http.trace_exporter") - with respx.mock() as auth, requests_mock.Mocker() as collector: - token_route = auth.post(_AGENTOPS_AUTH_ENDPOINT).respond(200, json={"token": "jwt-xyz"}) - collector.post(_AGENTOPS_ENDPOINT, status_code=200) - exporter = _agentops_exporter() - assert token_route.call_count == 0, "building the exporter must not fetch the JWT" - - results = [exporter.export(()), exporter.export(())] - - assert results == [SpanExportResult.SUCCESS, SpanExportResult.SUCCESS] - assert [json.loads(call.request.content) for call in token_route.calls] == [{"api_key": "ak"}] - assert [request.headers.get("Authorization") for request in collector.request_history] == [ - "Bearer jwt-xyz", - "Bearer jwt-xyz", - ] - - -def test_agentops_exporter_still_exports_without_auth_when_the_jwt_fetch_fails(): - pytest.importorskip("opentelemetry.exporter.otlp.proto.http.trace_exporter") - with respx.mock() as auth, requests_mock.Mocker() as collector: - auth.post(_AGENTOPS_AUTH_ENDPOINT).mock(side_effect=httpx.ConnectError("auth down")) - collector.post(_AGENTOPS_ENDPOINT, status_code=200) - - result = _agentops_exporter().export(()) - - assert result is SpanExportResult.SUCCESS - assert [request.headers.get("Authorization") for request in collector.request_history] == [None] +@pytest.mark.parametrize(("status", "payload"), ((500, {}), (200, {}), (200, {"token": 123}), (None, {}))) +def test_agentops_exporter_tolerates_fetch_failure( + requests_mock: Mocker, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + status: int | None, + payload: dict[str, object], +) -> None: + monkeypatch.delenv("OTEL_EXPORTER_OTLP_HEADERS", raising=False) + monkeypatch.delenv("OTEL_EXPORTER_OTLP_TRACES_HEADERS", raising=False) + route: Final = respx_mock.post(agentops_mod._AGENTOPS_AUTH_ENDPOINT) + auth: Final = ( + route.mock(side_effect=httpx.ConnectError("auth down")) if status is None else route.respond(status, json=payload) + ) + collector: Final = requests_mock.post(_AGENTOPS_ENDPOINT, status_code=200) + exporter: Final = _build_agentops_exporter( + ExporterSpec(kind=_AGENTOPS_EXPORTER_KIND, endpoint=_AGENTOPS_ENDPOINT, options={"api_key": "ak"}) + ) + try: + assert exporter.export([]) is SpanExportResult.SUCCESS + assert exporter.export([]) is SpanExportResult.SUCCESS + assert auth.call_count == 1 + assert [request.headers.get("Authorization") for request in collector.request_history] == [None, None] + finally: + exporter.shutdown() def test_fetch_jwt_uses_owned_client_not_shared_pool(monkeypatch): diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 561ac9b2b33..f837d3bf551 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import contextvars import copy import datetime import importlib.abc @@ -50,7 +51,7 @@ from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, set_callbacks, ) -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.cache_control_check import PROXY_CacheControlCheck @@ -8254,12 +8255,14 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur releasing.async_log_success_event.assert_awaited_once() -def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging: +def _async_request_logging_obj( + call_id: str, callbacks: list[CustomLogger], stream: bool = False, call_type: str = "acompletion" +) -> LitellmLogging: logging_obj: Final = LitellmLogging( model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=stream, - call_type="acompletion", + call_type=call_type, start_time=datetime.datetime.now(), litellm_call_id=call_id, function_id=call_id, @@ -13611,6 +13614,298 @@ def test_legacy_s3_logs_a_streamed_anthropic_messages_call_handed_off_from_the_a assert s3_writer.log_event.call_args.kwargs["response_obj"] is complete_response +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("response_kind", ["chat", "responses"]) +@pytest.mark.parametrize("accounting_first", [False, True]) +@pytest.mark.parametrize("duplicate_accounting", [False, True]) +@pytest.mark.parametrize("blocked_phase", ["hook", "delivery", "neither", "redacted"]) +async def test_slow_optional_logging_cannot_cancel_accounting( + monkeypatch: pytest.MonkeyPatch, + stream: bool, + response_kind: str, + accounting_first: bool, + duplicate_accounting: bool, + blocked_phase: str, +) -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + from litellm.proxy.proxy_server import general_settings + from litellm.proxy.spend_tracking.spend_event import SpendEvent, decode_spend_event + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress + + recorded: Final[asyncio.Queue[tuple[bytes, str]]] = asyncio.Queue() + hook_context: Final = contextvars.ContextVar("accounting_hook_context", default="outside") + delivered: Final = asyncio.Event() + deliveries: Final[asyncio.Queue[str]] = asyncio.Queue() + + async def store_event(line: bytes) -> None: + recorded.put_nowait((line, hook_context.get())) + + producer: Final = SpendEventProducer( + address=UnixAddress("/unused"), + on_unavailable="fallback", + buffer_size=1, + connect_timeout=1, + fallback=store_event, + ) + await producer.close(drain_timeout=1) + monkeypatch.setitem(general_settings, "store_prompts_in_spend_logs", True) + cancelled: Final = asyncio.Event() + never: Final = asyncio.Event() + + class OptionalLogger(CustomLogger): + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + hook_context.set("inside") + if blocked_phase == "hook": + try: + await never.wait() + finally: + cancelled.set() + return kwargs, result + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + delivered.set() + deliveries.put_nowait(json.dumps(kwargs, default=str) + str(response_obj)) + if blocked_phase == "delivery": + try: + await never.wait() + finally: + cancelled.set() + + optional: Final = OptionalLogger(message_logging=blocked_phase != "redacted") + accounting: Final = _ProxyDBLogger(producer) + accounting_callbacks: Final = [accounting, _ProxyDBLogger(producer)] if duplicate_accounting else [accounting] + callbacks: Final = [*accounting_callbacks, optional] if accounting_first else [optional, *accounting_callbacks] + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj: Final = _async_request_logging_obj( + "accounting-timeout", + callbacks, + stream=stream, + call_type="aresponses" if response_kind == "responses" else "acompletion", + ) + response: Final = ( + ModelResponse( + choices=[{"message": {"role": "assistant", "content": "private-marker"}}], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + if response_kind == "chat" + else ResponsesAPIResponse( + id="resp-private-instructions", + created_at=1, + model="gpt-4o-mini", + status="completed", + instructions="private-instructions-canary" if blocked_phase != "redacted" else None, + output=[ + { + "type": "message", + "id": "msg-private", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "private-marker", "annotations": []}], + } + ], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15), + ) + ) + original_response: Final = response.model_dump(exclude={"usage"}) + logging_obj.update_environment_variables( + litellm_params={ + "acompletion": True, + "metadata": {}, + "proxy_server_request": { + "body": { + "instructions": "private-request-instructions", + "system": "private-system-canary", + "tools": [{"description": "private-tool-canary"}], + } + }, + } + if blocked_phase == "hook" + else {"acompletion": True, "metadata": {}}, + optional_params={}, + ) + logging_result: Final = ( + ResponseCompletedEvent(type="response.completed", sequence_number=1, response=response) + if stream and isinstance(response, ResponsesAPIResponse) + else response + ) + worker: Final = LoggingWorker(timeout=0.5) + worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=logging_result)) + try: + await worker.flush() + line, context = recorded.get_nowait() + assert context == "inside" + assert delivered.is_set() + event: Final = decode_spend_event(line) + assert isinstance(event, SpendEvent) + assert event.response is not None + assert isinstance(event.response["usage"], dict) + assert event.response["usage"]["total_tokens"] == 15 + assert recorded.empty() + assert cancelled.is_set() is (blocked_phase in ("hook", "delivery")) + assert ("private-marker" in line.decode()) is (blocked_phase not in ("hook", "redacted")) + assert (REDACTED_BY_LITELLM in line.decode()) is (blocked_phase == "redacted") + assert ("private-instructions-canary" in line.decode()) is ( + response_kind == "responses" and blocked_phase in ("neither", "delivery") + ) + delivery: Final = deliveries.get_nowait() + assert not any( + value in delivery + for value in ( + "private-request-instructions", + "private-system-canary", + "private-tool-canary", + ) + ) + assert ("private-instructions-canary" in delivery) is ( + response_kind == "responses" and blocked_phase in ("neither", "delivery") + ) + assert response.model_dump(exclude={"usage"}) == original_response + finally: + await worker.stop() + + +@pytest.mark.asyncio +async def test_invalid_accounting_after_a_hook_timeout_never_exports_unmasked_data( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + + exported: Final = asyncio.Event() + never: Final = asyncio.Event() + + class InvalidMaskingHook(CustomLogger): + async def async_logging_hook( + self, + kwargs: dict[str, object], + result: object, + call_type: str, + ) -> tuple[dict[str, object], object]: + kwargs["stream"] = "invalid" + await never.wait() + return kwargs, result + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + exported.set() + + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj: Final = _async_request_logging_obj("invalid-accounting", [InvalidMaskingHook(), _ProxyDBLogger()]) + worker: Final = LoggingWorker(timeout=0.5) + worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=ModelResponse())) + try: + await worker.flush() + assert not exported.is_set() + assert "Cannot prepare body-free spend accounting after a logging hook timeout" in caplog.text + finally: + await worker.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("no_log", [False, True]) +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_masking_timeout_preserves_opt_outs_and_redacts_hidden_provider_content( + monkeypatch: pytest.MonkeyPatch, + no_log: bool, + metadata_field: str, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress + + recorded: Final[asyncio.Queue[str]] = asyncio.Queue() + exported: Final[asyncio.Queue[str]] = asyncio.Queue() + never: Final = asyncio.Event() + + async def store_event(line: bytes) -> None: + recorded.put_nowait(line.decode()) + + producer: Final = SpendEventProducer( + address=UnixAddress("/unused"), + on_unavailable="fallback", + buffer_size=1, + connect_timeout=1, + fallback=store_event, + ) + await producer.close(drain_timeout=1) + + class MaskingLogger(CustomLogger): + async def async_logging_hook( + self, + kwargs: dict[str, object], + result: object, + call_type: str, + ) -> tuple[dict[str, object], object]: + await never.wait() + return kwargs, result + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + exported.put_nowait(json.dumps(kwargs, default=str) + str(response_obj)) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "global_disable_no_log_param", False) + monkeypatch.setattr(litellm, "allow_dynamic_callback_disabling", True) + monkeypatch.setattr(proxy_server, "premium_user", True) + logging_obj: Final = _async_request_logging_obj( + "private-provider-content", [MaskingLogger(), _ProxyDBLogger(producer)] + ) + logging_obj.update_environment_variables( + litellm_params={ + "acompletion": True, + "no-log": no_log, + metadata_field: { + "hidden_params": { + "provider_specific_fields": { + "reasoning_content": "private-thinking-canary", + "thinking_blocks": [{"type": "thinking", "thinking": "private-thinking-canary"}], + "citations": [{"cited_text": "private-citation-canary"}], + } + } + }, + "proxy_server_request": {"headers": {"x-litellm-disable-callbacks": "dynamodb"}}, + }, + optional_params={}, + ) + worker: Final = LoggingWorker(timeout=0.5) + worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=ModelResponse())) + try: + await worker.flush() + assert recorded.qsize() == int(not no_log) + assert exported.qsize() == int(not no_log) + payloads: Final = tuple(recorded.get_nowait() for _ in range(recorded.qsize())) + tuple( + exported.get_nowait() for _ in range(exported.qsize()) + ) + assert all( + "private-thinking-canary" not in payload and "private-citation-canary" not in payload + for payload in payloads + ) + assert not logging_obj.should_run_callback( + "dynamodb", logging_obj.model_call_details["litellm_params"], "async_success" + ) + finally: + await worker.stop() + + _GO_PRICED_MESSAGE: Final = { "id": "msg_go", "type": "message",