diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 24328549094..eba6c862f7a 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -16,23 +16,32 @@ import asyncio import os import time import traceback -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar from urllib.parse import urlparse from litellm._logging import verbose_logger +from litellm.integrations.batch_utils import ( + BatchSendCancelled, + send_batch_with_413_split, + undelivered_after_http_error, +) from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( + MaskedHTTPStatusError, get_async_httpx_client, httpxSpecialProvider, ) +from litellm.types.integrations.azure_sentinel import AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com" DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default" +_QueuedPayload = TypeVar("_QueuedPayload", StandardLoggingPayload, StandardAuditLogPayload) + MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType( { "login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE, @@ -153,6 +162,8 @@ class AzureSentinelLogger(CustomBatchLogger): asyncio.create_task(self.periodic_flush()) self.log_queue: list[StandardLoggingPayload] = [] self.audit_log_queue: list[StandardAuditLogPayload] = [] + self.logs_awaiting_retry = False + self.audit_logs_awaiting_retry = False @staticmethod def _normalize_authority_host(authority_host: str) -> str: @@ -245,8 +256,8 @@ class AzureSentinelLogger(CustomBatchLogger): self.log_queue.append(standard_logging_payload) - if len(self.log_queue) >= self.batch_size: - await self.async_send_batch() + if len(self.log_queue) >= self.batch_size and not self.logs_awaiting_retry: + await self._threshold_send_logs() except Exception as e: verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc()) @@ -275,8 +286,8 @@ class AzureSentinelLogger(CustomBatchLogger): self.log_queue.append(standard_logging_payload) - if len(self.log_queue) >= self.batch_size: - await self.async_send_batch() + if len(self.log_queue) >= self.batch_size and not self.logs_awaiting_retry: + await self._threshold_send_logs() except Exception as e: verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc()) @@ -298,12 +309,24 @@ class AzureSentinelLogger(CustomBatchLogger): self.audit_log_queue.append(audit_log) - if len(self.audit_log_queue) >= self.batch_size: - await self.async_send_audit_batch() + if len(self.audit_log_queue) >= self.batch_size and not self.audit_logs_awaiting_retry: + await self._threshold_send_audit_logs() except Exception as e: verbose_logger.exception("Azure Sentinel Audit Log Layer Error - %s\n%s", e, traceback.format_exc()) + async def _threshold_send_logs(self) -> None: + async with self.flush_lock: + if self.logs_awaiting_retry or len(self.log_queue) < self.batch_size: + return + await self.async_send_batch() + + async def _threshold_send_audit_logs(self) -> None: + async with self.flush_lock: + if self.audit_logs_awaiting_retry or len(self.audit_log_queue) < self.batch_size: + return + await self.async_send_audit_batch() + async def async_send_batch(self): """ Sends the batch of logs to Azure Monitor Logs Ingestion API @@ -311,67 +334,110 @@ class AzureSentinelLogger(CustomBatchLogger): Raises: Raises a NON Blocking verbose_logger.exception if an error occurs """ - await self._async_send_batch_to_api( - log_queue=self.log_queue, - api_endpoint=self.api_endpoint, - log_type="logs", - ) + batch_to_send: Final = tuple(self.log_queue) + self.log_queue = [] # mutable-ok: queue ownership is detached before the async send + try: + undelivered: Final = await self._async_send_batch_to_api( + log_queue=batch_to_send, + api_endpoint=self.api_endpoint, + log_type="logs", + ) + except BatchSendCancelled as cancelled: + self.log_queue = self._requeue(cancelled.undelivered, self.log_queue, "logs") + self.logs_awaiting_retry = bool(self.log_queue) + raise asyncio.CancelledError() from cancelled + except asyncio.CancelledError: + self.log_queue = self._requeue(batch_to_send, self.log_queue, "logs") + self.logs_awaiting_retry = bool(self.log_queue) + raise + self.log_queue = self._requeue(undelivered, self.log_queue, "logs") + self.logs_awaiting_retry = bool(undelivered) and bool(self.log_queue) async def async_send_audit_batch(self): """ Sends the batch of audit logs to Azure Monitor Logs Ingestion API """ - await self._async_send_batch_to_api( - log_queue=self.audit_log_queue, - api_endpoint=self.audit_api_endpoint, - log_type="audit logs", + batch_to_send: Final = tuple(self.audit_log_queue) + self.audit_log_queue = [] # mutable-ok: queue ownership is detached before the async send + try: + undelivered: Final = await self._async_send_batch_to_api( + log_queue=batch_to_send, + api_endpoint=self.audit_api_endpoint, + log_type="audit logs", + ) + except BatchSendCancelled as cancelled: + self.audit_log_queue = self._requeue(cancelled.undelivered, self.audit_log_queue, "audit logs") + self.audit_logs_awaiting_retry = bool(self.audit_log_queue) + raise asyncio.CancelledError() from cancelled + except asyncio.CancelledError: + self.audit_log_queue = self._requeue(batch_to_send, self.audit_log_queue, "audit logs") + self.audit_logs_awaiting_retry = bool(self.audit_log_queue) + raise + self.audit_log_queue = self._requeue(undelivered, self.audit_log_queue, "audit logs") + self.audit_logs_awaiting_retry = bool(undelivered) and bool(self.audit_log_queue) + + def _requeue( + self, + undelivered: tuple[_QueuedPayload, ...], + queue: list[_QueuedPayload], + log_type: str, + ) -> list[_QueuedPayload]: + merged: Final = [*undelivered, *queue] # mutable-ok: queue trimming returns a mutable logger queue + overflow: Final = len(merged) - self.max_queue_size + if overflow <= 0: + return merged + + verbose_logger.warning( + "Azure Sentinel: %s queue exceeded max_queue_size=%s, dropped %s oldest records", + log_type, + self.max_queue_size, + overflow, ) + return merged[overflow:] async def _async_send_batch_to_api( self, - log_queue: list[StandardLoggingPayload | StandardAuditLogPayload], + log_queue: tuple[_QueuedPayload, ...], api_endpoint: str, log_type: str, - ) -> None: + ) -> tuple[_QueuedPayload, ...]: + if not log_queue: + return () + + verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type) try: - if not log_queue: - return - - verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type) - - # Get OAuth2 token bearer_token: Final = await self._get_oauth_token() + except MaskedHTTPStatusError as e: + return undelivered_after_http_error(log_queue, e.status_code, "Azure Sentinel OAuth token", str(e)) + except Exception as e: + verbose_logger.exception("Azure Sentinel Error getting OAuth token - %s", e) + return tuple(log_queue) - # Convert log queue to JSON array format expected by Logs Ingestion API - # Each log entry should be a JSON object in the array - body: Final = safe_dumps(log_queue) + headers: Final = { + "Authorization": f"Bearer {bearer_token}", + "Content-Type": "application/json", + } - # Set headers for Logs Ingestion API - headers: Final = { - "Authorization": f"Bearer {bearer_token}", - "Content-Type": "application/json", - } - - # Send the request - response = await self.async_httpx_client.post(url=api_endpoint, data=body.encode("utf-8"), headers=headers) - - if response.status_code not in [200, 204]: - verbose_logger.error( - "Azure Sentinel API error: status_code=%s, response=%s", - response.status_code, - response.text, - ) - raise Exception(f"Failed to send logs to Azure Sentinel: {response.status_code} - {response.text}") - - verbose_logger.debug( - "Azure Sentinel: Response from API status_code: %s", - response.status_code, + async def _send_batch(batch: Sequence[_QueuedPayload]): + body: Final = safe_dumps(batch) + return await self.async_httpx_client.post( + url=api_endpoint, + data=body.encode("utf-8"), + headers=headers, ) - except Exception as e: - verbose_logger.exception("Azure Sentinel Error sending batch API - %s\n%s", e, traceback.format_exc()) - finally: - log_queue.clear() + return await send_batch_with_413_split( + batch=log_queue, + send_batch=_send_batch, + exceeds_limits=lambda batch: ( + len(batch) > self.batch_size + or len(safe_dumps(batch).encode("utf-8")) > AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES + ), + success_status_codes=frozenset({200, 204}), + integration_name="Azure Sentinel", + drop_error_message="Azure Sentinel API Error - Payload too large for a single record", + non_success_handler=undelivered_after_http_error, + ) async def flush_queue(self): if self.flush_lock is None: diff --git a/litellm/integrations/batch_utils.py b/litellm/integrations/batch_utils.py new file mode 100644 index 00000000000..e1b48204e5d --- /dev/null +++ b/litellm/integrations/batch_utils.py @@ -0,0 +1,160 @@ +import asyncio +from collections.abc import Awaitable, Callable, Sequence +from typing import Final, Generic, TypeVar + +import httpx + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError + +_BatchItem = TypeVar("_BatchItem") + +_RETRYABLE_CLIENT_STATUS_CODES: Final = frozenset({408, 429}) + + +def is_retryable_status(status_code: int) -> bool: + return not 400 <= status_code < 500 or status_code in _RETRYABLE_CLIENT_STATUS_CODES + + +def undelivered_after_http_error( + batch: Sequence[_BatchItem], + status_code: int, + integration_name: str, + detail: str, +) -> tuple[_BatchItem, ...]: + """The records to requeue after a non-2xx: all of them on a status a retry can clear, none on + a 4xx that would only repeat, since retaining those retries a misconfiguration forever.""" + if is_retryable_status(status_code): + verbose_logger.error( + "%s API error: status_code=%s, will retry %s records - %s", + integration_name, + status_code, + len(batch), + detail, + ) + return tuple(batch) + verbose_logger.error( + "%s API error: status_code=%s is not retryable, dropped %s records - %s", + integration_name, + status_code, + len(batch), + detail, + ) + return () + + +def requeue_after_http_error( + batch: Sequence[_BatchItem], + status_code: int, + integration_name: str, + detail: str, +) -> tuple[_BatchItem, ...]: + verbose_logger.error( + "%s API error: status_code=%s, will retry %s records - %s", + integration_name, + status_code, + len(batch), + detail, + ) + return tuple(batch) + + +class BatchSendCancelled(asyncio.CancelledError, Generic[_BatchItem]): + """Cancellation of a batch send, carrying only the records the destination never accepted. + + A batch split under the size cap is delivered in pieces, so requeueing all of it after a + cancellation partway through would send the accepted pieces a second time. + """ + + def __init__(self, undelivered: tuple[_BatchItem, ...]) -> None: + super().__init__() + self.undelivered: Final = undelivered + + +async def _keep_the_remainder_on_cancel( + send: Awaitable[tuple[_BatchItem, ...]], + remainder: Sequence[_BatchItem], +) -> tuple[_BatchItem, ...]: + try: + return await send + except BatchSendCancelled as cancelled: + raise BatchSendCancelled((*cancelled.undelivered, *remainder)) from cancelled + + +async def send_batch_with_413_split( + batch: Sequence[_BatchItem], + send_batch: Callable[[Sequence[_BatchItem]], Awaitable[httpx.Response]], + exceeds_limits: Callable[[Sequence[_BatchItem]], bool], + success_status_codes: frozenset[int], + integration_name: str, + drop_error_message: str, + non_success_handler: Callable[ + [Sequence[_BatchItem], int, str, str], tuple[_BatchItem, ...] + ] = requeue_after_http_error, +) -> tuple[_BatchItem, ...]: + async def _halve() -> tuple[_BatchItem, ...]: + midpoint: Final = len(batch) // 2 + left_batch: Final = batch[:midpoint] + right_batch: Final = batch[midpoint:] + left_undelivered: Final = await _keep_the_remainder_on_cancel( + send_batch_with_413_split( + batch=left_batch, + send_batch=send_batch, + exceeds_limits=exceeds_limits, + success_status_codes=success_status_codes, + integration_name=integration_name, + drop_error_message=drop_error_message, + non_success_handler=non_success_handler, + ), + right_batch, + ) + if left_undelivered: + return (*left_undelivered, *right_batch) + return await send_batch_with_413_split( + batch=right_batch, + send_batch=send_batch, + exceeds_limits=exceeds_limits, + success_status_codes=success_status_codes, + integration_name=integration_name, + drop_error_message=drop_error_message, + non_success_handler=non_success_handler, + ) + + async def _handle_413() -> tuple[_BatchItem, ...]: + if len(batch) == 1: + verbose_logger.error(drop_error_message) + return () + return await _halve() + + if not batch: + return () + + try: + oversized: Final = exceeds_limits(batch) + except Exception as e: # noqa: BLE001 # any record that cannot be serialized is isolated and dropped alone + if len(batch) > 1: + return await _halve() + verbose_logger.exception("%s dropped a record that cannot be serialized - %s", integration_name, e) + return () + if oversized and len(batch) > 1: + return await _halve() + + try: + response: Final = await send_batch(batch) + except MaskedHTTPStatusError as e: + if e.status_code == 413: + return await _handle_413() + return non_success_handler(batch, e.status_code, integration_name, str(e)) + except asyncio.CancelledError as cancelled: + raise BatchSendCancelled(tuple(batch)) from cancelled + except Exception as e: + verbose_logger.exception("%s Error sending batch API - %s", integration_name, e) + return tuple(batch) + + if response.status_code == 413: + return await _handle_413() + if response.status_code not in success_status_codes: + return non_success_handler(batch, response.status_code, integration_name, response.text) + + verbose_logger.debug("%s delivered %s records, status_code=%s", integration_name, len(batch), response.status_code) + return () diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 866076a3c49..77b12d1e3fa 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -29,6 +29,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.integrations.batch_utils import BatchSendCancelled, requeue_after_http_error, send_batch_with_413_split from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.datadog.datadog_handler import ( get_datadog_base_url_from_env, @@ -43,7 +44,6 @@ from litellm.integrations.datadog.datadog_mock_client import ( ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.llms.custom_httpx.http_handler import ( - MaskedHTTPStatusError, _get_httpx_client, get_async_httpx_client, httpxSpecialProvider, @@ -396,6 +396,9 @@ class DataDogLogger( if self.is_mock_mode: verbose_logger.debug("[DATADOG MOCK] Batch of %s events successfully mocked", len(batch_to_send)) + except BatchSendCancelled as cancelled: + self.log_queue = list(cancelled.undelivered) + self.log_queue # mutable-ok: logger queue remains appendable + raise asyncio.CancelledError() from cancelled except Exception as e: self.log_queue = batch_to_send + self.log_queue verbose_logger.exception("Datadog Error sending batch API - %s\n%s", e, traceback.format_exc()) @@ -413,53 +416,16 @@ class DataDogLogger( that could not be delivered because of a non-413 (transient) error, so the caller re-queues only those and never the events already accepted by Datadog. """ - pending: Final[list[list]] = [batch] - while pending: - chunk = pending.pop() - if not chunk: - continue - if len(chunk) > 1 and self._exceeds_intake_limits(chunk): - mid = len(chunk) // 2 - pending.append(chunk[mid:]) - pending.append(chunk[:mid]) - continue - try: - response = await self.async_send_compressed_data(chunk) - except Exception as e: - if isinstance(e, MaskedHTTPStatusError) and e.status_code == 413: - response = e.response - else: - verbose_logger.exception("Datadog Error sending batch API - %s", e) - return self._undelivered(chunk, pending) - - if response.status_code == 413: - if len(chunk) == 1: - verbose_logger.error(DD_ERRORS.DATADOG_413_ERROR.value) - continue - mid = len(chunk) // 2 - pending.append(chunk[mid:]) - pending.append(chunk[:mid]) - continue - - if response.status_code != 202: - verbose_logger.error( - "Datadog: unexpected response status_code=%s, text=%s", - response.status_code, - response.text, - ) - return self._undelivered(chunk, pending) - - verbose_logger.debug( - "Datadog: delivered %s events, status_code=%s, text=%s", - len(chunk), - response.status_code, - response.text, - ) - return [] - - @staticmethod - def _undelivered(chunk: list, pending: list[list]) -> list: - return chunk + [event for remaining in reversed(pending) for event in remaining] + undelivered: Final = await send_batch_with_413_split( + batch=batch, + send_batch=self.async_send_compressed_data, + exceeds_limits=self._exceeds_intake_limits, + success_status_codes=frozenset({202}), + integration_name="Datadog", + drop_error_message=DD_ERRORS.DATADOG_413_ERROR.value, + non_success_handler=requeue_after_http_error, + ) + return list(undelivered) # mutable-ok: caller prepends records to the logger queue @staticmethod def _exceeds_intake_limits(chunk: Sequence[DatadogPayload]) -> bool: @@ -606,7 +572,7 @@ class DataDogLogger( ) return dd_payload - async def async_send_compressed_data(self, data: list) -> Response: + async def async_send_compressed_data(self, data: Sequence[DatadogPayload]) -> Response: """ Async helper to send compressed data to datadog self.intake_url diff --git a/litellm/types/integrations/azure_sentinel.py b/litellm/types/integrations/azure_sentinel.py index c7e04ded2bf..d83dd331e60 100644 --- a/litellm/types/integrations/azure_sentinel.py +++ b/litellm/types/integrations/azure_sentinel.py @@ -1,5 +1,9 @@ +from typing import Final + from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams +AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES: Final = 1_000_000 + class AzureSentinelInitParams(StandardCustomLoggerInitParams): """ diff --git a/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py b/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py index f645379a4f4..e2707d321bf 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py @@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, Mock, patch import httpx import pytest from httpx import Request, Response +from pydantic import BaseModel, computed_field from litellm.integrations.datadog.datadog import DataDogLogger from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError @@ -31,9 +32,7 @@ def _payloads(n, message=None): def _raised_413(): request = Request("POST", "https://example.com") response = Response(413, request=request, text="Payload Too Large") - return MaskedHTTPStatusError( - httpx.HTTPStatusError("413", request=request, response=response) - ) + return MaskedHTTPStatusError(httpx.HTTPStatusError("413", request=request, response=response)) def _make_send(max_ok, delivered, *, raise_413=True): @@ -85,9 +84,7 @@ async def test_async_send_batch_keeps_events_appended_during_send(datadog_env): status="info", ) ) - return Response( - 202, request=Request("POST", "https://example.com"), text="Accepted" - ) + return Response(202, request=Request("POST", "https://example.com"), text="Accepted") logger.async_send_compressed_data = AsyncMock(side_effect=_mock_send) @@ -172,9 +169,7 @@ async def test_413_returned_response_also_splits(datadog_env): logger.log_queue = _payloads(4) delivered: list = [] - logger.async_send_compressed_data = AsyncMock( - side_effect=_make_send(1, delivered, raise_413=False) - ) + logger.async_send_compressed_data = AsyncMock(side_effect=_make_send(1, delivered, raise_413=False)) await logger.async_send_batch() @@ -186,9 +181,7 @@ def _make_recording_send(sent_batches, delivered): async def _send(data): sent_batches.append(list(data)) delivered.extend(data) - return Response( - 202, request=Request("POST", "https://example.com"), text="Accepted" - ) + return Response(202, request=Request("POST", "https://example.com"), text="Accepted") return _send @@ -206,18 +199,13 @@ async def test_oversized_payload_splits_before_any_send(datadog_env): logger.log_queue = list(events) sent_batches: list = [] delivered: list = [] - logger.async_send_compressed_data = AsyncMock( - side_effect=_make_recording_send(sent_batches, delivered) - ) + logger.async_send_compressed_data = AsyncMock(side_effect=_make_recording_send(sent_batches, delivered)) await logger.async_send_batch() assert delivered == events assert len(sent_batches) == 3 - assert all( - len(safe_dumps(batch).encode("utf-8")) <= DD_MAX_PAYLOAD_SIZE_BYTES - for batch in sent_batches - ) + assert all(len(safe_dumps(batch).encode("utf-8")) <= DD_MAX_PAYLOAD_SIZE_BYTES for batch in sent_batches) assert logger.log_queue == [] @@ -232,9 +220,7 @@ async def test_batch_over_max_event_count_splits_before_any_send(datadog_env): logger.log_queue = list(events) sent_batches: list = [] delivered: list = [] - logger.async_send_compressed_data = AsyncMock( - side_effect=_make_recording_send(sent_batches, delivered) - ) + logger.async_send_compressed_data = AsyncMock(side_effect=_make_recording_send(sent_batches, delivered)) await logger.async_send_batch() @@ -281,9 +267,7 @@ async def test_partial_delivery_then_transient_error_requeues_only_undelivered( if messages == ['{"event": 2}', '{"event": 3}']: raise RuntimeError("transient network error") delivered.extend(messages) - return Response( - 202, request=Request("POST", "https://example.com"), text="Accepted" - ) + return Response(202, request=Request("POST", "https://example.com"), text="Accepted") logger.async_send_compressed_data = AsyncMock(side_effect=_send) @@ -304,9 +288,7 @@ async def test_unexpected_non_202_status_requeues(datadog_env): logger.log_queue = _payloads(2) logger.async_send_compressed_data = AsyncMock( - return_value=Response( - 200, request=Request("POST", "https://example.com"), text="OK" - ) + return_value=Response(200, request=Request("POST", "https://example.com"), text="OK") ) await logger.async_send_batch() @@ -502,3 +484,77 @@ async def test_flush_queue_returns_without_lock(datadog_env): await logger.flush_queue() logger.async_send_batch.assert_not_awaited() + + +class _RaisesWhileDumping(BaseModel): + @computed_field + @property + def rendered(self) -> str: + raise RuntimeError("this field cannot be rendered") + + +@pytest.mark.asyncio +async def test_event_whose_serialization_raises_is_dropped_alone(datadog_env): + """safe_dumps hands pydantic models to model_dump, so serialization can raise any exception + class. The intake-limit probe has to isolate that one event and drop it, not fail the whole + batch back onto the queue where it would poison every later flush.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(4) + logger.log_queue[1]["message"] = _RaisesWhileDumping() + delivered: list = [] + logger.async_send_compressed_data = AsyncMock(side_effect=_make_send(DD_MAX_BATCH_SIZE, delivered)) + + await logger.async_send_batch() + + assert delivered == ['{"event": 0}', '{"event": 2}', '{"event": 3}'] + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_cancellation_mid_split_requeues_only_the_undelivered_events(datadog_env): + """A cancelled split must keep the pieces Datadog never accepted, without resending the piece + it did, and must surface as a plain CancelledError so asyncio.wait_for still reads it as a + timeout on Python 3.12.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(4) + attempts: list = [] + + async def _send(data): + if len(data) > 2: + raise _raised_413() + attempts.append([event["message"] for event in data]) + if len(attempts) > 1: + raise asyncio.CancelledError + return Response(202, request=Request("POST", "https://example.com"), text="Accepted") + + logger.async_send_compressed_data = AsyncMock(side_effect=_send) + + with pytest.raises(asyncio.CancelledError) as excinfo: + await logger.async_send_batch() + + assert type(excinfo.value) is asyncio.CancelledError + assert attempts == [['{"event": 0}', '{"event": 1}'], ['{"event": 2}', '{"event": 3}']] + assert [event["message"] for event in logger.log_queue] == ['{"event": 2}', '{"event": 3}'] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [400, 403, 429, 500, 503]) +async def test_raised_intake_error_preserves_datadog_requeue_behavior(datadog_env, status_code): + """Datadog requeues every non-413 HTTP failure so a corrected key or endpoint can recover telemetry.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(2) + request = Request("POST", "https://example.com") + response = Response(status_code, request=request, text="rejected") + logger.async_send_compressed_data = AsyncMock( + side_effect=MaskedHTTPStatusError(httpx.HTTPStatusError(str(status_code), request=request, response=response)) + ) + + await logger.async_send_batch() + + assert [event["message"] for event in logger.log_queue] == ['{"event": 0}', '{"event": 1}'] diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 7335316548d..038197c06c5 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -2,18 +2,23 @@ Test Azure Sentinel logging integration """ +import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +from httpx import Request, Response +from pydantic import BaseModel, computed_field from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger +from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError +from litellm.types.integrations.azure_sentinel import AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload def _close_periodic_flush_task(coro): coro.close() - return None @pytest.mark.asyncio @@ -414,3 +419,837 @@ def test_azure_sentinel_authority_host_argument_outranks_the_scoped_env_var(_no_ assert logger.authority_host == "https://login.microsoftonline.com" assert logger.oauth_scope == "https://monitor.azure.com/.default" + + +def _standard_payloads(count, filler_bytes=0): + return [ + StandardLoggingPayload( + id=f"standard-{i}", + call_type="completion", + model="gpt-3.5-turbo", + status="success", + messages=[{"role": "user", "content": "x" * filler_bytes}], + response={"choices": [{"message": {"content": "Hi"}}]}, + ) + for i in range(count) + ] + + +def _audit_payloads(count, filler_bytes=0): + return [ + StandardAuditLogPayload( + id=f"audit-{i}", + updated_at="2026-05-06T04:39:00+00:00", + changed_by="user-1", + changed_by_api_key="sk-test", + action="created", + table_name="LiteLLM_TeamTable", + object_id="team-1", + before_value=None, + updated_values=json.dumps({"team_alias": "x" * filler_bytes}), + ) + for i in range(count) + ] + + +QUEUE_CASES = [ + pytest.param("log_queue", "async_send_batch", _standard_payloads, id="standard"), + pytest.param("audit_log_queue", "async_send_audit_batch", _audit_payloads, id="audit"), +] + + +def _token_response(): + response = MagicMock() + response.status_code = 200 + response.json = MagicMock(return_value={"access_token": "test-bearer-token", "expires_in": 3600}) + response.text = "Success" + return response + + +def _install_ingestion(logger, on_ingest): + """Route the OAuth call to a canned token and every ingestion call to `on_ingest(body_bytes)`.""" + + async def _post(*args, **kwargs): + if "oauth2/v2.0/token" in kwargs.get("url", ""): + return _token_response() + return await on_ingest(kwargs["data"]) + + logger.async_httpx_client.post = AsyncMock(side_effect=_post) + + +def _accepted(): + return Response(204, request=Request("POST", "https://example.com"), text="") + + +def _too_large(*, raised): + request = Request("POST", "https://example.com") + response = Response(413, request=request, text="Payload Too Large") + if raised: + raise MaskedHTTPStatusError(httpx.HTTPStatusError("413", request=request, response=response)) + return response + + +def _rejected(status_code, *, raised): + """litellm's http handler calls raise_for_status, so a real rejection arrives raised, not returned.""" + request = Request("POST", "https://example.com") + response = Response(status_code, request=request, text=f"rejected with {status_code}") + if raised: + raise MaskedHTTPStatusError(httpx.HTTPStatusError(str(status_code), request=request, response=response)) + return response + + +def _awaiting_retry(logger, queue_attr): + return getattr(logger, "logs_awaiting_retry" if queue_attr == "log_queue" else "audit_logs_awaiting_retry") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_splits_a_batch_that_would_exceed_the_ingestion_cap( + queue_attr, send_method, build_payloads +): + """Azure Monitor rejects a body over 1MB uncompressed, so an oversize batch has to be split + before it is sent instead of being posted whole and lost.""" + logger = _build_logger() + records = build_payloads(4, filler_bytes=400_000) + setattr(logger, queue_attr, list(records)) + + sent_bodies = [] + + async def _on_ingest(data): + sent_bodies.append(data) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert len(sent_bodies) > 1 + assert all(len(body) <= AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES for body in sent_bodies) + delivered = [record["id"] for body in sent_bodies for record in json.loads(body.decode("utf-8"))] + assert delivered == [record["id"] for record in records] + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raised", [True, False], ids=["raised", "returned"]) +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_halves_the_batch_on_413(queue_attr, send_method, build_payloads, raised): + """A 413 the size estimate did not predict must halve the batch and retry, not drop it. + + litellm's http handler raises MaskedHTTPStatusError on a 4xx, so the raised path is the one + a real Azure Monitor 413 takes, and both are covered here. + """ + logger = _build_logger() + records = build_payloads(4) + setattr(logger, queue_attr, list(records)) + + delivered = [] + + async def _on_ingest(data): + body = json.loads(data.decode("utf-8")) + if len(body) > 1: + return _too_large(raised=raised) + delivered.extend(record["id"] for record in body) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert delivered == [record["id"] for record in records] + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_drops_only_the_lone_record_that_still_413s(queue_attr, send_method, build_payloads): + """One undeliverable record must not take its siblings down with it or wedge the queue.""" + logger = _build_logger() + records = build_payloads(4) + poison = records[2]["id"] + setattr(logger, queue_attr, list(records)) + + delivered = [] + + async def _on_ingest(data): + body = json.loads(data.decode("utf-8")) + if any(record["id"] == poison for record in body): + return _too_large(raised=True) + delivered.extend(record["id"] for record in body) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await asyncio.wait_for(getattr(logger, send_method)(), timeout=10) + + assert delivered == [record["id"] for record in records if record["id"] != poison] + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_requeues_only_what_a_transient_failure_left_undelivered( + queue_attr, send_method, build_payloads +): + """Records Azure Monitor already accepted must not be sent twice, and the rest must survive + for the next flush instead of being cleared.""" + logger = _build_logger() + records = build_payloads(4) + setattr(logger, queue_attr, list(records)) + + delivered = [] + + async def _on_ingest(data): + body = json.loads(data.decode("utf-8")) + if len(body) > 2: + return _too_large(raised=True) + if any(record["id"] == records[2]["id"] for record in body): + raise httpx.ConnectError("connection reset") + delivered.extend(record["id"] for record in body) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert delivered == [records[0]["id"], records[1]["id"]] + assert getattr(logger, queue_attr) == records[2:] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_requeues_the_batch_on_a_non_success_status(queue_attr, send_method, build_payloads): + """A 500 from ingestion is retryable, so the batch has to stay queued.""" + logger = _build_logger() + records = build_payloads(3) + setattr(logger, queue_attr, list(records)) + + async def _on_ingest(data): + return Response(500, request=Request("POST", "https://example.com"), text="Internal Server Error") + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert getattr(logger, queue_attr) == records + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_requeues_the_batch_when_the_oauth_token_call_fails( + queue_attr, send_method, build_payloads +): + """Losing the token is transient, so the batch must not be dropped on the way to the wire.""" + logger = _build_logger() + records = build_payloads(2) + setattr(logger, queue_attr, list(records)) + + ingestion_calls = [] + + async def _post(*args, **kwargs): + if "oauth2/v2.0/token" in kwargs.get("url", ""): + failed = MagicMock() + failed.status_code = 401 + failed.text = "Unauthorized" + return failed + ingestion_calls.append(kwargs["url"]) + return _accepted() + + logger.async_httpx_client.post = AsyncMock(side_effect=_post) + + await getattr(logger, send_method)() + + assert ingestion_calls == [] + assert getattr(logger, queue_attr) == records + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_caps_the_retry_queue_at_max_queue_size(queue_attr, send_method, build_payloads): + """Retrying forever against an unreachable workspace must not grow the queue without bound, + so the oldest records go once the queue is over its limit.""" + logger = _build_logger(max_queue_size=3) + records = build_payloads(4) + setattr(logger, queue_attr, list(records)) + + async def _on_ingest(data): + raise httpx.ConnectError("connection reset") + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert getattr(logger, queue_attr) == records[1:] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_keeps_records_queued_during_a_send(queue_attr, send_method, build_payloads): + """The queue is detached before sending, so a record logged mid-flush is kept and lands behind + anything the failed send hands back.""" + logger = _build_logger() + records = build_payloads(2) + late_record = build_payloads(1)[0] + late_record["id"] = "logged-during-send" + setattr(logger, queue_attr, list(records)) + + async def _on_ingest(data): + getattr(logger, queue_attr).append(late_record) + raise httpx.ConnectError("connection reset") + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert getattr(logger, queue_attr) == [*records, late_record] + + +def _poison(record): + """A mixed-type set makes safe_dumps raise TypeError while sorting it, so the record can never be serialized.""" + field = "messages" if "messages" in record else "updated_values" + record[field] = {1, "a"} + return record + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_drops_only_the_record_that_cannot_be_serialized(queue_attr, send_method, build_payloads): + """A record that raises during serialization used to escape the send, which killed the periodic + flush task for good and lost the already-detached batch with it. It has to be isolated and + dropped alone, with the flush completing normally.""" + logger = _build_logger() + records = build_payloads(4) + poison = _poison(records[2])["id"] + setattr(logger, queue_attr, list(records)) + + delivered = [] + + async def _on_ingest(data): + delivered.extend(record["id"] for record in json.loads(data.decode("utf-8"))) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await asyncio.wait_for(logger.flush_queue(), timeout=10) + + assert delivered == [record["id"] for record in records if record["id"] != poison] + assert getattr(logger, queue_attr) == [] + + +async def _log(logger, queue_attr, record): + if queue_attr == "log_queue": + await logger.async_log_success_event({"standard_logging_object": record}, None, None, None) + return + await logger.async_log_audit_log_event(record) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_retries_on_the_flush_timer_not_on_every_record_while_the_destination_is_down( + queue_attr, send_method, build_payloads +): + """Requeued records keep the queue at or over batch_size, so without a guard every new record + re-sent the whole growing queue. While a retry is pending only the periodic flush may send, and + a successful flush hands the trigger back to the batch size.""" + logger = _build_logger(batch_size=3) + records = build_payloads(11) + + attempts = [] + destination_down = True + + async def _on_ingest(data): + attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))]) + if destination_down: + raise httpx.ConnectError("connection reset") + return _accepted() + + _install_ingestion(logger, _on_ingest) + + for record in records[:8]: + await _log(logger, queue_attr, record) + + assert attempts == [[record["id"] for record in records[:3]]] + assert getattr(logger, queue_attr) == records[:8] + + destination_down = False + await logger.flush_queue() + for record in records[8:]: + await _log(logger, queue_attr, record) + + assert [record_id for attempt in attempts[1:-1] for record_id in attempt] == [record["id"] for record in records[:8]] + assert all(len(attempt) <= 3 for attempt in attempts[1:-1]) + assert attempts[-1] == [record["id"] for record in records[8:]] + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_threshold_send_waits_for_an_in_flight_timer_flush( + queue_attr, send_method, build_payloads +): + """A batch-size send that overlapped the periodic flush could finish after it and requeue its + newer records in front of the older ones, so the max_queue_size trim would then drop the + newest records instead of the oldest. Both paths have to take the flush lock, and a waiter + that gets the lock after a failed flush stands down instead of resending the whole queue.""" + logger = _build_logger(batch_size=2) + records = build_payloads(4) + setattr(logger, queue_attr, list(records[:2])) + + attempts = [] + timer_send_started = asyncio.Event() + release_timer_send = asyncio.Event() + + async def _on_ingest(data): + attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))]) + if len(attempts) == 1: + timer_send_started.set() + await release_timer_send.wait() + raise httpx.ConnectError("connection reset") + + _install_ingestion(logger, _on_ingest) + + timer_flush = asyncio.create_task(logger.flush_queue()) + await asyncio.wait_for(timer_send_started.wait(), timeout=10) + await _log(logger, queue_attr, records[2]) + threshold_send = asyncio.create_task(_log(logger, queue_attr, records[3])) + await asyncio.sleep(0) + release_timer_send.set() + await asyncio.wait_for(timer_flush, timeout=10) + await asyncio.wait_for(threshold_send, timeout=10) + + assert attempts == [[record["id"] for record in records[:2]]] + assert getattr(logger, queue_attr) == records + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_concurrent_threshold_sends_collapse_into_one_attempt_while_the_destination_is_down( + queue_attr, send_method, build_payloads +): + """Records logged while a threshold send is blocked on the wire all see the retry flag still + unset and queue up on the flush lock. Each waiter has to recheck under the lock, or every one + of them resends the growing queue as soon as the first attempt fails.""" + logger = _build_logger(batch_size=2) + records = build_payloads(6) + + attempts = [] + first_send_started = asyncio.Event() + release_first_send = asyncio.Event() + destination_down = True + + async def _on_ingest(data): + attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))]) + if len(attempts) == 1: + first_send_started.set() + await release_first_send.wait() + if destination_down: + raise httpx.ConnectError("connection reset") + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await _log(logger, queue_attr, records[0]) + first_send = asyncio.create_task(_log(logger, queue_attr, records[1])) + await asyncio.wait_for(first_send_started.wait(), timeout=10) + waiters = [asyncio.create_task(_log(logger, queue_attr, record)) for record in records[2:]] + await asyncio.sleep(0) + release_first_send.set() + await asyncio.wait_for(asyncio.gather(first_send, *waiters), timeout=10) + + assert attempts == [[record["id"] for record in records[:2]]] + assert getattr(logger, queue_attr) == records + + destination_down = False + await logger.flush_queue() + + assert [record_id for attempt in attempts[1:] for record_id in attempt] == [record["id"] for record in records] + assert all(len(attempt) <= 2 for attempt in attempts[1:]) + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_requeues_a_cancelled_send( + queue_attr, send_method, build_payloads +): + """Cancellation after detaching a batch must preserve the detached records for a later flush.""" + logger = _build_logger() + records = build_payloads(2) + setattr(logger, queue_attr, list(records)) + + async def _on_ingest(data): + raise asyncio.CancelledError + + _install_ingestion(logger, _on_ingest) + + with pytest.raises(asyncio.CancelledError) as excinfo: + await getattr(logger, send_method)() + + assert type(excinfo.value) is asyncio.CancelledError + assert getattr(logger, queue_attr) == records + assert _awaiting_retry(logger, queue_attr) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_requeues_a_send_cancelled_before_it_reached_the_wire( + queue_attr, send_method, build_payloads +): + """Cancellation can land on the token call, before any record was sent, and the detached batch + has to survive that too.""" + logger = _build_logger() + records = build_payloads(2) + setattr(logger, queue_attr, list(records)) + logger.async_httpx_client.post = AsyncMock(side_effect=asyncio.CancelledError) + + with pytest.raises(asyncio.CancelledError) as excinfo: + await getattr(logger, send_method)() + + assert type(excinfo.value) is asyncio.CancelledError + assert getattr(logger, queue_attr) == records + assert _awaiting_retry(logger, queue_attr) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_does_not_resend_the_half_delivered_before_a_cancelled_split( + queue_attr, send_method, build_payloads +): + """A batch over the size cap goes out in pieces, so a cancellation partway through must requeue + only the pieces the destination never accepted, or the accepted ones land in Sentinel twice.""" + logger = _build_logger() + records = build_payloads(8, filler_bytes=400_000) + setattr(logger, queue_attr, list(records)) + + attempts = [] + cancel_after_the_first_piece = True + + async def _on_ingest(data): + attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))]) + if cancel_after_the_first_piece and len(attempts) > 1: + raise asyncio.CancelledError + return _accepted() + + _install_ingestion(logger, _on_ingest) + + with pytest.raises(asyncio.CancelledError) as excinfo: + await getattr(logger, send_method)() + + assert type(excinfo.value) is asyncio.CancelledError + assert attempts == [[record["id"] for record in records[:2]], [record["id"] for record in records[2:4]]] + assert getattr(logger, queue_attr) == records[2:] + + cancel_after_the_first_piece = False + await logger.flush_queue() + + assert [record_id for attempt in attempts[2:] for record_id in attempt] == [record["id"] for record in records[2:]] + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_threshold_waiter_does_not_send_a_sub_batch_after_success( + queue_attr, send_method, build_payloads +): + """A successful threshold send can leave one record behind, so a waiter must not send it + before the next record completes a batch.""" + logger = _build_logger(batch_size=2) + records = build_payloads(3) + + attempts = [] + first_send_started = asyncio.Event() + release_first_send = asyncio.Event() + + async def _on_ingest(data): + attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))]) + if len(attempts) == 1: + first_send_started.set() + await release_first_send.wait() + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await _log(logger, queue_attr, records[0]) + first_send = asyncio.create_task(_log(logger, queue_attr, records[1])) + await asyncio.wait_for(first_send_started.wait(), timeout=10) + waiter = asyncio.create_task(_log(logger, queue_attr, records[2])) + await asyncio.sleep(0) + release_first_send.set() + await asyncio.wait_for(asyncio.gather(first_send, waiter), timeout=10) + + assert attempts == [[record["id"] for record in records[:2]]] + assert getattr(logger, queue_attr) == [records[2]] + + +@pytest.mark.asyncio +async def test_azure_sentinel_threshold_send_only_sends_the_queue_that_crossed_the_threshold(): + """The standard and audit queues retry independently: crossing the audit threshold must not + resend standard records that are waiting for the periodic flush.""" + logger = _build_logger(batch_size=2) + standard_records = _standard_payloads(2) + audit_records = _audit_payloads(2) + logger.log_queue = list(standard_records) + logger.logs_awaiting_retry = True + + attempts = [] + + async def _on_ingest(data): + attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))]) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + for record in audit_records: + await logger.async_log_audit_log_event(record) + + assert attempts == [[record["id"] for record in audit_records]] + assert logger.audit_log_queue == [] + assert logger.log_queue == standard_records + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [408, 429, 500, 503]) +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_keeps_the_batch_when_ingestion_raises_a_retryable_status( + queue_attr, send_method, build_payloads, status_code +): + """A 5xx, a timeout or a throttle can clear on the next flush, so the whole batch stays queued + and the awaiting-retry flag hands the send back to the timer.""" + logger = _build_logger() + records = build_payloads(3) + setattr(logger, queue_attr, list(records)) + + async def _on_ingest(data): + return _rejected(status_code, raised=True) + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert getattr(logger, queue_attr) == records + assert _awaiting_retry(logger, queue_attr) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raised", [True, False], ids=["raised", "returned"]) +@pytest.mark.parametrize("status_code", [400, 403, 404]) +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_drops_the_batch_when_ingestion_rejects_it_for_good( + queue_attr, send_method, build_payloads, status_code, raised +): + """A permanent 4xx is dropped, the flag is cleared and the next records go out on their own.""" + logger = _build_logger(batch_size=2) + rejected_records = build_payloads(2) + later_records = build_payloads(4)[2:] + setattr(logger, queue_attr, list(rejected_records)) + + delivered = [] + destination_rejects = True + + async def _on_ingest(data): + if destination_rejects: + return _rejected(status_code, raised=raised) + delivered.extend(record["id"] for record in json.loads(data.decode("utf-8"))) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert getattr(logger, queue_attr) == [] + assert not _awaiting_retry(logger, queue_attr) + + destination_rejects = False + for record in later_records: + await _log(logger, queue_attr, record) + + assert delivered == [record["id"] for record in later_records] + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_keeps_the_whole_batch_when_the_first_piece_of_a_split_fails( + queue_attr, send_method, build_payloads +): + """When the first half of a split hits a retryable error the untried second half must be kept + too, in the original order, instead of being sent ahead of records that are still pending.""" + logger = _build_logger() + records = build_payloads(4) + setattr(logger, queue_attr, list(records)) + + attempts = [] + + async def _on_ingest(data): + body = json.loads(data.decode("utf-8")) + attempts.append([record["id"] for record in body]) + if len(body) > 2: + return _too_large(raised=True) + return _rejected(503, raised=True) + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert attempts == [[record["id"] for record in records], [record["id"] for record in records[:2]]] + assert getattr(logger, queue_attr) == records + assert _awaiting_retry(logger, queue_attr) + + +class _RaisesWhileDumping(BaseModel): + @computed_field + @property + def rendered(self) -> str: + raise RuntimeError("this field cannot be rendered") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_drops_only_the_record_whose_serialization_raises_an_unexpected_error( + queue_attr, send_method, build_payloads +): + """Serialization can fail with any exception class, not just TypeError or ValueError, because + safe_dumps hands pydantic models to model_dump. A record that raises anything has to be isolated + and dropped alone, or the flush dies with the whole batch.""" + logger = _build_logger() + records = build_payloads(4) + poison = records[1] + poison["messages" if "messages" in poison else "updated_values"] = _RaisesWhileDumping() + setattr(logger, queue_attr, list(records)) + + delivered = [] + + async def _on_ingest(data): + delivered.extend(record["id"] for record in json.loads(data.decode("utf-8"))) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await asyncio.wait_for(logger.flush_queue(), timeout=10) + + assert delivered == [record["id"] for record in records if record is not poison] + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_send_cancelled_by_a_timeout_surfaces_as_a_timeout( + queue_attr, send_method, build_payloads +): + """The logging worker bounds each flush with asyncio.wait_for, which on Python 3.12 only turns + an exact CancelledError into TimeoutError. A subclass carrying the undelivered records would + escape the worker as an unhandled error, so the send must re-raise the plain class.""" + logger = _build_logger() + records = build_payloads(2) + setattr(logger, queue_attr, list(records)) + + async def _on_ingest(data): + await asyncio.sleep(60) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(getattr(logger, send_method)(), timeout=0.05) + + assert getattr(logger, queue_attr) == records + assert _awaiting_retry(logger, queue_attr) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_never_sends_more_than_batch_size_records_in_one_request( + queue_attr, send_method, build_payloads +): + """A recovery flush can find far more than batch_size records queued. Splitting on the count + first keeps each request at the configured size and bounds how much of the queue is serialized + just to measure it.""" + logger = _build_logger(batch_size=2) + records = build_payloads(5) + setattr(logger, queue_attr, list(records)) + + attempts = [] + + async def _on_ingest(data): + attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))]) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert attempts == [ + [records[0]["id"], records[1]["id"]], + [records[2]["id"]], + [records[3]["id"], records[4]["id"]], + ] + assert getattr(logger, queue_attr) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status_code, expected_queue", + [pytest.param(503, "kept", id="503-kept"), pytest.param(401, "dropped", id="401-dropped")], +) +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_oauth_rejection_follows_the_same_retry_rule_as_ingestion( + queue_attr, send_method, build_payloads, status_code, expected_queue +): + """The token endpoint raises through the same http handler as ingestion. A 5xx there is + transient and keeps the batch, a 401 means the client secret is wrong and would fail every + retry, so the batch is dropped instead of wedging the queue.""" + logger = _build_logger() + records = build_payloads(2) + setattr(logger, queue_attr, list(records)) + + ingestion_calls = [] + + async def _post(*args, **kwargs): + if "oauth2/v2.0/token" in kwargs.get("url", ""): + return _rejected(status_code, raised=True) + ingestion_calls.append(kwargs["url"]) + return _accepted() + + logger.async_httpx_client.post = AsyncMock(side_effect=_post) + + await getattr(logger, send_method)() + + assert ingestion_calls == [] + assert getattr(logger, queue_attr) == (records if expected_queue == "kept" else []) + assert _awaiting_retry(logger, queue_attr) is (expected_queue == "kept") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES) +async def test_azure_sentinel_does_not_stay_in_retry_mode_when_the_queue_cap_trims_everything( + queue_attr, send_method, build_payloads +): + """With max_queue_size at 0 the cap drops every requeued record, so there is nothing for the + timer to retry. The flag must follow the retained queue, or every later threshold send is + skipped until the timer happens to fire.""" + logger = _build_logger(batch_size=2, max_queue_size=0) + lost_records = build_payloads(2) + later_records = build_payloads(4)[2:] + setattr(logger, queue_attr, list(lost_records)) + + delivered = [] + destination_down = True + + async def _on_ingest(data): + if destination_down: + raise httpx.ConnectError("connection reset") + delivered.extend(record["id"] for record in json.loads(data.decode("utf-8"))) + return _accepted() + + _install_ingestion(logger, _on_ingest) + + await getattr(logger, send_method)() + + assert getattr(logger, queue_attr) == [] + assert not _awaiting_retry(logger, queue_attr) + + destination_down = False + for record in later_records: + await _log(logger, queue_attr, record) + + assert delivered == [record["id"] for record in later_records]