From 2e1a98f5212f29c42fe5df15af06a0a5c55ac160 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 4 Oct 2026 23:13:43 -0700 Subject: [PATCH] fix(logging): deduplicate streaming failure callbacks (#44442) * update logic that marks a logging callback as complete * test(logging): cover streaming failure dedupe in mark_logging_complete Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): cover streaming failure dedupe in S3 and DataDog Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): keep has_run_logging as a deprecated alias Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): audit streaming failure dedupe across surfaces, fallbacks and sink outage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): assert anthropic upstream path in streaming failure audit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): count only provider posts in streaming audit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): assert the sink outage rejects uploads in burst audit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): configure datadog retries with router override Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Mrinal Chanshetty Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 19 +- .../_experimental/mcp_server/operations.py | 8 +- tests/e2e/coverage_registry/logging.yaml | 1 + tests/e2e/logging/test_datadog_log_e2e.py | 69 ++- .../test_s3_v2_streaming_failure_dedupe.py | 408 ++++++++++++++++++ tests/unit/integrations/test_s3_v2.py | 60 +++ .../test_litellm_logging.py | 160 ++++++- .../test_mcp_server_tool_calls_and_headers.py | 6 +- 8 files changed, 710 insertions(+), 21 deletions(-) create mode 100644 tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5a3a17f338c..b2c22880f5a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2259,11 +2259,11 @@ class Logging(LiteLLMLoggingBaseClass): except Exception: return True - def has_run_logging( + def mark_logging_complete( self, event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], ) -> None: - if self.stream is not None and self.stream is True: + if self.stream is not None and self.stream is True and event_type in ["async_success", "sync_success"]: """ Ignore check on stream, as there can be multiple chunks """ @@ -2271,6 +2271,13 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details[f"has_logged_{event_type}"] = True return + def has_run_logging( + self, + event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], + ) -> None: + """Deprecated alias of mark_logging_complete, kept for callers of the old name""" + self.mark_logging_complete(event_type=event_type) + def should_run_callback(self, callback: litellm.CALLBACK_TYPES, litellm_params: dict, event_hook: str) -> bool: if litellm.global_disable_no_log_param: return True @@ -2862,7 +2869,7 @@ class Logging(LiteLLMLoggingBaseClass): call_type=self.call_type, ) - self.has_run_logging(event_type="sync_success") + self.mark_logging_complete(event_type="sync_success") for callback in callbacks: try: should_run = self.should_run_callback( @@ -3443,7 +3450,7 @@ class Logging(LiteLLMLoggingBaseClass): ) self._handle_callback_failure(callback=callback) - self.has_run_logging(event_type="async_success") + self.mark_logging_complete(event_type="async_success") for callback in callbacks: # check if callback can run for this request @@ -3736,7 +3743,7 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=result, ) - self.has_run_logging(event_type="sync_failure") + self.mark_logging_complete(event_type="sync_failure") for callback in callbacks: try: should_run = self.should_run_callback( @@ -3928,7 +3935,7 @@ class Logging(LiteLLMLoggingBaseClass): result: Final = None # result sent to all loggers, init this to None incase it's not created - self.has_run_logging(event_type="async_failure") + self.mark_logging_complete(event_type="async_failure") for callback in callbacks: try: litellm_params = self.model_call_details.get("litellm_params", {}) diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index d83a4a72e2f..1dea44a84f0 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -2286,8 +2286,8 @@ async def _run_post_mcp_call_guardrails( def suppress_completed_success_logging(logging_obj: LiteLLMLoggingObj) -> None: """An interim ``InputRequiredResult`` is not a completed call, so the ``@client`` wrapper on ``call_mcp_tool`` must not run the success handlers for it when the coroutine returns.""" - logging_obj.has_run_logging(event_type="sync_success") - logging_obj.has_run_logging(event_type="async_success") + logging_obj.mark_logging_complete(event_type="sync_success") + logging_obj.mark_logging_complete(event_type="async_success") async def _fire_mcp_tool_call_logging( @@ -2335,8 +2335,8 @@ async def _fire_mcp_tool_call_logging( await logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time) return result - logging_obj.has_run_logging(event_type="sync_success") - logging_obj.has_run_logging(event_type="async_success") + logging_obj.mark_logging_complete(event_type="sync_success") + logging_obj.mark_logging_complete(event_type="async_success") tool_error: Final = MCPToolResultError(error_message) logging_obj.failure_handler(tool_error, "", start_time, end_time) await logging_obj.async_failure_handler(tool_error, "", start_time, end_time) diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index 5424fb2d12f..9a851587fce 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -6,6 +6,7 @@ - {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"} - {id: logging.datadog.stream.exports_metric, module: logging, tier: P0, event: stream, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses], source: "integrations/datadog/datadog.py", rationale: "Streaming aggregates usage after the last chunk; delivery and cost must survive that path"} - {id: logging.datadog.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "integrations/datadog/datadog.py", rationale: "Failure metrics for alerting/SLO"} +- {id: logging.datadog.stream_failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "litellm_core_utils/litellm_logging.py", rationale: "Streamed failures re-enter the failure handler once per retry on the same logging object; dedup must hold on the stream path too (#42988)"} - {id: logging.prometheus.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/prometheus.py", rationale: "Standard OSS metrics; per-key cardinality (existing e2e)"} - {id: logging.prometheus.success.records_queue_time, module: logging, tier: P1, event: success, assertions: [records_queue_time], exercised_on: [chat_completions], source: "integrations/prometheus.py / LIT-2034", fail_before_fix: proven, rationale: "Queue time feeds saturation alerting; the family stayed registered while no observation was ever recorded, so presence alone is not the contract"} - {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"} diff --git a/tests/e2e/logging/test_datadog_log_e2e.py b/tests/e2e/logging/test_datadog_log_e2e.py index a4821ed058b..2d57181c4bf 100644 --- a/tests/e2e/logging/test_datadog_log_e2e.py +++ b/tests/e2e/logging/test_datadog_log_e2e.py @@ -22,13 +22,12 @@ import math import time import pytest -from pydantic import BaseModel, ConfigDict - from datadog_reader import DdLogEvent, DdLogsReader from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker from lifecycle import ResourceManager from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok, readiness_details_body -from models import LiteLLMParamsBody +from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride +from pydantic import BaseModel, ConfigDict pytestmark = pytest.mark.e2e @@ -398,3 +397,67 @@ class TestDataDogFailureDelivery: assert payload.error_str is not None and "AnthropicException" in payload.error_str, ( f"the event must carry the provider error, got error_str={payload.error_str!r}" ) + + @pytest.mark.covers("logging.datadog.stream_failure.exports_metric", exercised_on=["chat_completions"]) + def test_failed_chat_completions_stream_emits_one_error_event( + self, client: LoggingClient, dd_logs: DdLogsReader, resources: ResourceManager + ) -> None: + """A STREAMED /chat/completions call that fails at the provider after + its configured retries must still reach the DataDog logs intake as + exactly one error-grade event: the retry loop invokes the failure + handler once per attempt on the same logging object, so a dedup that + only works for non-streaming calls multiplies every retried stream + failure by its attempt count (issue #42988). + + The deployment fails with a connect error, not an auth error: the + router does not retry AuthenticationError when the model group has a + single deployment, so an invalid key would never reach the retry loop + this test exercises. api_base is an unroutable address, so every + attempt fails the same retryable way.""" + _assert_datadog_configured(client) + + model_name = f"dd-err-stream-{unique_marker()}" + model_id = client.create_model( + model_name, + LiteLLMParamsBody( + model="anthropic/claude-haiku-4-5", + api_key=INVALID_UPSTREAM_API_KEY, + api_base="http://localhost:1", + ), + ) + resources.defer(lambda: client.delete_model(model_id)) + key = client.key_with_alias(f"dd-err-stream-key-{unique_marker()}", models=[model_name]) + resources.defer(lambda: client.delete_key(key)) + + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + outcome = client.proxy.transport.stream( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=model_name, + messages=[ChatMessage(role="user", content="trigger an upstream connect failure")], + stream=True, + max_tokens=16, + router_settings_override=RouterSettingsOverride(num_retries=2), + ), + ) + assert not outcome.ok, "the call must fail; the deployment's upstream is unreachable" + assert outcome.status_code != -1, ( + "network failure between the test and the proxy while provoking the provider " + "failure; retrying now could double-log the failure payload and falsely trip " + f"the exactly-one assertion - fix the rig connectivity first: {outcome.body[:200]}" + ) + if "AnthropicException" in outcome.body or time.monotonic() >= deadline: + break + time.sleep(client.proxy.poll_interval) + assert "AnthropicException" in outcome.body, ( + "never saw the upstream provider failure before the deadline; the deployment may still " + f"be propagating - last outcome {outcome.status_code}: {outcome.body[:200]}" + ) + + events = dd_logs.poll_events_for_query(f"@model_group:{model_name}") + payload = _assert_exactly_one_failure_event(events, model_group=model_name) + assert payload.error_str is not None and "AnthropicException" in payload.error_str, ( + f"the event must carry the provider error, got error_str={payload.error_str!r}" + ) diff --git a/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py b/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py new file mode 100644 index 00000000000..ee347592d37 --- /dev/null +++ b/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py @@ -0,0 +1,408 @@ +from __future__ import annotations + +import json +import uuid +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final, Literal + +import httpx +import pytest +from _s3_v2_support import RecordingS3Sink, call_surface, collect_payloads, matched_ids, s3_config, surface_reply +from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +Surface = Literal["chat", "messages", "responses"] +SURFACES: Final[tuple[Surface, ...]] = ("chat", "messages", "responses") +FAILURE_BODY: Final[bytes] = json.dumps( + { + "error": { + "message": "synthetic upstream failure", + "type": "server_error", + "param": None, + "code": "synthetic_failure", + } + } +).encode() + + +def _failure_provider(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + return Reply(status=503, body=FAILURE_BODY) + + +def _surface_target(surface: Surface) -> str: + return { + "chat": "/v1/chat/completions", + "messages": "/v1/messages", + "responses": "/v1/responses", + }[surface] + + +def _surface_request( + candidate: Gateway, + surface: Surface, + openai_model: str, + anthropic_model: str, + key: str, + marker: str, + stream: bool, +) -> httpx.Response: + if surface == "chat": + body: Final = { + "model": openai_model, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + } + return candidate.request("POST", _surface_target(surface), body, key=key) + if surface == "messages": + body: Final = { + "model": anthropic_model, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + } + return candidate.request("POST", _surface_target(surface), body, key=key) + body: Final = {"model": openai_model, "input": marker, "stream": stream} + return candidate.request("POST", _surface_target(surface), body, key=key) + + +def _responses_stream_id(response: httpx.Response) -> str: + events: Final = tuple( + object_value(json.loads(line.removeprefix("data: "))) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + completed: Final = tuple( + object_value(event["response"]) + for event in events + if event.get("type") == "response.completed" + ) + assert len(completed) == 1, events + response_id: Final = completed[0]["id"] + assert isinstance(response_id, str), response_id + return response_id + + +def _register_models(scenario: Scenario, upstream_url: str) -> tuple[str, str]: + openai_model: Final = scenario.model( + model="openai/gpt-4o-mini", + api_base=upstream_url + "/v1", + api_key="synthetic-provider-key", + ) + anthropic_model: Final = scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", + api_base=upstream_url, + api_key="synthetic-provider-key", + ) + return openai_model, anthropic_model + + +def _request_key(payload: dict[str, JsonValue], response: httpx.Response, marker: str) -> str | None: + call_id: Final = response.headers.get("x-litellm-call-id") + if call_id is not None and call_id in {str(payload.get("id")), str(payload.get("litellm_call_id"))}: + return call_id + if marker in json.dumps(payload): + return marker + return None + + +def _assert_upstream_requests(requests: tuple[Request, ...], surface: Surface, expected: int) -> None: + target: Final = _surface_target(surface) + posts: Final = tuple(request for request in requests if request.method == "POST") + observed: Final = tuple(request.target for request in posts) + assert len(posts) == expected, f"expected {expected} upstream POSTs, observed {len(posts)}: {observed}" + assert all(request.target == target for request in posts), f"expected {target}, observed {observed}" + + +def _assert_one_payload( + sink: RecordingS3Sink, + response: httpx.Response, + marker: str, + status: str, + upstream_posts: int, +) -> tuple[dict[str, JsonValue], ...]: + first_payloads: Final = collect_payloads(sink, 1) + assert first_payloads[0]["status"] == status + # A full three-flush-interval window is needed to detect late duplicate uploads. + payloads: Final = eventually( + lambda: sink.payloads(), + lambda observed: len(observed) >= 2, + seconds=6, + return_last_on_timeout=True, + ) + assert len(payloads) == 1, ( + f"expected one {status} payload after {upstream_posts} upstream POSTs, " + f"observed {len(payloads)} payloads" + ) + assert ( + _request_key(payloads[0], response, marker) == response.headers.get("x-litellm-call-id") + or marker in json.dumps(payloads[0]) + ), f"payload did not match request id or marker: {payloads[0]!r}" + return payloads + + +@pytest.mark.parametrize( + ("surface", "stream"), + [ + pytest.param("chat", False, id="chat-nonstream"), + pytest.param("chat", True, id="chat-stream"), + pytest.param("messages", False, id="messages-nonstream"), + pytest.param("messages", True, id="messages-stream"), + pytest.param("responses", False, id="responses-nonstream"), + pytest.param("responses", True, id="responses-stream"), + ], +) +def test_retried_failure_uploads_one_s3_object( + gateway: Gateway, + tmp_path: Path, + surface: Surface, + stream: bool, +) -> None: + marker: Final = f"s3-a-{surface}-{uuid.uuid4().hex}" + sink: Final = RecordingS3Sink() + with wire_server(_failure_provider) as upstream, wire_server(sink.respond) as bucket: + config: Final = s3_config(tmp_path, bucket.url, {}, settings={"num_retries": 2}) + with ( + owned_proxy( + gateway, + tmp_path, + {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, + config=config, + workers=2, + ) as candidate, + candidate.scenario() as scenario, + ): + openai_model, anthropic_model = _register_models(scenario, upstream.url) + key: Final = scenario.key(models=[openai_model, anthropic_model]) + response: Final = _surface_request( + candidate, + surface, + openai_model, + anthropic_model, + key, + marker, + stream, + ) + requests: Final = upstream.drain() + _assert_upstream_requests(requests, surface, 3) + assert not 200 <= response.status_code < 300, f"{response.status_code}: {response.text}" + assert "synthetic upstream failure" in response.text, response.text + _assert_one_payload(sink, response, marker, "failure", len(requests)) + + +@pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) +def test_fallback_chain_failure_uploads_one_s3_object( + gateway: Gateway, + tmp_path: Path, + stream: bool, +) -> None: + marker: Final = f"s3-b-{uuid.uuid4().hex}" + sink: Final = RecordingS3Sink() + with ( + wire_server(_failure_provider) as primary_upstream, + wire_server(_failure_provider) as secondary_upstream, + wire_server(sink.respond) as bucket, + ): + config: Final = s3_config(tmp_path, bucket.url, {}, settings={"num_retries": 0}) + with ( + owned_proxy( + gateway, + tmp_path, + {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, + config=config, + workers=2, + ) as candidate, + candidate.scenario() as scenario, + ): + primary_model: Final = scenario.model( + model="openai/gpt-4o-mini", + api_base=primary_upstream.url + "/v1", + api_key="synthetic-primary-key", + ) + secondary_model: Final = scenario.model( + model="openai/gpt-4o-mini", + api_base=secondary_upstream.url + "/v1", + api_key="synthetic-secondary-key", + ) + key: Final = scenario.key(models=[primary_model, secondary_model]) + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": primary_model, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + "fallbacks": [secondary_model], + "num_retries": 0, + }, + key=key, + ) + primary_requests: Final = primary_upstream.drain() + secondary_requests: Final = secondary_upstream.drain() + _assert_upstream_requests(primary_requests, "chat", 1) + _assert_upstream_requests(secondary_requests, "chat", 1) + assert not 200 <= response.status_code < 300, f"{response.status_code}: {response.text}" + assert "synthetic upstream failure" in response.text, response.text + _assert_one_payload(sink, response, marker, "failure", len(primary_requests) + len(secondary_requests)) + + +@pytest.mark.parametrize("surface", SURFACES, ids=SURFACES) +def test_streaming_success_uploads_one_s3_object( + gateway: Gateway, + tmp_path: Path, + surface: Surface, +) -> None: + marker: Final = f"s3-c-{surface}-{uuid.uuid4().hex}" + sink: Final = RecordingS3Sink() + with wire_server(surface_reply) as upstream, wire_server(sink.respond) as bucket: + config: Final = s3_config(tmp_path, bucket.url, {}, settings={"num_retries": 0}) + with ( + owned_proxy( + gateway, + tmp_path, + {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, + config=config, + workers=2, + ) as candidate, + candidate.scenario() as scenario, + ): + openai_model, anthropic_model = _register_models(scenario, upstream.url) + key: Final = scenario.key(models=[openai_model, anthropic_model]) + captured_responses: Final[list[httpx.Response]] = [] + candidate.client.event_hooks["response"].append(captured_responses.append) + surface_response_id, call_id = call_surface( + candidate, + f"{surface}_stream", + openai_model, + anthropic_model, + key, + marker, + ) + if surface == "responses": + assert len(captured_responses) == 1, captured_responses + response_id: Final = ( + _responses_stream_id(captured_responses[0]) if surface == "responses" else surface_response_id + ) + payloads: Final = collect_payloads(sink, 1) + payloads_after_window: Final = eventually( + lambda: sink.payloads(), + lambda observed: len(observed) >= 2, + seconds=6, + return_last_on_timeout=True, + ) + assert len(payloads_after_window) == 1, payloads_after_window + assert payloads[0]["status"] == "success" + if surface == "responses": + assert payloads[0]["litellm_call_id"] == call_id + else: + assert payloads[0]["id"] == response_id + assert matched_ids(payloads, ((response_id, call_id),)) == frozenset({str(payloads[0]["id"])}) + + +def test_failure_burst_through_sink_outage_lands_each_request_once( + gateway: Gateway, + tmp_path: Path, +) -> None: + requests_per_variant: Final = 4 + jobs: Final = tuple( + (surface, stream, f"s3-d-{surface}-{'stream' if stream else 'nonstream'}-{index}") + for surface in SURFACES + for stream in (False, True) + for index in range(requests_per_variant) + ) + sink: Final = RecordingS3Sink() + with wire_server(_failure_provider) as upstream, wire_server(sink.respond) as bucket: + config: Final = s3_config(tmp_path, bucket.url, {}, settings={"num_retries": 2}) + with ( + owned_proxy( + gateway, + tmp_path, + {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, + config=config, + workers=2, + ) as candidate, + candidate.scenario() as scenario, + ): + openai_model, anthropic_model = _register_models(scenario, upstream.url) + key: Final = scenario.key(models=[openai_model, anthropic_model]) + sink.fail_until = float("inf") + + def call(job: tuple[Surface, bool, str]) -> httpx.Response: + surface, stream, marker = job + return _surface_request(candidate, surface, openai_model, anthropic_model, key, marker, stream) + + with ThreadPoolExecutor(max_workers=len(jobs)) as pool: + responses: Final = tuple(pool.map(call, jobs)) + + upstream_requests: Final = tuple( + request for request in upstream.drain() if request.method == "POST" + ) + route_counts: Final = Counter(request.target for request in upstream_requests) + expected_route_counts: Final = { + _surface_target(surface): 2 * requests_per_variant * 3 for surface in SURFACES + } + assert len(upstream_requests) == len(jobs) * 3, ( + f"expected {len(jobs) * 3} upstream POSTs, observed {len(upstream_requests)}: {route_counts}" + ) + assert route_counts == expected_route_counts, f"unexpected upstream routes: {route_counts}" + assert all(not 200 <= response.status_code < 300 for response in responses) + assert all("synthetic upstream failure" in response.text for response in responses) + eventually( + lambda: sink.attempts, + lambda attempts: attempts > len(sink.store), + seconds=30, + ) + rejected: Final = sink.attempts - len(sink.store) + assert rejected >= 1, f"expected at least one rejected sink upload, rejected={rejected}" + assert len(sink.store) == 0, ( + f"expected all sink uploads to fail before recovery, rejected={rejected}, " + f"attempts={sink.attempts}" + ) + sink.fail_until = 0.0 + payloads: Final = eventually( + lambda: sink.payloads(), + lambda observed: len(observed) >= len(jobs), + seconds=90, + return_last_on_timeout=True, + ) + payloads_after_window: Final = eventually( + lambda: sink.payloads(), + lambda observed: len(observed) > len(payloads), + seconds=6, + return_last_on_timeout=True, + ) + payload_matches: Final = tuple( + tuple( + index + for index, (response, job) in enumerate(zip(responses, jobs)) + if _request_key(payload, response, job[2]) is not None + ) + for payload in payloads_after_window + ) + assert all(len(matches) == 1 for matches in payload_matches), ( + f"payloads could not be matched to one request: rejected={rejected}, " + f"matches={payload_matches}, payloads={payloads_after_window}" + ) + landed_keys: Final = tuple( + _request_key(payload, responses[matches[0]], jobs[matches[0]][2]) + for payload, matches in zip(payloads_after_window, payload_matches) + ) + duplicate_keys: Final = tuple(key for key in set(landed_keys) if landed_keys.count(key) > 1) + expected_keys: Final = tuple( + response.headers.get("x-litellm-call-id") or job[2] for response, job in zip(responses, jobs) + ) + missing_keys: Final = tuple(key for key in expected_keys if key not in landed_keys) + assert not duplicate_keys, ( + f"duplicate request ids landed: {duplicate_keys}; " + f"missing={missing_keys}; rejected={rejected}; upstream_posts={len(upstream_requests)} " + f"sink_payloads={len(payloads_after_window)}" + ) + assert not missing_keys, ( + f"requests missing after sink recovery: {missing_keys}; " + f"rejected={rejected}; upstream_posts={len(upstream_requests)} " + f"sink_payloads={len(payloads_after_window)}" + ) diff --git a/tests/unit/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py index 963586d9532..ac905335f81 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -17,7 +17,10 @@ import httpx import pytest import respx +import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.s3_v2 import S3BatchUploadError, S3Logger +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.integrations.s3_v2 import s3BatchLoggingElement from litellm.types.utils import StandardLoggingPayload @@ -5069,3 +5072,60 @@ async def test_send_batch_time_grows_linearly_with_the_batch() -> None: quadrupled: Final = await _timed_send_batch(8_000) assert quadrupled / baseline < 8, f"2k took {baseline:.3f}s, 8k took {quadrupled:.3f}s" + + +class QueueOnlyS3LoggerWithoutCredentialsOrPeriodicFlush(S3Logger): + def __init__(self) -> None: + CustomLogger.__init__(self) + self.log_queue: list[s3BatchLoggingElement] = [] + self.batch_size = 100 + self.s3_strip_base64_files = False + self.s3_use_team_prefix = False + self.s3_use_key_prefix = False + self.s3_path = "" + self.s3_log_prompts_only = None + self.s3_batch_file_upload = False + self._upload_semaphore = asyncio.Semaphore(1) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("concurrent", [False, True]) +@pytest.mark.asyncio +async def test_repeated_failure_notifications_upload_once( + monkeypatch: pytest.MonkeyPatch, + stream: bool, + concurrent: bool, +) -> None: + sink: Final = QueueOnlyS3LoggerWithoutCredentialsOrPeriodicFlush() + upload: Final = AsyncMock(return_value=True) + monkeypatch.setattr(sink, "async_upload_data_to_s3", upload) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "_async_failure_callback", [sink]) + logging_obj: Final = Logging( + model="openai/gpt-4o-mini", + messages=[], + stream=stream, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="synthetic-request", + function_id="synthetic-function", + ) + error: Final = litellm.ServiceUnavailableError( + message="synthetic deployment-selection failure", + llm_provider="openai", + model="openai/gpt-4o-mini", + ) + + if concurrent: + await asyncio.gather( + *(logging_obj.async_failure_handler(error, "synthetic traceback") for _ in range(3)) + ) + else: + for _ in range(3): + await logging_obj.async_failure_handler(error, "synthetic traceback") + + assert len(sink.log_queue) == 1 + assert all(entry.payload["status"] == "failure" for entry in sink.log_queue) + await sink.async_send_batch() + await asyncio.sleep(0) + assert upload.await_count == 1 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 4b25da2ff79..c726127f892 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1540,11 +1540,161 @@ def test_logging_prevent_double_logging(logging_obj): This is to avoid double logging. """ logging_obj.stream = False - logging_obj.has_run_logging(event_type="sync_success") - assert logging_obj.should_run_logging(event_type="sync_success") == False - assert logging_obj.should_run_logging(event_type="sync_failure") == True - assert logging_obj.should_run_logging(event_type="async_success") == True - assert logging_obj.should_run_logging(event_type="async_failure") == True + logging_obj.mark_logging_complete(event_type="sync_success") + assert logging_obj.should_run_logging(event_type="sync_success") is False + assert logging_obj.should_run_logging(event_type="sync_failure") is True + assert logging_obj.should_run_logging(event_type="async_success") is True + assert logging_obj.should_run_logging(event_type="async_failure") is True + + +class _FailureCountingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.async_failure_events: asyncio.Queue[None] = asyncio.Queue() + self.sync_failure_events: asyncio.Queue[None] = asyncio.Queue() + + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + self.async_failure_events.put_nowait(None) + + def log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + self.sync_failure_events.put_nowait(None) + + +def _register_failure_counting_logger( + monkeypatch: pytest.MonkeyPatch, + logger: _FailureCountingLogger, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", [logger]) + monkeypatch.setattr(litellm, "_async_failure_callback", [logger]) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + "event", + ["async_success", "sync_success", "async_failure", "sync_failure"], +) +def test_mark_logging_complete_flags_by_stream_and_event( + logging_obj: LitellmLogging, + stream: bool, + event: Literal["async_success", "sync_success", "async_failure", "sync_failure"], +) -> None: + events: Final[tuple[Literal["async_success", "sync_success", "async_failure", "sync_failure"], ...]] = ( + "async_success", + "sync_success", + "async_failure", + "sync_failure", + ) + logging_obj.stream = stream + + logging_obj.mark_logging_complete(event_type=event) + + expected_to_run: Final = stream and event in ("async_success", "sync_success") + assert logging_obj.should_run_logging(event_type=event) is expected_to_run + assert all(logging_obj.should_run_logging(event_type=other_event) for other_event in events if other_event != event) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_has_run_logging_alias_marks_logging_complete( + logging_obj: LitellmLogging, + stream: bool, +) -> None: + logging_obj.stream = stream + + logging_obj.has_run_logging(event_type="async_failure") + assert logging_obj.should_run_logging(event_type="async_failure") is False + + logging_obj.has_run_logging(event_type="async_success") + assert logging_obj.should_run_logging(event_type="async_success") is stream + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("concurrent", [False, True]) +async def test_async_failure_handler_dispatches_once_on_repeated_notifications( + logging_obj: LitellmLogging, + monkeypatch: pytest.MonkeyPatch, + stream: bool, + concurrent: bool, +) -> None: + failure_logger: Final = _FailureCountingLogger() + _register_failure_counting_logger(monkeypatch, failure_logger) + logging_obj.stream = stream + logging_obj.model_call_details["litellm_params"] = {} + error: Final = litellm.ServiceUnavailableError("503", "anthropic", "claude") + failure_calls: Final = tuple(logging_obj.async_failure_handler(error, "tb") for _ in range(3)) + + if concurrent: + await asyncio.gather(*failure_calls) + else: + for failure_call in failure_calls: + await failure_call + + assert failure_logger.async_failure_events.qsize() == 1 + + +def test_sync_failure_handler_dispatches_once_on_repeated_streaming_notifications( + logging_obj: LitellmLogging, + monkeypatch: pytest.MonkeyPatch, +) -> None: + failure_logger: Final = _FailureCountingLogger() + _register_failure_counting_logger(monkeypatch, failure_logger) + logging_obj.stream = True + logging_obj.call_type = "completion" + logging_obj.model_call_details["litellm_params"] = {} + error: Final = litellm.ServiceUnavailableError("503", "anthropic", "claude") + + for _ in range(3): + logging_obj.failure_handler(error, "tb") + + assert failure_logger.sync_failure_events.qsize() == 1 + + +@pytest.mark.asyncio +async def test_streaming_sync_and_async_failure_dedupe_independently( + logging_obj: LitellmLogging, + monkeypatch: pytest.MonkeyPatch, +) -> None: + failure_logger: Final = _FailureCountingLogger() + _register_failure_counting_logger(monkeypatch, failure_logger) + logging_obj.stream = True + logging_obj.call_type = "completion" + logging_obj.model_call_details["litellm_params"] = {} + error: Final = litellm.ServiceUnavailableError("503", "anthropic", "claude") + + await logging_obj.async_failure_handler(error, "tb") + await logging_obj.async_failure_handler(error, "tb") + logging_obj.failure_handler(error, "tb") + logging_obj.failure_handler(error, "tb") + + assert failure_logger.async_failure_events.qsize() == 1 + assert failure_logger.sync_failure_events.qsize() == 1 + + +def test_streaming_success_is_not_marked_complete(logging_obj: LitellmLogging) -> None: + logging_obj.stream = True + + logging_obj.mark_logging_complete(event_type="async_success") + logging_obj.mark_logging_complete(event_type="sync_success") + + assert logging_obj.should_run_logging(event_type="async_success") is True + assert logging_obj.should_run_logging(event_type="sync_success") is True + assert "has_logged_async_success" not in logging_obj.model_call_details + assert "has_logged_sync_success" not in logging_obj.model_call_details @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 545b2757ffd..556fe0b2536 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -9191,8 +9191,8 @@ async def test_fire_mcp_tool_call_logging_iserror_logs_failure(): tool_error = logging_obj.async_failure_handler.await_args.args[0] assert isinstance(tool_error, MCPToolResultError) assert str(tool_error) == "upstream exploded" - logging_obj.has_run_logging.assert_any_call(event_type="sync_success") - logging_obj.has_run_logging.assert_any_call(event_type="async_success") + logging_obj.mark_logging_complete.assert_any_call(event_type="sync_success") + logging_obj.mark_logging_complete.assert_any_call(event_type="async_success") proxy_logging_mock.post_call_failure_hook.assert_awaited_once() hook_kwargs = proxy_logging_mock.post_call_failure_hook.await_args.kwargs assert hook_kwargs["route"] == "/mcp/call_tool" @@ -9768,7 +9768,7 @@ async def test_call_mcp_tool_modern_interim_result_passes_through_without_comple logging_obj.async_post_mcp_tool_call_hook.assert_not_awaited() proxy_logging_mock.post_mcp_call_hook.assert_not_awaited() proxy_logging_mock.post_call_failure_hook.assert_not_awaited() - assert sorted(c.kwargs["event_type"] for c in logging_obj.has_run_logging.call_args_list) == [ + assert sorted(c.kwargs["event_type"] for c in logging_obj.mark_logging_complete.call_args_list) == [ "async_success", "sync_success", ], "the @client wrapper would otherwise log the interim result as a completed success on return"