From b318231fe9d8b2b98c61fe8ee7b2babc98939211 Mon Sep 17 00:00:00 2001 From: oss-agent-shin Date: Wed, 6 May 2026 15:50:06 -0700 Subject: [PATCH] Add Azure Sentinel audit log support (#27280) * Add Azure Sentinel audit log callback support Co-authored-by: ishaan-berri * Fix Azure Sentinel audit log batching Co-authored-by: ishaan-berri * Fix Azure Sentinel CI checks Co-authored-by: ishaan-berri --------- Co-authored-by: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com> Co-authored-by: ishaan-berri --- .../azure_sentinel/azure_sentinel.py | 139 ++++++++++--- .../custom_logger_registry.py | 2 + litellm/proxy/_types.py | 13 ++ .../integrations/test_azure_sentinel.py | 186 +++++++++++++++++- 4 files changed, 305 insertions(+), 35 deletions(-) diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index dd508e6c6c..0cfd49cda3 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -14,16 +14,18 @@ For batching specific details see CustomBatchLogger class import asyncio import os +import time import traceback -from typing import List, Optional +from typing import List, Optional, Union from litellm._logging import verbose_logger 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 ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload class AzureSentinelLogger(CustomBatchLogger): @@ -39,6 +41,7 @@ class AzureSentinelLogger(CustomBatchLogger): tenant_id: Optional[str] = None, client_id: Optional[str] = None, client_secret: Optional[str] = None, + audit_stream_name: Optional[str] = None, **kwargs, ): """ @@ -57,57 +60,77 @@ class AzureSentinelLogger(CustomBatchLogger): If not provided, will use AZURE_SENTINEL_CLIENT_ID or AZURE_CLIENT_ID env var. client_secret (str, optional): Azure Client Secret for OAuth2 authentication. If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var. + audit_stream_name (str, optional): Stream name from DCR for audit logs. + If not provided, audit logs use the standard stream name. """ self.async_httpx_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) - self.dcr_immutable_id = dcr_immutable_id or os.getenv( + resolved_dcr_immutable_id = dcr_immutable_id or os.getenv( "AZURE_SENTINEL_DCR_IMMUTABLE_ID" ) - self.stream_name = stream_name or os.getenv( - "AZURE_SENTINEL_STREAM_NAME", "Custom-LiteLLM" + resolved_stream_name = ( + stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM" ) - self.endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT") - self.tenant_id = ( + resolved_audit_stream_name = audit_stream_name or resolved_stream_name + resolved_endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT") + resolved_tenant_id = ( tenant_id or os.getenv("AZURE_SENTINEL_TENANT_ID") or os.getenv("AZURE_TENANT_ID") ) - self.client_id = ( + resolved_client_id = ( client_id or os.getenv("AZURE_SENTINEL_CLIENT_ID") or os.getenv("AZURE_CLIENT_ID") ) - self.client_secret = ( + resolved_client_secret = ( client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET") ) - if not self.dcr_immutable_id: + if not resolved_dcr_immutable_id: raise ValueError( "AZURE_SENTINEL_DCR_IMMUTABLE_ID is required. Set it as an environment variable or pass dcr_immutable_id parameter." ) - if not self.endpoint: + if not resolved_endpoint: raise ValueError( "AZURE_SENTINEL_ENDPOINT is required. Set it as an environment variable or pass endpoint parameter." ) - if not self.tenant_id: + if not resolved_tenant_id: raise ValueError( "AZURE_SENTINEL_TENANT_ID or AZURE_TENANT_ID is required. Set it as an environment variable or pass tenant_id parameter." ) - if not self.client_id: + if not resolved_client_id: raise ValueError( "AZURE_SENTINEL_CLIENT_ID or AZURE_CLIENT_ID is required. Set it as an environment variable or pass client_id parameter." ) - if not self.client_secret: + if not resolved_client_secret: raise ValueError( "AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET is required. Set it as an environment variable or pass client_secret parameter." ) + self.dcr_immutable_id = resolved_dcr_immutable_id + self.stream_name = resolved_stream_name + self.audit_stream_name = resolved_audit_stream_name + self.endpoint = resolved_endpoint + self.tenant_id = resolved_tenant_id + self.client_id = resolved_client_id + self.client_secret = resolved_client_secret + # Build API endpoint: {Endpoint}/dataCollectionRules/{DCR Immutable ID}/streams/{Stream Name}?api-version=2023-01-01 - self.api_endpoint = f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01" + self.api_endpoint = self._build_api_endpoint( + endpoint=resolved_endpoint, + dcr_immutable_id=resolved_dcr_immutable_id, + stream_name=resolved_stream_name, + ) + self.audit_api_endpoint = self._build_api_endpoint( + endpoint=resolved_endpoint, + dcr_immutable_id=resolved_dcr_immutable_id, + stream_name=resolved_audit_stream_name, + ) # OAuth2 scope for Azure Monitor self.oauth_scope = "https://monitor.azure.com/.default" @@ -118,6 +141,13 @@ class AzureSentinelLogger(CustomBatchLogger): super().__init__(**kwargs, flush_lock=self.flush_lock) asyncio.create_task(self.periodic_flush()) self.log_queue: List[StandardLoggingPayload] = [] + self.audit_log_queue: List[StandardAuditLogPayload] = [] + + @staticmethod + def _build_api_endpoint( + endpoint: str, dcr_immutable_id: str, stream_name: str + ) -> str: + return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01" async def _get_oauth_token(self) -> str: """ @@ -126,9 +156,6 @@ class AzureSentinelLogger(CustomBatchLogger): Returns: Bearer token string """ - # Check if we have a valid cached token - import time - if ( self.oauth_token and self.oauth_token_expires_at @@ -170,9 +197,6 @@ class AzureSentinelLogger(CustomBatchLogger): if not self.oauth_token: raise Exception("OAuth2 token response did not contain access_token") - # Cache token expiry time - import time - self.oauth_token_expires_at = time.time() + expires_in return self.oauth_token @@ -246,6 +270,34 @@ class AzureSentinelLogger(CustomBatchLogger): ) pass + async def async_log_audit_log_event( + self, audit_log: StandardAuditLogPayload + ) -> None: + """ + Async log LiteLLM audit log events to Azure Sentinel. + + Audit logs are queued separately from standard LLM logs so mixed callback + usage never sends schema-mismatched records in the same ingestion batch. + """ + try: + verbose_logger.debug( + "Azure Sentinel: Logging audit event id=%s action=%s table=%s", + audit_log.get("id"), + audit_log.get("action"), + audit_log.get("table_name"), + ) + + self.audit_log_queue.append(audit_log) + + if len(self.audit_log_queue) >= self.batch_size: + await self.async_send_audit_batch() + + except Exception as e: + verbose_logger.exception( + f"Azure Sentinel Audit Log Layer Error - {str(e)}\n{traceback.format_exc()}" + ) + pass + async def async_send_batch(self): """ Sends the batch of logs to Azure Monitor Logs Ingestion API @@ -253,22 +305,42 @@ 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", + ) + + 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", + ) + + async def _async_send_batch_to_api( + self, + log_queue: List[Union[StandardLoggingPayload, StandardAuditLogPayload]], + api_endpoint: str, + log_type: str, + ) -> None: try: - if not self.log_queue: + if not log_queue: return verbose_logger.debug( - "Azure Sentinel - about to flush %s events", len(self.log_queue) + "Azure Sentinel - about to flush %s %s", len(log_queue), log_type ) - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - # Get OAuth2 token bearer_token = await self._get_oauth_token() # Convert log queue to JSON array format expected by Logs Ingestion API # Each log entry should be a JSON object in the array - body = safe_dumps(self.log_queue) + body = safe_dumps(log_queue) # Set headers for Logs Ingestion API headers = { @@ -278,7 +350,7 @@ class AzureSentinelLogger(CustomBatchLogger): # Send the request response = await self.async_httpx_client.post( - url=self.api_endpoint, data=body.encode("utf-8"), headers=headers + url=api_endpoint, data=body.encode("utf-8"), headers=headers ) if response.status_code not in [200, 204]: @@ -301,4 +373,15 @@ class AzureSentinelLogger(CustomBatchLogger): f"Azure Sentinel Error sending batch API - {str(e)}\n{traceback.format_exc()}" ) finally: - self.log_queue.clear() + log_queue.clear() + + async def flush_queue(self): + if self.flush_lock is None: + return + + async with self.flush_lock: + if self.log_queue: + await self.async_send_batch() + if self.audit_log_queue: + await self.async_send_audit_batch() + self.last_flush_time = time.time() diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index f873bfeece..fd402b90d8 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -14,6 +14,7 @@ from litellm import _custom_logger_compatible_callbacks_literal from litellm.integrations.agentops import AgentOps from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.argilla import ArgillaLogger +from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger from litellm.integrations.bitbucket import BitBucketPromptManager from litellm.integrations.braintrust_logging import BraintrustLogger @@ -73,6 +74,7 @@ class CustomLoggerRegistry: "opik": OpikLogger, "argilla": ArgillaLogger, "opentelemetry": OpenTelemetry, + "azure_sentinel": AzureSentinelLogger, "azure_storage": AzureBlobStorageLogger, "humanloop": HumanloopLogger, # OTEL compatible loggers diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c6653a722d..2c97647979 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3348,6 +3348,19 @@ class AllCallbacks(LiteLLMPydanticObjectBase): ], ) + azure_sentinel: CallbackOnUI = CallbackOnUI( + litellm_callback_name="azure_sentinel", + ui_callback_name="Azure Sentinel", + litellm_callback_params=[ + "AZURE_SENTINEL_DCR_IMMUTABLE_ID", + "AZURE_SENTINEL_ENDPOINT", + "AZURE_SENTINEL_TENANT_ID", + "AZURE_SENTINEL_CLIENT_ID", + "AZURE_SENTINEL_CLIENT_SECRET", + "AZURE_SENTINEL_STREAM_NAME", + ], + ) + openmeter: CallbackOnUI = CallbackOnUI( litellm_callback_name="openmeter", ui_callback_name="OpenMeter", diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 031b85211f..30b246202f 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -2,13 +2,18 @@ Test Azure Sentinel logging integration """ -import datetime -from unittest.mock import AsyncMock, patch +import json +from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload + + +def _close_periodic_flush_task(coro): + coro.close() + return None @pytest.mark.asyncio @@ -20,7 +25,7 @@ async def test_azure_sentinel_oauth_and_send_batch(): test_client_id = "test-client-id" test_client_secret = "test-client-secret" - with patch("asyncio.create_task"): + with patch("asyncio.create_task", side_effect=_close_periodic_flush_task): logger = AzureSentinelLogger( dcr_immutable_id=test_dcr_id, endpoint=test_endpoint, @@ -42,9 +47,6 @@ async def test_azure_sentinel_oauth_and_send_batch(): # Add to queue logger.log_queue.append(standard_payload) - # Mock OAuth token response - from unittest.mock import MagicMock - mock_token_response = MagicMock() mock_token_response.status_code = 200 mock_token_response.json = MagicMock( @@ -91,3 +93,173 @@ async def test_azure_sentinel_oauth_and_send_batch(): # Verify queue is cleared assert len(logger.log_queue) == 0 + + +@pytest.mark.asyncio +async def test_azure_sentinel_queues_audit_log_event(): + """Test that Azure Sentinel supports direct audit log callbacks""" + with patch("asyncio.create_task", side_effect=_close_periodic_flush_task): + logger = AzureSentinelLogger( + dcr_immutable_id="dcr-test123456789", + endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com", + tenant_id="test-tenant-id", + client_id="test-client-id", + client_secret="test-client-secret", + ) + + logger.batch_size = 2 + logger.async_send_audit_batch = AsyncMock() + + audit_log = StandardAuditLogPayload( + id="audit-123", + 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='{"team_alias": "sentinel-demo"}', + ) + + await logger.async_log_audit_log_event(audit_log) + + assert logger.audit_log_queue == [audit_log] + logger.async_send_audit_batch.assert_not_called() + + await logger.async_log_audit_log_event(audit_log) + + assert logger.audit_log_queue == [audit_log, audit_log] + logger.async_send_audit_batch.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_azure_sentinel_sends_audit_log_payload_to_ingestion_api(): + """Test that queued audit logs are sent to Azure Monitor Logs Ingestion""" + with patch("asyncio.create_task", side_effect=_close_periodic_flush_task): + logger = AzureSentinelLogger( + dcr_immutable_id="dcr-test123456789", + endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com", + tenant_id="test-tenant-id", + client_id="test-client-id", + client_secret="test-client-secret", + ) + + audit_log = StandardAuditLogPayload( + id="audit-123", + 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='{"team_alias": "sentinel-demo"}', + ) + await logger.async_log_audit_log_event(audit_log) + + mock_token_response = MagicMock() + mock_token_response.status_code = 200 + mock_token_response.json = MagicMock( + return_value={ + "access_token": "test-bearer-token", + "expires_in": 3600, + } + ) + mock_token_response.text = "Success" + + mock_api_response = MagicMock() + mock_api_response.status_code = 204 + mock_api_response.text = "Success" + + async def mock_post(*args, **kwargs): + if "oauth2/v2.0/token" in kwargs.get("url", ""): + return mock_token_response + return mock_api_response + + logger.async_httpx_client.post = AsyncMock(side_effect=mock_post) + + await logger.flush_queue() + + api_call_args = logger.async_httpx_client.post.call_args_list[-1] + body = json.loads(api_call_args.kwargs["data"].decode("utf-8")) + assert body == [audit_log] + assert "dcr-test123456789" in api_call_args.kwargs["url"] + assert "Custom-LiteLLM" in api_call_args.kwargs["url"] + assert len(logger.audit_log_queue) == 0 + + +@pytest.mark.asyncio +async def test_azure_sentinel_flushes_standard_and_audit_logs_separately(): + """Test mixed callback roles do not send schema-mismatched batches.""" + with patch("asyncio.create_task", side_effect=_close_periodic_flush_task): + logger = AzureSentinelLogger( + dcr_immutable_id="dcr-test123456789", + stream_name="Custom-LiteLLM-Standard", + audit_stream_name="Custom-LiteLLM-Audit", + endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com", + tenant_id="test-tenant-id", + client_id="test-client-id", + client_secret="test-client-secret", + ) + + standard_payload = StandardLoggingPayload( + id="standard-123", + call_type="completion", + model="gpt-3.5-turbo", + status="success", + messages=[{"role": "user", "content": "Hello"}], + response={"choices": [{"message": {"content": "Hi"}}]}, + ) + audit_log = StandardAuditLogPayload( + id="audit-123", + 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='{"team_alias": "sentinel-demo"}', + ) + + logger.log_queue.append(standard_payload) + await logger.async_log_audit_log_event(audit_log) + + mock_token_response = MagicMock() + mock_token_response.status_code = 200 + mock_token_response.json = MagicMock( + return_value={ + "access_token": "test-bearer-token", + "expires_in": 3600, + } + ) + mock_token_response.text = "Success" + + mock_api_response = MagicMock() + mock_api_response.status_code = 204 + mock_api_response.text = "Success" + + async def mock_post(*args, **kwargs): + if "oauth2/v2.0/token" in kwargs.get("url", ""): + return mock_token_response + return mock_api_response + + logger.async_httpx_client.post = AsyncMock(side_effect=mock_post) + + await logger.flush_queue() + + ingestion_calls = [ + call + for call in logger.async_httpx_client.post.call_args_list + if "dataCollectionRules" in call.kwargs["url"] + ] + assert len(ingestion_calls) == 2 + + standard_call, audit_call = ingestion_calls + assert "Custom-LiteLLM-Standard" in standard_call.kwargs["url"] + assert json.loads(standard_call.kwargs["data"].decode("utf-8")) == [ + standard_payload + ] + assert "Custom-LiteLLM-Audit" in audit_call.kwargs["url"] + assert json.loads(audit_call.kwargs["data"].decode("utf-8")) == [audit_log]