scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(azure_sentinel): split batches under the 1MB ingestion cap (#39880)

Browse source 
* fix(azure_sentinel): split batches under the 1MB ingestion cap and keep undelivered records queued

Azure Monitor rejects any Logs Ingestion body over 1MB with a 413. The Sentinel logger
posted the whole queue as one body and cleared it in a finally block, so an oversize
batch, a transient 5xx, or a failed token call dropped every queued record, and records
logged while a send was in flight were cleared with it. Both the standard and the audit
queue share the sender.

Move Datadog's proactive size split and 413 halving into a shared helper,
litellm/integrations/batch_utils.send_batch_with_413_split, and route Sentinel through it
with a 1MB size check. A lone record that still 413s is dropped, everything a transient
failure leaves undelivered goes back to the front of its queue, and the retry queue is
capped at max_queue_size so an unreachable workspace cannot grow memory without bound

* fix(azure_sentinel): retry undelivered records on the flush timer only

Requeued records made every later event cross the batch_size threshold, so a
down ingestion endpoint got one full-queue resend per request. Threshold sends
now go through flush_queue, so they take the flush lock instead of racing the
timer, and they stand down while records are awaiting retry.

A record that cannot be serialized raised out of the size probe and killed the
periodic flush task. The probe now runs inside the failure handling, so the
batch is split and only the record that cannot be serialized is dropped.

* fix(azure_sentinel): decide threshold sends under the flush lock

Concurrent callbacks all read logs_awaiting_retry before the first send
finished, so each one resent the whole queue once that send failed. The
flag and the batch_size threshold are now rechecked while holding the
flush lock, and each queue sends only itself instead of going through
flush_queue, which was retrying the other queue too.

* test(azure_sentinel): cover successful threshold waiters

* fix(azure_sentinel): preserve cancelled batches for retry

* fix(azure_sentinel): requeue only the undelivered part of a cancelled split

A batch over the ingestion cap goes out in pieces, so a cancellation partway
through requeued pieces the destination had already accepted and sent them a
second time on the next flush

The split helper now raises a cancellation carrying the records it never
delivered, and Azure Sentinel requeues those instead of the whole batch

* fix(azure_sentinel): drop batches a permanent rejection will never accept

A non-413 4xx from the ingestion endpoint or from the OAuth token call means the request
will fail the same way on every retry, so requeueing it held the batch, and every record
logged behind it, until the queue cap dropped them. Retryable statuses (5xx, 408, 429)
still keep the whole batch, and a shared classifier gives Datadog the same rule

The serialization probe now catches any exception, not just TypeError and ValueError,
because safe_dumps hands pydantic models to model_dump and can raise anything. It also
splits on record count, so a recovery flush sends batch_size records per request instead
of serializing the whole requeued queue to measure it

Both integrations re-raise a cancelled send as exactly asyncio.CancelledError. Python
3.12's asyncio.wait_for only translates the exact class into TimeoutError, so the
BatchSendCancelled subclass escaped the logging worker as an unhandled error

The awaiting-retry flag now follows the queue that survived the max_queue_size trim, so
a deployment with the cap at zero is not left waiting for a timer flush with nothing
queued to retry

* chore(logging): document mutable queue ownership

Annotate the queue detach and requeue constructions required by the logger's appendable queue contract so the type-discipline budget stays clean

* fix(datadog): preserve non-413 retry behavior

Keep Datadog's existing contract of requeuing every non-413 HTTP failure while Azure Sentinel applies its permanent-client-error policy through the shared splitter

* fix(batch_utils): requeue by default and let Sentinel opt into dropping

The shared splitter's default non-success handler is now requeue_after_http_error, the behavior Datadog had before the extraction, so a caller that omits the argument keeps its records. Azure Sentinel passes undelivered_after_http_error explicitly to drop permanent 4xx rejections

Also drops an explicit return None the strict ruff gate flags in the test helper
This commit is contained in:
yucheng-berri 2026-09-05 17:15:36 -07:00 • committed by GitHub
parent 56a61cf016
commit d515a285b1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1221 additions and 130 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

170
litellm/integrations/azure_sentinel/azure_sentinel.py
View file

@ -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:

160
litellm/integrations/batch_utils.py Normal file
View file

@ -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 ()

64
litellm/integrations/datadog/datadog.py
View file

@ -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

4
litellm/types/integrations/azure_sentinel.py
View file

@ -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):
"""

112
tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py
View file

@ -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}']

841
tests/test_litellm/integrations/test_azure_sentinel.py
View file

@ -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]