diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9931452fc6e..1dd012383cb 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -723,6 +723,7 @@ class Logging(LiteLLMLoggingBaseClass): # enqueue closure here instead of firing it immediately. self._defer_async_logging: bool = False self._enqueue_deferred_logging: Callable[[], None] | None = None + self._async_success_scheduled: bool = False self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None @@ -2261,6 +2262,15 @@ class Logging(LiteLLMLoggingBaseClass): except Exception: return True + def claim_async_success_log(self) -> bool: + """One async success log per request. The innermost @client wrapper always exits first, + since the outer one is awaiting it, so it claims the log here and the outer wrapper gets + False and schedules nothing.""" + if self._async_success_scheduled: + return False + self._async_success_scheduled = True + return True + def mark_logging_complete( self, event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], diff --git a/litellm/utils.py b/litellm/utils.py index ff8dd321495..0b612828c7a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1373,11 +1373,14 @@ def _schedule_async_success_logging( Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses bridge) each exit through here with the same logging object and their own shape of the same - response. The immediate path already logs one request once, since the first task marks - ``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same - first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log - reads usage from, and a later wrapper never swaps in its client-shaped translation. + response. The innermost wrapper always exits first, since the outer one is awaiting it, so it + claims the async success log here, synchronously, and the outer wrapper returns without + enqueuing: one request queues one handler, whatever the handler tasks do later. The deferred + slot keeps the same innermost-wins rule: the provider-shaped result is the one the spend log + reads usage from, and the outer wrapper never swaps in its client-shaped translation. """ + if not logging_obj.claim_async_success_log(): + return def _enqueue_async_logging() -> None: asyncio.create_task( diff --git a/tests/integration/spend/test_responses_bridge_spend_once.py b/tests/integration/spend/test_responses_bridge_spend_once.py new file mode 100644 index 00000000000..9eddd832ca0 --- /dev/null +++ b/tests/integration/spend/test_responses_bridge_spend_once.py @@ -0,0 +1,111 @@ +import json +from base64 import b64decode +from collections.abc import Callable +from hashlib import sha256 +from typing import Final +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +PROMPT_TOKENS: Final = 30 +COMPLETION_TOKENS: Final = 5 +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +EXPECTED_SPEND: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE + + +def _responses_reply(prompt: str, response_id: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply( + body=json.dumps({"object": "list", "data": [{"id": "gpt-6.1-sol", "object": "model"}]}).encode() + ) + assert request.method == "POST", request.method + assert request.target == "/v1/responses", request.target + body: Final = object_value(json.loads(request.body)) + encoded_input: Final = json.dumps(body["input"]) + assert prompt in encoded_input, encoded_input + assert body.get("stream") is not True, body + response: Final = { + "id": response_id, + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": body["model"], + "output": [ + { + "type": "message", + "id": f"msg_{uuid4().hex}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": { + "input_tokens": PROMPT_TOKENS, + "output_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + } + return Reply(body=json.dumps(response).encode()) + + return respond + + +def _rows_for_key(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT litellm_call_id, request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ) + + +@pytest.mark.timeout(120) +def test_bridged_chat_completion_above_log_offload_threshold_logs_and_charges_once(gateway: Gateway) -> None: + prompt_prefix: Final = f"spend once {uuid4().hex[:8]} " + long_prompt: Final = (prompt_prefix + "lorem ipsum " * 30_000)[:300_000] + upstream_id: Final = f"resp_{uuid4().hex}" + with ( + wire_server(_responses_reply(long_prompt, upstream_id)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/responses/gpt-6.1-sol", + api_base=f"{wire.url}/v1", + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ) + key: Final = scenario.key(models=[model]) + digest: Final = sha256(key.encode()).hexdigest() + + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": long_prompt}], "stream": False}, + key=key, + ) + assert response.status_code == 200, response.text + assert "ok" in response.text, response.text + assert response.json()["id"] == upstream_id, response.text + call_id: Final = response.headers["x-litellm-call-id"] + + rows: Final = eventually( + lambda: _rows_for_key(key), + lambda values: any(row["litellm_call_id"] == call_id for row in values), + seconds=70, + ) + call_rows: Final = [row for row in rows if row["litellm_call_id"] == call_id] + assert len(call_rows) == 1, rows + assert float(str(call_rows[0]["spend"])) == pytest.approx(EXPECTED_SPEND), rows + logged_id: Final = b64decode(str(call_rows[0]["request_id"]).removeprefix("resp_")).decode() + assert logged_id.startswith("litellm:") and f";response_id:{upstream_id}" in logged_id, rows + + key_spend: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(str(values[0]["spend"])) >= EXPECTED_SPEND, + seconds=70, + ) + assert float(str(key_spend[0]["spend"])) == pytest.approx(EXPECTED_SPEND), key_spend diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index b42a6f99ee5..5e85a5903c6 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -5975,7 +5975,7 @@ def _interactions_logging_obj(stream: bool, call_type: str = "acreate"): messages=[], stream=stream, call_type=call_type, - start_time=time.time(), + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc).timestamp(), litellm_call_id="interactions-call-id", function_id="interactions-fn-id", ) @@ -9536,3 +9536,61 @@ def test_get_custom_logger_compatible_class_does_not_match_generic_api_logger( assert logging_module.get_custom_logger_compatible_class(integration) is None finally: logging_module._in_memory_loggers.clear() + + +@pytest.mark.asyncio +async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch): + """ + The create's in_progress success handler can still be parked on awaited work when + the settlement arrives. The completed result is the only event carrying usage and + cost, so it has to reach the success callbacks whatever the one-success-log-per-request + dedupe does, or the settlement is marked billed with no spend row behind it. + """ + from litellm.litellm_core_utils import litellm_logging + from litellm.types.interactions import InteractionsAPIResponse + + class CountingLogger(CustomLogger): + def __init__(self): + super().__init__() + self.logged_results: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_results.append(response_obj) + + counting_logger: Final = CountingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger]) + + in_progress_parked: Final = asyncio.Event() + release_in_progress: Final = asyncio.Event() + real_truncate: Final = litellm_logging.truncate_base64_in_messages_async + calls: list[int] = [] + + async def parked_then_real_truncate(messages): + calls.append(1) + if len(calls) == 1: + in_progress_parked.set() + await release_in_progress.wait() + return await real_truncate(messages) + + monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", parked_then_real_truncate) + + logging_obj: Final = _interactions_logging_obj(stream=False) + in_progress: Final = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + completed: Final = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + first: Final = asyncio.create_task(logging_obj.async_success_handler(result=in_progress)) + await in_progress_parked.wait() + completion: Final = asyncio.create_task(logging_obj.async_log_background_interaction_completion(result=completed)) + await asyncio.sleep(0) + await asyncio.sleep(0) + release_in_progress.set() + await first + await completion + + assert counting_logger.logged_results == [completed, in_progress], counting_logger.logged_results diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 8bd12f7581c..a7d74da489f 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -6711,3 +6711,67 @@ def test_function_setup_never_logs_the_ocr_data_uri_payload() -> None: assert logged == [{"role": "user", "content": f"data:application/pdf;base64 ({len(payload)} chars)"}] assert payload not in str(logged) + + +@pytest.mark.asyncio +async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None: + """Chat over the Responses bridge exits two @client wrappers with one logging object. Issue + #44500: both exits enqueued a success handler, and on a large prompt the first one yielded to + the worker-thread base64 offload before it marked ``has_logged_async_success``, so the second + passed the check too and the request was logged and billed twice. The schedule step claims the + log for the object synchronously, so only the inner provider-shaped result is ever logged.""" + from litellm.litellm_core_utils import litellm_logging + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.utils import _dispatch_success_logging + + class CountingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_results: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + if kwargs["litellm_call_id"] == "bridge-call-id": + self.logged_results.append(response_obj) + + counting_logger: Final = CountingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger]) + real_truncate: Final = litellm_logging.truncate_base64_in_messages_async + + async def yielding_truncate(messages): + await asyncio.sleep(0) + return await real_truncate(messages) + + monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", yielding_truncate) + + messages: Final = [{"role": "user", "content": "hello"}] + logging_obj: Final = Logging( + model="gpt-5.6-luna", + messages=messages, + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc).timestamp(), + litellm_call_id="bridge-call-id", + function_id="bridge-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="gpt-5.6-luna", custom_llm_provider="openai", input=messages + ) + inner_result: Final = ModelResponse(id="inner") + outer_result: Final = ModelResponse(id="outer") + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + for result in (inner_result, outer_result): + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=now, + end_time=now, + is_completion_with_fallbacks=False, + is_litellm_internal_call=False, + ) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + assert len(counting_logger.logged_results) == 1, counting_logger.logged_results + assert counting_logger.logged_results[0] is inner_result