From ba1bde70e4b45d4f2bf5e9dd4b49858e7d9ac691 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 00:20:25 +0000 Subject: [PATCH] feat(guardrails/rubrik): prompt moderation, response-text blocking, streaming buffer, failure logging (#35722) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(guardrails/rubrik): prompt moderation, response-text blocking, streaming buffer, failure logging (#34019) * feat(guardrails/rubrik): add prompt moderation, response-text blocking, streaming buffer, failure logging - Add `pre_call` prompt moderation via `/v1/before_prompt/openai/v1` webhook: structured messages are flattened and sent before the LLM is called; blocked prompts surface a `ModifyResponseException` with the refusal text. - Extend `post_call` response moderation to cover assistant text in addition to tool calls; text blocks (wholesale replacement) are distinguished from tool-block explanations (appended) via `startswith` diffing. - Add `streaming_end_of_stream_only = True` and `streaming_buffer_until_moderated = True` so streamed responses are withheld until end-of-stream moderation passes (requires litellm >= BerriAI/litellm#31389; older versions fall back to detect-only). - Add `_MalformedToolBlockingResponseError` for structurally invalid service responses; `_guarded` logs at CRITICAL so operators notice misconfiguration. - Add `max_queue_size = 10_000`, `_enforce_max_queue_size`, and drop-oldest backpressure so a webhook outage cannot grow the retry queue unboundedly. - Add `flush_queue` override that snapshots once for both send and drain, preventing duplicate delivery on concurrent flush calls. - Make `_log_batch_to_rubrik` re-raise on error so `flush_queue` preserves undelivered events for the next retry. - Add `async_post_call_failure_hook` to log blocked requests (`ModifyResponseException`) with a best-effort fallback payload for prompt blocks (where no `standard_logging_object` exists yet). - Add `_correlation_id` / `_apply_correlation_id` / `_prepend_system_prompt` helpers; `_prepare_log_payload` now applies them for all providers (not just Anthropic) so every log correlates by `litellm_call_id`. - Add `get_supported_event_hooks` classmethod advertising `[pre_call, post_call]`. - Use dedicated `httpx.AsyncClient` (`moderation_client`) for webhook calls with explicit pool limits, separate from the shared logging client. - Drop module-level `rubrik_handler` singleton (inappropriate for a library). - Update `initialize_guardrail` docstring to explain `pre_call` vs `post_call` mode. - Update tests: rename `tool_blocking_client` → `moderation_client`, `tool_blocking_endpoint` → `response_moderation_endpoint`, `_flush_task` → `_periodic_flush_task`; migrate `TestExtractBlockedTools` to `TestExtractResponseBlock` for the new combined text+tool block API; add tests for prompt moderation, text blocking, streaming flags, and failure payload construction. Co-Authored-By: Claude Sonnet 4.6 (1M context) * test(guardrails/rubrik): add tests to reach 100% coverage 50 new tests across 18 classes covering previously-untested paths: - Prompt moderation: passthrough, block, no-messages skip, message flattening (content-list → string), payload construction with tools/user/correlation_key/litellm_call_id fallback, refusal extraction - async_post_call_failure_hook: non-matching exception no-op, missing stash warning, valid stash → enqueue, AttributeError in payload build, flush exception handling - Block payload building: standard_logging_object present vs fallback path, missing start_time - async_log_success_event: _rubrik_blocked=True skip path - aclose: task cancel + moderation_client.aclose() - Edge cases: sampling rate clamp warning, unknown input_type passthrough, empty-inputs early return, model_call_details warning, _stash_block_context, duck-typed tool-call normalization, request_data["tools"] preference over optional_params, system-prompt exception handler, flush-at-batch-size, enqueue exception swallowing, queue empty/lock-None guards, non-dict JSON response TypeError Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): use get_async_httpx_client, ruff format - Replace bare httpx.AsyncClient with get_async_httpx_client (required by ensure_async_clients_test; avoids per-request client creation) - aclose() calls close() (AsyncHTTPHandler interface, not aclose()) - ruff format on rubrik.py and guardrail_hooks/rubrik/__init__.py - Update 3 tests for AsyncHTTPHandler type (isinstance check, close()) osv-scan and documentation CI failures are pre-existing on the base branch and unrelated to this PR. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): fix UP006 strict ruff violation get_supported_event_hooks return type used List[...] (UP006) instead of list[...]. Replace with the built-in generic and remove the now-unused List import from typing. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): fix 3 reportArgumentType basedpyright violations Use `# pyright: ignore[reportArgumentType]` (not `# type: ignore`) to suppress the three errors basedpyright reports in --outputjson mode: - convert_content_list_to_str call (dict vs AllMessageValues) - _apply_correlation_id call (StandardLoggingPayload vs dict[str, Any]) - _prepend_system_prompt call (same) Also tighten _apply_correlation_id and _prepend_system_prompt signatures from bare `dict` to `dict[str, Any]`. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): don't close shared HTTP client in aclose() moderation_client and async_httpx_client both come from LiteLLM's global HTTP-client cache (get_async_httpx_client keys on llm_provider + params). Two RubrikLogger instances with the same parameters share the same underlying AsyncHTTPHandler object. Calling close() in aclose() closed the shared connection pool for all instances, breaking any subsequent moderation request on other loggers. aclose() now only cancels the periodic flush task and lets LiteLLM manage the shared client lifecycle. Tests updated to assert close() is NOT called. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): use Counter for duplicate tool-call ID detection Set-based comparison lost ID multiplicity: two original tool calls with the same ID both appeared "allowed" even when the service returned only one (e.g. one allowed + one prohibited sharing an ID). Replace with Counter so returned_id_counts[id] >= required_id_counts[id] must hold for every ID. Matches the approach in the original _extract_blocked_tools. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): respect default_on=true when omitted from config LitellmParams.__init__ converts an omitted default_on to False before initialize_guardrail receives it, so litellm_params.default_on is always bool and never None. The is-None guard in RubrikLogger.__init__ therefore never fired on the proxy path, leaving prompt/response moderation inactive for any config that omitted default_on. Fix: read the raw guardrail dict (before LitellmParams coercion) to distinguish an explicit `default_on: false` from the absent-means-True default. When the key is absent from the raw config, default_on=True is used; when it is explicitly set (either True or False), that value wins. Co-Authored-By: Claude Sonnet 4.6 (1M context) * style: ruff format rubrik.py after Counter import addition Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): detect ID-less tool call removal; fix UP045 ID-less tool calls (tc.id is falsy) were excluded from required_id_counts, so the Counter comparison never caught their removal. Add a cardinality check (len(returned) < len(original)) that fires on any removal regardless of ID presence, combined with the Counter check for duplicate-ID attacks. Also fix 5 UP045 violations (Optional[X] → X | None) introduced by our new code against the daily-branch baseline. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): filter optional_params through ModelParamHelper in fallback payload _build_fallback_payload forwarded the raw optional_params dict as model_parameters. optional_params can contain extra_headers, api_key, and other upstream provider credentials that must not reach the Rubrik webhook. The normal standard_logging_object path already filters through ModelParamHelper.get_standard_logging_model_parameters(), which allowlists only safe LLM API parameters. Apply the same filter here. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): scope failure hook by guardrail_name; moderate text-completions Guard async_post_call_failure_hook by guardrail_name so multiple Rubrik instances don't cross-log: the failure hook is called for every registered callback; without the check the first instance pops the stash and the originating instance finds None and silently skips logging. Now each instance only handles blocks raised by itself. Also moderate /v1/completions prompts: _moderate_prompt returned early when structured_messages was absent. For text-completion requests litellm supplies inputs["texts"] with no structured_messages. Added a fallback that synthesises a user-message from texts so the before_prompt webhook can evaluate text-completion prompts. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(lint): add reason comments to pyright: ignore suppressions type-discipline budget requires each # pyright: ignore[...] to carry an explanatory comment. Add reasons to the three bare suppressions on lines 483, 651, 652. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): include tool-call arguments in prompt moderation _flatten_messages_for_moderation only sent the content field, silently dropping tool_calls[].function.arguments and function_call.arguments. An attacker could embed prohibited text in tool-call arguments inside assistant history turns and bypass prompt moderation entirely. Now collects all attacker-controlled text per message: text content via convert_content_list_to_str, plus all tool_calls[].function.arguments and the deprecated function_call.arguments, joined with newlines before being sent to the before_prompt webhook. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): tighten append detection to prevent prefix bypass startswith(sent_content) allowed any replacement whose text shares the original as a prefix (e.g. "Hello" → "Hello, blocked.") to be classified as a tool-block append rather than a text block, bypassing detection. Use startswith(f"{sent_content}\n\n") to require the exact two-newline separator the webhook uses between original text and appended tool-block explanations. Also add `returned_content != sent_content` to text_blocked so an unchanged passthrough is never classified as a block. Co-Authored-By: Claude Sonnet 4.6 (1M context) * fix(guardrails/rubrik): default_on=False when omitted (follow existing pattern) Remove the custom raw-dict lookup that was defaulting default_on to True when omitted from the guardrail config. Follow the standard litellm convention: omitted resolves to False (users must explicitly opt in with default_on: true). - initialize_guardrail: pass litellm_params.default_on directly - RubrikLogger.__init__: is-None guard defaults to False not True - Test updated to assert the correct False default Co-Authored-By: Claude Sonnet 4.6 (1M context) --------- Co-authored-by: Claude Sonnet 4.6 (1M context) * chore(rubrik): keep the ported guardrail within staging lint budgets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: credit the original author of the rubrik guardrail work Co-authored-by: Joseph Barker <156112794+seph-barker@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: keep this mirror PR's diff limited to the rubrik files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Joseph Barker <156112794+seph-barker@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 (1M context) Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/rubrik.py | 1040 +++++++++++++---- .../guardrail_hooks/rubrik/__init__.py | 12 + .../test_litellm/integrations/test_rubrik.py | 956 +++++++++++++-- 3 files changed, 1691 insertions(+), 317 deletions(-) diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 2e49da45ce9..4bcbe8bae37 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -1,12 +1,14 @@ -"""Rubrik LiteLLM Plugin for tool blocking and batch logging.""" +"""Rubrik LiteLLM Plugin for prompt/response moderation and batch logging.""" import asyncio import os import random import time -import urllib.parse import uuid from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Literal, Optional import httpx @@ -18,6 +20,10 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.litellm_core_utils.core_helpers import safe_deep_copy +from litellm.litellm_core_utils.model_param_helper import ModelParamHelper +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -35,15 +41,16 @@ if TYPE_CHECKING: Logging as LiteLLMLoggingObj, ) -_ENDPOINT_ANTHROPIC_MESSAGES = "/v1/messages" -_WEBHOOK_PATH_TOOL_BLOCKING = "/v1/after_completion/openai/v1" +_WEBHOOK_PATH_RESPONSE_MODERATION = "/v1/after_completion/openai/v1" +_WEBHOOK_PATH_PROMPT_MODERATION = "/v1/before_prompt/openai/v1" _WEBHOOK_PATH_LOGGING_BATCH = "/v1/litellm/batch" _MAX_QUEUE_SIZE = 10_000 _DROP_WARNING_INTERVAL_SECONDS = 60.0 +_EMPTY_MAPPING: Mapping[str, Any] = MappingProxyType({}) class _MalformedToolBlockingResponseError(Exception): - """Raised when the tool blocking service returns a structurally invalid + """Raised when the response moderation service returns a structurally invalid response (e.g. empty ``choices``). Distinct from transient network/HTTP errors so callers can surface a @@ -52,11 +59,15 @@ class _MalformedToolBlockingResponseError(Exception): """ -class RubrikLogger(CustomGuardrail, CustomBatchLogger): - @classmethod - def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: - return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] +@dataclass +class BlockedResponseResult: + """Returned by _extract_response_block when the response was blocked + (response text replaced, or at least one tool call removed).""" + explanation: str + + +class RubrikLogger(CustomGuardrail, CustomBatchLogger): def __init__( self, api_key: str | None = None, @@ -67,21 +78,82 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): kwargs.setdefault("guardrail_name", "rubrik") # `initialize_guardrail` always passes these kwargs explicitly, with # value `None` when the user omits `mode` / `default_on` from the - # guardrail config. Coerce None (omitted) to the desired default - # while preserving any explicit value the caller did set -- - # in particular `default_on=False` if the user wants the guardrail - # off by default. + # guardrail config. Follow the standard litellm convention: omitted + # resolves to False (off by default, user must opt in explicitly). kwargs["event_hook"] = kwargs.get("event_hook") or GuardrailEventHooks.post_call if kwargs.get("default_on") is None: - kwargs["default_on"] = True - kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) + kwargs["default_on"] = False super().__init__( flush_lock=self.flush_lock, + supported_event_hooks=list(self.get_supported_event_hooks()), **kwargs, ) verbose_logger.debug("initializing rubrik logger") + # Defining ``apply_guardrail`` routes streaming responses through + # litellm's ``unified_guardrail.async_post_call_streaming_iterator_hook``. + # By default that hook samples intermediate chunks + # (``streaming_sampling_rate``, default 5) and also moderates at + # end-of-stream, so a streamed response costs ~ceil(N/5)+1 Rubrik + # webhook round-trips. litellm reads this attribute via + # ``getattr(guardrail, "streaming_end_of_stream_only", False)``; when + # True it yields chunks unprocessed and only moderates the fully + # assembled response once at end of stream. + self.streaming_end_of_stream_only = True + + # ``streaming_end_of_stream_only`` is detect-only: it releases every + # chunk to the client *before* moderating, so a block can only append a + # trailing message -- the original content has already been delivered. + # ``streaming_buffer_until_moderated`` (litellm >= BerriAI/litellm#31389) + # withholds all chunks until end-of-stream moderation passes, then + # releases the original response (clean) or only the block message + # (blocked). On older litellm this attribute is ignored and we fall + # back to the detect-only behavior above. + self.streaming_buffer_until_moderated = True + + self._parse_sampling_rate() + + self.key = api_key or os.getenv("RUBRIK_API_KEY") + if not self.key: + verbose_logger.warning("Rubrik: No API key configured. Requests will be unauthenticated.") + + self._parse_batch_size() + + # Cap the in-memory retry queue so a Rubrik webhook outage cannot let + # authenticated traffic accumulate prompt/response payloads until the + # proxy runs out of memory. Once the cap is reached, oldest events are + # dropped to make room for fresh ones (drop-oldest backpressure). + self.max_queue_size = _MAX_QUEUE_SIZE + self._dropped_since_warning = 0 + self._last_drop_warning_time = 0.0 + + _webhook_url = api_base or os.getenv("RUBRIK_WEBHOOK_URL") + if not _webhook_url: + raise ValueError("Rubrik webhook URL not configured. Set RUBRIK_WEBHOOK_URL or pass api_base.") + + _webhook_url = _webhook_url.rstrip("/").removesuffix("/v1") + self._setup_clients(_webhook_url) + + self._headers: Mapping[str, str] = MappingProxyType( + {"Content-Type": "application/json", "Authorization": f"Bearer {self.key}"} + if self.key + else {"Content-Type": "application/json"} + ) + + self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + """Return the guardrail event hooks this integration supports. + + Prompt moderation (``pre_call``) evaluates the user's message before + the LLM is called. Response moderation (``post_call``) evaluates the + assistant's reply and tool calls after the LLM returns. + """ + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + + def _parse_sampling_rate(self) -> None: self.sampling_rate = 1.0 rbrk_sampling_rate = os.getenv("RUBRIK_SAMPLING_RATE") if rbrk_sampling_rate is not None: @@ -93,80 +165,54 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): except ValueError: verbose_logger.warning(f"Invalid RUBRIK_SAMPLING_RATE: {rbrk_sampling_rate!r}, using 1.0") - self.key = api_key or os.getenv("RUBRIK_API_KEY") - if not self.key: - verbose_logger.warning("Rubrik: No API key configured. Requests will be unauthenticated.") + def _parse_batch_size(self) -> None: _batch_size = os.getenv("RUBRIK_BATCH_SIZE") - if _batch_size: try: - self.batch_size = int(_batch_size) + parsed_size = int(_batch_size) + if parsed_size <= 0: + verbose_logger.warning(f"RUBRIK_BATCH_SIZE={_batch_size!r} must be > 0, using default") + else: + self.batch_size = parsed_size except ValueError: verbose_logger.warning(f"Invalid RUBRIK_BATCH_SIZE: {_batch_size!r}, using default") - # Cap the in-memory retry queue so a Rubrik webhook outage cannot let - # authenticated traffic accumulate prompt/response payloads until the - # proxy runs out of memory. Once the cap is reached, oldest events are - # dropped to make room for fresh ones (drop-oldest backpressure). - self.max_queue_size = _MAX_QUEUE_SIZE - self._dropped_since_warning = 0 - self._last_drop_warning_time = 0.0 - - _webhook_url = api_base or os.getenv("RUBRIK_WEBHOOK_URL") - - if _webhook_url is None: - raise ValueError("Rubrik webhook URL not configured. Set RUBRIK_WEBHOOK_URL or pass api_base.") - - _webhook_url = _webhook_url.rstrip("/").removesuffix("/v1") - self.tool_blocking_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_TOOL_BLOCKING}" - self.logging_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_LOGGING_BATCH}" + def _setup_clients(self, webhook_url: str) -> None: + self.response_moderation_endpoint = f"{webhook_url}{_WEBHOOK_PATH_RESPONSE_MODERATION}" + self.prompt_moderation_endpoint = f"{webhook_url}{_WEBHOOK_PATH_PROMPT_MODERATION}" + self.logging_endpoint = f"{webhook_url}{_WEBHOOK_PATH_LOGGING_BATCH}" self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) - self.tool_blocking_client = get_async_httpx_client( + self.moderation_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback, params={"timeout": httpx.Timeout(5.0, connect=2.0)}, ) - self._headers: dict[str, str] = {"Content-Type": "application/json"} - if self.key: - self._headers["Authorization"] = f"Bearer {self.key}" - - # Periodic flush is started lazily on the first log event so that - # low-traffic deployments still get their batches drained even when the - # logger is instantiated outside a running event loop (sync init). - self._flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() - def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None: """Start the periodic flush task only when an event loop is already running.""" try: loop = asyncio.get_running_loop() except RuntimeError: - verbose_logger.debug( - "Rubrik logger init: no running event loop, periodic flush will start on first log event." - ) return None return loop.create_task(self.periodic_flush()) def _ensure_periodic_flush_task(self) -> None: - # Synchronous helper: in asyncio's cooperative model there is no await - # between the check and assignment, so two callers cannot race here. - if self._flush_task is None or self._flush_task.done(): - self._flush_task = self._start_periodic_flush_task() + if self._periodic_flush_task is None or self._periodic_flush_task.done(): + self._periodic_flush_task = self._start_periodic_flush_task() async def aclose(self): - """Close the dedicated HTTP clients used by this logger.""" - # Cancel the periodic flush task before closing the HTTP clients so - # the loop doesn't wake up and try to POST via a closed client. - if self._flush_task is not None and not self._flush_task.done(): - self._flush_task.cancel() - try: - await self._flush_task - except (asyncio.CancelledError, Exception): - pass - self._flush_task = None - await self.tool_blocking_client.close() - await self.async_httpx_client.close() + """Cancel the periodic flush task. + + ``moderation_client`` and ``async_httpx_client`` are shared objects + from LiteLLM's global HTTP-client cache (``get_async_httpx_client`` + uses the same cache key for all instances with equal parameters). + Closing them here would close the shared connection pool for every + other logger instance; let LiteLLM manage their lifecycle instead. + """ + task = getattr(self, "_periodic_flush_task", None) + if task is not None: + task.cancel() # -- Guardrail hook -------------------------------------------------------- @@ -177,67 +223,104 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - """Validate tool calls against the blocking service (fail-open).""" - if input_type != "response": - return inputs + """Moderate prompts (request) and responses (response); fail-open. - tool_calls = inputs.get("tool_calls") - if not tool_calls: - return inputs + - ``request``: evaluate the prompt via the before_prompt webhook and + block disallowed prompts before the model is called. + - ``response``: evaluate the assistant's response text and tool calls + via the after_completion webhook and block on a policy violation. + litellm's guardrail-translation layer normalizes Anthropic and OpenAI + requests/responses into ``inputs`` before this runs, so a single code + path covers both wire formats. The configured guardrail ``mode`` + selects which surface(s) run. + """ + if input_type == "request": + return await self._guarded( + self._moderate_prompt(inputs, request_data, logging_obj), + inputs, + "Prompt moderation", + ) + if input_type == "response": + return await self._guarded( + self._moderate_response(inputs, request_data, logging_obj), + inputs, + "Response moderation", + ) + return inputs + + @staticmethod + async def _guarded( + coro: Any, + inputs: GenericGuardrailAPIInputs, + label: str, + ) -> GenericGuardrailAPIInputs: + """Await a moderation coroutine fail-open: re-raise an intentional + block, log at critical for malformed service responses, and swallow + any other error returning ``inputs`` unchanged.""" try: - return await self._check_tool_calls(inputs, tool_calls, request_data, logging_obj) + return await coro except ModifyResponseException: raise except _MalformedToolBlockingResponseError as e: - # Distinct from transient errors: the service responded but the - # payload was structurally invalid, which usually indicates a - # misconfigured webhook or a breaking change in its response - # format. Log loudly so operators notice their tool-blocking - # policy is not actually being enforced. + # The service responded but the payload was structurally invalid, + # which usually indicates a misconfigured webhook or a breaking + # change in its response format. Log loudly so operators notice + # their moderation policy is not actually being enforced. verbose_logger.critical( - "Tool blocking service returned a malformed response: %s. " - "Tool calls are NOT being checked -- verify the webhook " - "configuration. Returning original response unchanged.", + "Response moderation service returned a malformed response: %s. " + "Requests are NOT being checked -- verify the webhook " + "configuration. Returning original inputs unchanged.", e, exc_info=True, ) return inputs except Exception as e: verbose_logger.error( - f"Tool blocking hook failed: {e}. Returning original response unchanged.", + f"{label} hook failed: {e}. Returning original inputs unchanged.", exc_info=True, ) return inputs - async def _check_tool_calls( + async def _moderate_response( self, inputs: GenericGuardrailAPIInputs, - tool_calls: Any, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"], ) -> GenericGuardrailAPIInputs: - """Send tool calls to blocking service, raise if any are blocked.""" - message_tool_calls = self._normalize_tool_calls(tool_calls) + """Send response text + tool calls to the after_completion webhook and + raise if either the response text or any tool call is blocked.""" + tool_calls = inputs.get("tool_calls") + texts = inputs.get("texts") + if not tool_calls and not texts: + return inputs - call_details = getattr(logging_obj, "model_call_details", {}) if logging_obj else {} - response = request_data.get("response") - request_id = getattr(response, "id", None) if response else None + message_tool_calls = self._normalize_tool_calls(tool_calls or ()) + sent_content = self._join_texts(texts) + + call_details = getattr(logging_obj, "model_call_details", _EMPTY_MAPPING) if logging_obj else _EMPTY_MAPPING if logging_obj and not call_details: verbose_logger.warning( "Rubrik: logging_obj present but model_call_details is empty -- request context will be missing" ) - response_data = self._build_tool_call_payload(message_tool_calls, request_id) - req_data = self._extract_request_data(call_details) + # The moderation payload's ``id`` becomes the tool-blocking log's + # correlation key (the S3 filename), so it must match the failure + # (response) log written for the same blocked request. Both use + # ``litellm_call_id`` -- see ``_correlation_id``. + request_id = self._correlation_id(call_details, request_data) - service_response = await self._post_to_tool_blocking_service(response_data, req_data) - blocked_explanation = self._extract_blocked_tools(service_response, message_tool_calls) + response_data = self._build_response_moderation_payload(message_tool_calls, sent_content, request_id) + req_data = self._extract_request_data(call_details, request_data) - if blocked_explanation is not None: + service_response = await self._post_to_response_moderation_endpoint(response_data, req_data) + blocked = self._extract_response_block(service_response, message_tool_calls, sent_content) + + if blocked: model = self._resolve_model(request_data, call_details) + self._stash_block_context(logging_obj, request_data) raise ModifyResponseException( - message=blocked_explanation, + message=blocked.explanation, model=model, request_data=request_data, guardrail_name=self.guardrail_name, @@ -245,43 +328,125 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return inputs - @staticmethod - def _normalize_tool_calls(tool_calls: Any) -> list[ChatCompletionMessageToolCall]: - """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" - result = [] - for tc in tool_calls: - if isinstance(tc, ChatCompletionMessageToolCall): - result.append(tc) - elif isinstance(tc, dict): - func = tc.get("function", {}) - result.append( - ChatCompletionMessageToolCall( - id=tc.get("id", ""), - type=tc.get("type", "function"), - function=Function( - name=func.get("name", ""), - arguments=func.get("arguments", ""), - ), - ) - ) - elif hasattr(tc, "id") and hasattr(tc, "function"): - result.append( - ChatCompletionMessageToolCall( - id=tc.id or "", - type=getattr(tc, "type", None) or "function", - function=tc.function, - ) - ) - else: - raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}") - return result + async def _moderate_prompt( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> GenericGuardrailAPIInputs: + """Send the (normalized) prompt to the before_prompt webhook and raise + if the prompt is blocked.""" + messages = inputs.get("structured_messages") + if not messages: + # For non-chat request types (e.g. /v1/completions), litellm + # supplies the prompt as ``texts`` with no structured_messages. + # Synthesise a user-message so the webhook can evaluate the prompt. + texts = inputs.get("texts") + if texts: + joined = "\n".join(t for t in texts if t) + if joined: + messages = [{"role": "user", "content": joined}] + if not messages: + return inputs + + payload = self._build_prompt_moderation_payload(inputs, request_data) + service_response = await self._post_to_prompt_moderation_endpoint(payload) + refusal = self._extract_prompt_refusal(service_response) + if refusal is None: + return inputs + + model = inputs.get("model") or request_data.get("model") or "unknown" + self._stash_block_context(logging_obj, request_data) + raise ModifyResponseException( + message=refusal, + model=model, + request_data=request_data, + guardrail_name=self.guardrail_name, + ) @staticmethod - def _build_tool_call_payload( - tool_calls: list[ChatCompletionMessageToolCall], + def _stash_block_context( + logging_obj: Optional["LiteLLMLoggingObj"], + request_data: dict, + ) -> None: + """Stash signals so the deferred success-event skips this request and + ``async_post_call_failure_hook`` can build the failure payload. + + - Sets a flag on ``logging_obj.model_call_details`` so the deferred + success-event handler short-circuits. + - Stashes a reference to ``logging_obj`` on ``request_data`` under a + custom key. ``ProxyLogging.post_call_failure_hook`` pops only + ``litellm_logging_obj`` before iterating callbacks, so this key + survives. + + When ``logging_obj`` is ``None`` the success-event has no way to + observe the block (the flag has nowhere to live), so we log an error + instead of silently dropping the signal. + """ + if logging_obj is None: + verbose_logger.error( + "Rubrik: moderation block fired with logging_obj=None for " + f"litellm_call_id={request_data.get('litellm_call_id')}; " + "cannot suppress success event or attach failure payload." + ) + request_data["_rubrik_logging_obj"] = None + return + logging_obj.model_call_details["_rubrik_blocked"] = True + request_data["_rubrik_logging_obj"] = logging_obj + + @staticmethod + def _normalize_tool_calls(tool_calls: Any) -> tuple[ChatCompletionMessageToolCall, ...]: + """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" + return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls) + + @staticmethod + def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall: + if isinstance(tc, ChatCompletionMessageToolCall): + return tc + if isinstance(tc, dict): + func = tc.get("function") or _EMPTY_MAPPING + return ChatCompletionMessageToolCall( + id=tc.get("id", ""), + type=tc.get("type", "function"), + function=Function( + name=func.get("name", ""), + arguments=func.get("arguments", ""), + ), + ) + if hasattr(tc, "id") and hasattr(tc, "function"): + return ChatCompletionMessageToolCall( + id=tc.id or "", + type=getattr(tc, "type", None) or "function", + function=tc.function, + ) + raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}") + + @staticmethod + def _join_texts(texts: Any) -> str: + """Join response text segments into the single content string the + webhook evaluates. Empty when there is no assistant text.""" + if not texts: + return "" + return "\n".join(t for t in texts if t) + + @staticmethod + def _build_response_moderation_payload( + tool_calls: Sequence[ChatCompletionMessageToolCall], + content: str, request_id: str | None, - ) -> dict[str, Any]: - """Build a full OpenAI ChatCompletion-format dict for the blocking service.""" + ) -> Mapping[str, Any]: + """Build an OpenAI ChatCompletion-format dict (assistant text + tool + calls) for the after_completion webhook. + + ``content`` is sent so the webhook can moderate the response text; + ``None`` when the assistant produced no text (tool-call-only response). + """ + message: dict[str, Any] = { + "role": "assistant", + "content": content or None, + } + if tool_calls: + message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls) return { "id": request_id or f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", @@ -290,42 +455,133 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "choices": [ { "index": 0, - "message": { - "role": "assistant", - "content": None, - "tool_calls": [tc.model_dump(exclude_none=True) for tc in tool_calls], - }, - "finish_reason": "tool_calls", + "message": message, + "finish_reason": "tool_calls" if tool_calls else "stop", } ], } @staticmethod - def _extract_request_data(call_details: dict[str, Any]) -> dict[str, Any]: - """Extract original request data from model_call_details.""" - if not call_details: - return {} - litellm_params = call_details.get("litellm_params", {}) or {} + def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]: + """Collapse each message's content to a plain string for the webhook. + + litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape, + but a turn sent as content-parts (``[{"type": "text", ...}]``) stays a + list. The before_prompt webhook reads ``content`` as a string and drops + non-string content, so we flatten text parts here (images skipped, per + ``convert_content_list_to_str``) -- otherwise block-content prompts + would pass through unmoderated. Builds a new list; never mutates the + shared ``structured_messages``. + """ + return tuple( + { + "role": message.get("role"), + "content": "\n".join(p for p in RubrikLogger._moderation_text_parts(message) if p), + } + for message in messages or () + if isinstance(message, dict) + ) + + @staticmethod + def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]: + """Every attacker-controlled text segment of a message: its content plus + the arguments of any tool call or deprecated function call.""" + fc = message.get("function_call") + return ( + # Base text content (flattens Anthropic content-part arrays) + convert_content_list_to_str(message), # pyright: ignore[reportArgumentType] # dict[str,Any] is AllMessageValues at runtime + *( + str((tc.get("function") or _EMPTY_MAPPING).get("arguments") or "") + for tc in message.get("tool_calls") or () + if isinstance(tc, dict) + ), + str((fc.get("arguments") if isinstance(fc, dict) else None) or ""), + ) + + @staticmethod + def _build_prompt_moderation_payload( + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, Any], + ) -> Mapping[str, Any]: + """Build the bare OpenAI request the before_prompt webhook consumes. + + Unlike the after_completion envelope, this endpoint takes a raw OpenAI + chat-completions request. ``structured_messages`` is litellm's + OpenAI-normalized view of the prompt, so this works for Anthropic + ``/v1/messages`` requests too. Optional fields are sent only when + present so the payload stays clean. + """ + payload: dict[str, Any] = { + "model": inputs.get("model") or request_data.get("model") or "", + "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), + } + tools = inputs.get("tools") + if tools is not None: + payload["tools"] = tools + user = request_data.get("user") + if user: + payload["user"] = user + # Fall back to litellm_call_id, the stable cross-provider join key the + # response/tool path uses (see _correlation_id). LiteLLM does not + # populate request_data["correlation_key"]; it carries litellm_call_id. + # The before_prompt webhook skips the *_prompt_moderation.json S3 write + # when correlation_key is empty, so without this the block fires but no + # log is ever written. An explicit correlation_key still wins. + correlation_key = request_data.get("correlation_key") or request_data.get("litellm_call_id") + if correlation_key: + payload["correlation_key"] = correlation_key + return payload + + @staticmethod + def _extract_request_data( + call_details: Mapping[str, Any], + request_data: Mapping[str, Any] | None, + ) -> Mapping[str, Any]: + """Extract original request data from model_call_details for the + response moderation service envelope. + + Includes the agent's declared ``tools`` (OpenAI-format) when available + so the webhook's hallucination evaluator can compare returned tool calls + against the declared tool list. + """ + if not call_details and not request_data: + return _EMPTY_MAPPING + call_details = call_details or _EMPTY_MAPPING + request_data = request_data or _EMPTY_MAPPING + optional_params = call_details.get("optional_params") or _EMPTY_MAPPING + + # Use ``in`` rather than truthy ``or`` so an explicit empty list + # (caller declared the agent has NO tools) is forwarded as-is. + # The response moderation service uses that signal to flag tool-call + # hallucinations -- ``or`` would mask it by falling through to + # optional_params. + if "tools" in request_data: + tools = request_data["tools"] + else: + tools = optional_params.get("tools") + + # The response moderation service consumes only messages/model/tools. + # Don't forward proxy_server_request -- in litellm >=1.83 its ``body`` + # snapshot carries a UserAPIKeyAuth instance that breaks json.dumps, + # silently fail-opening the guardrail. return { "messages": call_details.get("messages"), "model": call_details.get("model"), - "proxy_server_request": RubrikLogger._sanitize_proxy_server_request( - litellm_params.get("proxy_server_request") - ), + "tools": tools, } @staticmethod def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: """Allowlist only routing fields (``url``, ``method``) when forwarding - ``proxy_server_request`` to the external Rubrik webhook, dropping - inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw + ``proxy_server_request`` to an external webhook, dropping inbound + ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw request ``body`` so proxy credentials are not exfiltrated.""" if not isinstance(proxy_server_request, dict): return proxy_server_request return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request} @staticmethod - def _resolve_model(request_data: dict[str, Any], call_details: dict[str, Any]) -> str: + def _resolve_model(request_data: Mapping[str, Any], call_details: Mapping[str, Any]) -> str: """Get the model name for the ModifyResponseException.""" response = request_data.get("response") if response and hasattr(response, "model"): @@ -334,8 +590,70 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Logging hooks --------------------------------------------------------- - async def _prepare_log_payload(self, kwargs: dict, event_type: str) -> StandardLoggingPayload | None: - """Shared logic for success and failure logging.""" + @staticmethod + def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | None = None) -> str | None: + """The id that joins a blocked request's two S3 logs by filename: the + moderation (``_blocking``) log and the failure (response) log. + + Always ``litellm_call_id``. It is assigned at request start and is + present identically in both the guardrail path (``model_call_details`` + / ``request_data``) and the failure-hook path. Unlike ``response.id`` + or ``standard_logging_object["id"]`` it is immune to the race where a + block fires before the response/logging object is populated, so the + two logs correlate for every provider (OpenAI and Anthropic alike). + """ + return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id") + + @classmethod + def _apply_correlation_id(cls, payload: dict[str, Any], source: Mapping[str, Any]) -> None: + """Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log + shares its S3 filename id with the moderation (``_blocking``) and + failure logs for the same request -- for every provider. + + ``standard_logging_object["id"]`` is the provider response id + (``response_obj.get("id", litellm_call_id)``), a ``chatcmpl-*`` value + for OpenAI, which would not correlate. ``litellm_call_id`` is assigned + at request start and is identical across all log paths. Falls back to + the existing id when ``litellm_call_id`` is somehow absent rather than + writing a null filename key. + + ``source`` may be ``model_call_details`` directly or a ``kwargs`` dict + that aliases it -- same shape either way. + """ + correlated = cls._correlation_id(source) + if correlated: + payload["id"] = correlated + + @staticmethod + def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> None: + """Prepend ``source["system"]`` onto ``payload["messages"]``. + + Builds a NEW messages list rather than mutating ``payload["messages"]`` + in place. The fallback branch of ``_prepare_block_failure_payload`` + aliases ``call_details["messages"]`` directly, so an in-place + ``list.insert(0, ...)`` would mutate the shared source dict. + + No-op if no system prompt is present. Tolerates list/dict/str + message shapes; on unexpected shape, leaves payload alone. + """ + system_prompt = source.get("system") + if not system_prompt: + return + try: + system_scaffold = {"role": "system", "content": system_prompt} + messages = payload.get("messages") + if isinstance(messages, list): + payload["messages"] = (system_scaffold, *messages) + elif isinstance(messages, (dict, str)): + payload["messages"] = (system_scaffold, messages) + except Exception as e: + verbose_logger.warning( + f"Rubrik: failed to prepend system prompt: {e}", + exc_info=True, + ) + + async def _prepare_log_payload(self, kwargs: Mapping[str, Any], event_type: str) -> StandardLoggingPayload | None: + """Shared logic for success logging (sampled).""" if random.random() > self.sampling_rate: verbose_logger.debug(f"Skipping Rubrik {event_type} logging (sampling_rate={self.sampling_rate})") return None @@ -343,59 +661,17 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # Deep-copy so mutations don't affect other callbacks sharing this object standard_logging_payload: StandardLoggingPayload = safe_deep_copy(kwargs["standard_logging_object"]) - # For Anthropic /v1/messages requests, LiteLLM creates a separate - # ModelResponse (with a generated chatcmpl-* id) for logging, which - # differs from the original Anthropic msg-* id on the response dict. - # Normalize to litellm_call_id so that the logging and tool-blocking - # endpoints see the same request identifier. - litellm_params = kwargs.get("litellm_params", {}) or {} - proxy_request = litellm_params.get("proxy_server_request", {}) or {} - url_path = urllib.parse.urlparse(proxy_request.get("url", "")).path - if url_path.endswith(_ENDPOINT_ANTHROPIC_MESSAGES): - _litellm_call_id = kwargs.get("litellm_call_id") - if _litellm_call_id: - standard_logging_payload["id"] = _litellm_call_id # type: ignore[literal-required] - - if "system" in kwargs: - system_prompt_msg_list = kwargs["system"] - try: - if system_prompt_msg_list: - system_scaffold = { - "role": "system", - "content": system_prompt_msg_list, - } - if isinstance(standard_logging_payload["messages"], list): - standard_logging_payload["messages"].insert(0, system_scaffold) - elif isinstance(standard_logging_payload["messages"], (dict, str)): - standard_logging_payload["messages"] = [ - system_scaffold, - standard_logging_payload["messages"], - ] - except Exception as e: - verbose_logger.warning( - f"Rubrik: failed to prepend system prompt: {e}", - exc_info=True, - ) + self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime + self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime return standard_logging_payload - async def _enqueue_log_event(self, kwargs: dict, event_type: str): - try: - self._ensure_periodic_flush_task() - payload = await self._prepare_log_payload(kwargs, event_type) - if payload is None: - return - - self.log_queue.append(payload) - self._enforce_max_queue_size() - - if len(self.log_queue) >= self.batch_size: - await self.flush_queue() - except Exception as e: - verbose_logger.error( - f"Rubrik {event_type} logging hook failed: {e}. Skipping logging for this event.", - exc_info=True, - ) + async def _append_and_maybe_flush(self, payload) -> None: + self._ensure_periodic_flush_task() + self.log_queue.append(payload) + self._enforce_max_queue_size() + if len(self.log_queue) >= self.batch_size: + await self.flush_queue() def _enforce_max_queue_size(self) -> None: overflow = len(self.log_queue) - self.max_queue_size @@ -415,18 +691,213 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): self._dropped_since_warning = 0 self._last_drop_warning_time = now + async def _enqueue_log_event(self, kwargs: Mapping[str, Any], event_type: str): + try: + payload = await self._prepare_log_payload(kwargs, event_type) + if payload is None: + return + await self._append_and_maybe_flush(payload) + except Exception as e: + verbose_logger.error( + f"Rubrik {event_type} logging hook failed: {e}. Skipping logging for this event.", + exc_info=True, + ) + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + # Blocked requests are logged via async_post_call_failure_hook; + # skip here to avoid double-logging the pre-block response. + if kwargs.get("_rubrik_blocked"): + verbose_logger.debug( + f"Rubrik: skipping success event for blocked request litellm_call_id={kwargs.get('litellm_call_id')}" + ) + return await self._enqueue_log_event(kwargs, "success") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + # Log regular LLM failures (timeouts, upstream errors, etc.) to Rubrik. + # NOTE: ``ModifyResponseException`` blocks are NOT routed here; they + # bypass ``Logging.async_failure_handler`` entirely and reach + # ``async_post_call_failure_hook`` instead. So there is no risk of + # double-logging a block through this path. await self._enqueue_log_event(kwargs, "failure") + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: Any, + traceback_str: str | None = None, + ) -> None: + """Log blocked requests signalled via ``ModifyResponseException`` + (prompt blocks, response/tool blocks, streaming blocks). + + Carries the stashed ``_rubrik_logging_obj``. For every other + exception we no-op; LiteLLM's standard failure plumbing handles those. + """ + if not isinstance(original_exception, ModifyResponseException): + return + + # Guard by guardrail_name so that when multiple Rubrik instances are + # registered, only the instance that raised the block handles it. + # The failure hook is called for every registered callback; without + # this check the first instance pops the stash and the originating + # instance finds None and silently skips logging. + if getattr(original_exception, "guardrail_name", None) != self.guardrail_name: + return + + logging_obj = request_data.pop("_rubrik_logging_obj", None) + if logging_obj is None: + # Legitimate when a non-Rubrik guardrail raised the block; + # problematic if Rubrik did and the stash was lost (e.g. + # ``_stash_block_context`` ran with ``logging_obj=None``). Either + # way we cannot build the payload. + verbose_logger.warning( + "Rubrik: block exception without stashed logging_obj. " + f"litellm_call_id={request_data.get('litellm_call_id')}, " + f"model={request_data.get('model')}, " + f"user_id={getattr(user_api_key_dict, 'user_id', None)}, " + f"raising_guardrail=" + f"{getattr(original_exception, 'guardrail_name', None)}" + ) + return + + call_id: str | None = None + await self._build_and_enqueue_block_event(logging_obj, original_exception, call_id) + + async def _build_and_enqueue_block_event( + self, + logging_obj: "LiteLLMLoggingObj", + exception: "ModifyResponseException", + call_id: str | None, + ) -> None: + try: + call_details = logging_obj.model_call_details + # Do NOT pop "_rubrik_blocked" here. The deferred success-handler + # task may still be iterating callbacks, and popping mid-iteration + # (between two awaited callback invocations) would cause this + # plugin's success-event callback to read the flag as absent and + # log the pre-block response -- the exact bug this hook exists to + # prevent. The flag dies with model_call_details when the request + # completes; there's nothing to clean up. + call_id = call_details.get("litellm_call_id") + payload = self._prepare_block_failure_payload(logging_obj, exception) + except (AttributeError, KeyError, TypeError) as e: + verbose_logger.error( + f"Rubrik: failed to build blocked-tool payload for " + f"litellm_call_id={call_id}: {e}. Event will NOT be logged.", + exc_info=True, + ) + return + + try: + await self._append_and_maybe_flush(payload) + except Exception as e: + verbose_logger.error( + f"Rubrik: failed to enqueue blocked-tool event for litellm_call_id={call_id}: {e}.", + exc_info=True, + ) + + def _prepare_block_failure_payload( + self, + logging_obj: "LiteLLMLoggingObj", + exception: "ModifyResponseException", + ) -> StandardLoggingPayload: + """Build a failure-style payload using the exception text as response. + + Blocked-tool events are security-relevant and **bypass sampling**: + every block is logged. + + The deferred success-handler runs as a separately-scheduled task and + races with this hook, so ``standard_logging_object`` on + ``model_call_details`` may not yet be populated. If present we reuse + it; otherwise we fall back to a best-effort payload built from the + fields available at block time. + + For prompt blocks the LLM is never called, so ``standard_logging_object`` + is never populated. The fallback therefore must carry enough fields to + pass the log processor's ``LogEntry`` schema (``BaseLogEntry`` requires + ``metadata``, ``model_id``, ``model_group``, ``model_parameters``, + ``startTime``, ``endTime``, and ``completionStartTime``). Without a + parseable payload the log processor discards the entry with a parse + error and no session is created, so prompt-moderation violations are + silently dropped even though the ``_prompt_moderation.json`` forensic + log is written correctly. + + Field sourcing for the fallback path: + - ``model`` / ``model_group``: ``call_details["model"]`` -- this is the + model-group name (e.g. "gpt-4o") set by the proxy before the guardrail + fires. The router writes ``metadata["model_group"]`` only inside + ``acompletion()``, which hasn't run yet for a prompt block. + - ``model_id``: not available before the LLM returns hidden_params; + defaults to empty string. + - ``user_api_key_hash``: ``call_details["metadata"]["user_api_key"]`` -- + the hashed token written by ``add_user_information_to_request_data`` + before ``pre_call_hook`` fires. + - time fields: ``call_details["start_time"]`` reused for all three; + end/completion times are meaningless for a prompt block. + """ + call_details = logging_obj.model_call_details + exception_text = f"{type(exception).__name__}: {exception.message}" + + base = call_details.get("standard_logging_object") + if base is not None: + payload: dict = safe_deep_copy(base) + else: + verbose_logger.debug( + "Rubrik: standard_logging_object not yet on model_call_details " + f"for litellm_call_id={call_details.get('litellm_call_id')}; " + "using best-effort fallback payload." + ) + payload = self._build_fallback_payload(call_details) + + payload["response"] = exception_text + + # Pin the correlation key to litellm_call_id so this failure log shares + # its S3 filename id with the moderation (``_blocking``) log for the + # same request. The copied ``standard_logging_object["id"]`` is + # ``response_obj.get("id", litellm_call_id)`` -- a provider ``chatcmpl-*`` + # value for OpenAI -- which would not correlate; overwrite it. + payload["id"] = self._correlation_id(call_details) or f"chatcmpl-{uuid.uuid4()}" + self._prepend_system_prompt(payload, call_details) + + return payload # type: ignore[return-value] + + @staticmethod + def _build_fallback_payload(call_details: Mapping[str, Any]) -> dict[str, Any]: + _metadata: Mapping[str, Any] = call_details.get("metadata") or _EMPTY_MAPPING + # Convert datetime to a Unix float so json.dumps can serialize it. + # httpx's json= parameter uses stdlib json.dumps with no custom encoder. + _raw_start = call_details.get("start_time") + _start = _raw_start.timestamp() if _raw_start is not None else None + return { + "id": call_details.get("litellm_call_id"), + "model": call_details.get("model") or "", + # model_group is set by the router inside acompletion(), which + # hasn't run for a prompt block; use the model name instead. + "model_group": call_details.get("model") or "", + # model_id comes from response.hidden_params -- unavailable here. + "model_id": "", + "model_parameters": ModelParamHelper.get_standard_logging_model_parameters( + call_details.get("optional_params") or _EMPTY_MAPPING # pyright: ignore[reportArgumentType] # helper only reads the mapping + ), + "startTime": _start, + "endTime": _start, + "completionStartTime": _start, + "messages": call_details.get("messages") or (), + "metadata": { + # "user_api_key" is the hashed token written by + # add_user_information_to_request_data before guardrails fire. + "user_api_key_hash": _metadata.get("user_api_key_hash") or _metadata.get("user_api_key") or "", + }, + "status": "failure", + } + # -- Batch logging --------------------------------------------------------- async def _log_batch_to_rubrik(self, data): - # NOTE: this method intentionally re-raises on failure so the parent - # CustomBatchLogger.flush_queue keeps the unsent events in the queue - # for the next flush attempt instead of silently dropping them. + # NOTE: this method intentionally re-raises on failure so flush_queue + # can preserve the unsent events for the next flush attempt instead of + # silently dropping them. try: response = await self.async_httpx_client.post( url=self.logging_endpoint, @@ -452,10 +923,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): if not self.log_queue: return - log_queue_snapshot = list(self.log_queue) - verbose_logger.debug("Rubrik: Flushing batch of %s events", len(log_queue_snapshot)) await self._log_batch_to_rubrik( - data=log_queue_snapshot, + data=self.log_queue, ) async def flush_queue(self): @@ -463,8 +932,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Overrides the base implementation so the same snapshot drives both the HTTP send and the queue truncation. This avoids the subtle - coupling where the base class captures `len(self.log_queue)` - separately from the snapshot taken inside `async_send_batch`, + coupling where the base class captures ``len(self.log_queue)`` + separately from the snapshot taken inside ``async_send_batch``, which could otherwise drift in a future refactor and cause duplicate deliveries to Rubrik. """ @@ -485,70 +954,141 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): del self.log_queue[: len(snapshot)] self.last_flush_time = time.time() - # -- Tool blocking service ------------------------------------------------- + # -- Webhook services ------------------------------------------------------ - async def _post_to_tool_blocking_service( + async def _post_json(self, endpoint: str, payload: Mapping[str, Any], service_name: str) -> Mapping[str, Any]: + """POST ``payload`` to a Rubrik webhook and return its dict response. + + Raises: + Exception: If the service is unavailable or returns an error. + TypeError: If the response JSON is not a dict. + """ + verbose_logger.debug(f"Sending request to {service_name}: {endpoint}") + http_response = await self.moderation_client.post( + endpoint, + json=payload, + headers=self._headers, + ) + http_response.raise_for_status() + result = http_response.json() + if not isinstance(result, dict): + raise TypeError( + f"{service_name} returned non-dict JSON " + f"({type(result).__name__}); expected OpenAI chat completion " + "shape or empty object." + ) + return result + + async def _post_to_response_moderation_endpoint( self, - response_data: dict[str, Any], - request_data: dict[str, Any], - ) -> dict[str, Any]: - """Post a payload to the tool blocking service and return the response. + response_data: Mapping[str, Any], + request_data: Mapping[str, Any], + ) -> Mapping[str, Any]: + """Post the ``{request, response}`` envelope to the after_completion + webhook and return its (possibly rewritten) response. Args: response_data: The OpenAI-formatted response payload to send. request_data: Original LLM request data to include alongside the response for additional context. Empty dict if unavailable. - - Raises: - Exception: If the service is unavailable or returns an error. """ - envelope = { - "request": request_data, - "response": response_data, - } - verbose_logger.debug(f"Sending request to tool blocking service: {self.tool_blocking_endpoint}") - http_response = await self.tool_blocking_client.post( - self.tool_blocking_endpoint, - json=envelope, - headers=self._headers, + envelope = {"request": request_data, "response": response_data} + return await self._post_json( + self.response_moderation_endpoint, + envelope, + "Response moderation service", ) - http_response.raise_for_status() - result: dict[str, Any] = http_response.json() - return result + + async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Post a bare OpenAI request to the before_prompt webhook. + + Returns ``{}`` (passthrough) or a synthetic chat.completion (block). + """ + return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service") @staticmethod - def _extract_blocked_tools( - service_response: dict[str, Any], - all_tool_calls: list[ChatCompletionMessageToolCall], - ) -> str | None: - """Return the blocking explanation if any tool calls were blocked. + def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None: + """Return the refusal text when the prompt was blocked, else None. - Compares the service response (which contains only allowed tools) against - the full set of tool calls. Returns ``None`` if all tools are allowed, or - the explanation string (prefixed with newlines) otherwise. + The before_prompt webhook returns ``{}`` (passthrough) or a synthetic + chat.completion whose ``choices[0].message.content`` is the refusal + explanation. + """ + choices = service_response.get("choices") + if not choices: + return None + message = choices[0].get("message") or _EMPTY_MAPPING + content = message.get("content") + return content or "Request blocked by policy." + + @staticmethod + def _extract_response_block( + service_response: Mapping[str, Any], + all_tool_calls: Sequence[ChatCompletionMessageToolCall], + sent_content: str, + ) -> BlockedResponseResult | None: + """Detect whether the webhook moderated the response text or tool calls. + + The after_completion webhook rewrites the response in place with no + explicit "blocked" flag, so we infer a block by diffing what we sent + against what came back: + + - Tool block: a tool call we sent is absent from the returned (allowed) + set. + - Text block: the returned content was REPLACED wholesale (a text + violation), as opposed to having a tool-block explanation APPENDED to + the original content. We tell them apart with ``startswith``, which + mirrors the webhook's own append-vs-replace behavior. + + Returns None when nothing was moderated. A text block supersedes a tool + block (mirroring the webhook, which drops tool calls on a text block). Expects service_response in OpenAI chat completion format: {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} """ - choices = service_response.get("choices", []) + choices = service_response.get("choices") or () if not choices: - raise _MalformedToolBlockingResponseError("Tool blocking service returned empty response") + raise _MalformedToolBlockingResponseError("Response moderation service returned empty response") - message = choices[0].get("message", {}) - returned_tool_calls = message.get("tool_calls") or [] - blocking_explanation = message.get("content", "") + message = choices[0].get("message") or _EMPTY_MAPPING + returned_tool_calls = message.get("tool_calls") or () + returned_content = message.get("content") or "" - allowed_id_counts: Counter = Counter( - tc["id"] for tc in returned_tool_calls if isinstance(tc, dict) and tc.get("id") - ) - required_id_counts: Counter = Counter(tc.id for tc in all_tool_calls if tc.id) - - all_allowed = len(returned_tool_calls) >= len(all_tool_calls) and all( - allowed_id_counts.get(tc_id, 0) >= count for tc_id, count in required_id_counts.items() + # Use Counter so duplicate IDs are handled correctly: if the model + # emits two calls with the same ID (one allowed, one prohibited) and + # the service returns only the allowed one, a set-based check would + # miss the block. Counter preserves multiplicity. + returned_id_counts: Counter[str] = Counter(tc["id"] for tc in returned_tool_calls if tc.get("id")) + required_id_counts: Counter[str] = Counter(tc.id for tc in all_tool_calls if tc.id) + # Cardinality check catches ID-less tool calls (not counted in + # required_id_counts because tc.id is falsy); Counter check catches + # duplicate-ID attacks where one occurrence is silently removed. + tools_blocked = len(returned_tool_calls) < len(all_tool_calls) or not all( + returned_id_counts.get(tc_id, 0) >= count for tc_id, count in required_id_counts.items() ) - if all_allowed: - return None + # The webhook either replaces content wholesale (text block) or appends + # a tool-block explanation to the original text. ``appended`` tells the + # two apart, and is reused below to recover just the explanation. A text + # block requires there to have been assistant text to block. + # Use the documented ``\n\n`` separator to distinguish a tool-block + # append from a text replacement that shares the original as a prefix. + # Without the separator, a replacement like "Hello, blocked." where the + # original was "Hello" would be classified as an append (not a text + # block) and silently pass through to the client. + appended = bool(sent_content) and returned_content.startswith(f"{sent_content}\n\n") + text_blocked = bool(sent_content) and returned_content != sent_content and not appended - explanation = blocking_explanation or "Tool call blocked by policy." - return f"\n\n{explanation}" + if text_blocked: + return BlockedResponseResult(explanation=returned_content or "Response blocked by policy.") + + if tools_blocked: + if appended: + # Recover just the appended explanation: drop the original text + # and the leading separator the webhook inserted before it. + explanation = returned_content[len(sent_content) :].lstrip("\n") + else: + explanation = returned_content + return BlockedResponseResult(explanation=explanation or "Tool call blocked by policy.") + + return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py index 4ad29bbeae8..2f2228ae312 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py @@ -10,6 +10,18 @@ if TYPE_CHECKING: def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RubrikLogger: + """Create and register a RubrikLogger instance. + + The ``mode`` field in the guardrail config controls which surfaces are + moderated: + - ``pre_call`` (or a mode that includes it): prompt moderation via the + ``/v1/before_prompt/openai/v1`` webhook. + - ``post_call`` (the default when ``mode`` is omitted): response and tool + call moderation via the ``/v1/after_completion/openai/v1`` webhook. + + Both hooks are active when ``mode`` covers both ``pre_call`` and + ``post_call``. + """ import litellm rubrik_callback = RubrikLogger( diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index 922d2fe8a15..7f589dc15bf 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -1,8 +1,8 @@ """ Tests for the Rubrik LiteLLM plugin. -Covers initialization, apply_guardrail tool blocking (all allowed, all blocked, -partial blocking, fail-open), batch logging, and Anthropic format handling. +Covers initialization, apply_guardrail (prompt moderation + response/tool +blocking), batch logging, and Anthropic format handling. """ import os @@ -13,8 +13,10 @@ import httpx import pytest from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.integrations.rubrik import RubrikLogger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.integrations.rubrik import ( + RubrikLogger, + _MalformedToolBlockingResponseError, +) from tests.test_litellm.integrations.rubrik_test_helpers import ( make_inputs_with_tools, @@ -50,19 +52,19 @@ class TestInitialization: with patch("asyncio.create_task", Mock()): handler = RubrikLogger() assert ( - handler.tool_blocking_endpoint + handler.response_moderation_endpoint == "http://localhost:8080/v1/after_completion/openai/v1" ) assert handler.logging_endpoint == "http://localhost:8080/v1/litellm/batch" assert handler.key == "test-api-key" - assert isinstance(handler.tool_blocking_client, AsyncHTTPHandler) + assert handler.moderation_client is not None def test_init_with_constructor_params(self): with patch("asyncio.create_task", Mock()): handler = RubrikLogger(api_key="ctor-key", api_base="http://ctor-host:9090") assert handler.key == "ctor-key" assert ( - handler.tool_blocking_endpoint + handler.response_moderation_endpoint == "http://ctor-host:9090/v1/after_completion/openai/v1" ) @@ -82,7 +84,7 @@ class TestInitialization: with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080/"}): with patch("asyncio.create_task", Mock()): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://localhost:8080/v1/after_completion/openai/v1" ) @@ -90,13 +92,13 @@ class TestInitialization: with patch("asyncio.create_task", Mock()): with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v1"}): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://host/v1/after_completion/openai/v1" ) with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v11"}): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://host/v11/v1/after_completion/openai/v1" ) @@ -155,10 +157,10 @@ class TestInitialization: # Do NOT patch asyncio.create_task — the real call should be # guarded and fall back gracefully when there is no event loop. handler = RubrikLogger() - assert handler.tool_blocking_endpoint.startswith("http://localhost:8080") + assert handler.response_moderation_endpoint.startswith("http://localhost:8080") # Without a running loop at init, the periodic flush task should be # deferred so batches still get drained once a log event arrives. - assert handler._flush_task is None + assert handler._periodic_flush_task is None @pytest.mark.asyncio async def test_periodic_flush_task_started_lazily_on_first_log(self, mock_env): @@ -170,7 +172,7 @@ class TestInitialization: side_effect=RuntimeError("no running loop"), ): handler = RubrikLogger() - assert handler._flush_task is None + assert handler._periodic_flush_task is None kwargs = { "standard_logging_object": { @@ -183,8 +185,8 @@ class TestInitialization: with patch.object(handler, "_log_batch_to_rubrik", AsyncMock()): await handler.async_log_success_event(kwargs, None, None, None) - assert handler._flush_task is not None - handler._flush_task.cancel() + assert handler._periodic_flush_task is not None + handler._periodic_flush_task.cancel() def test_event_hook_defaults_to_post_call_when_none_passed(self, mock_env): """`initialize_guardrail` always passes ``event_hook=litellm_params.mode`` @@ -204,14 +206,13 @@ class TestInitialization: handler = RubrikLogger(event_hook=GuardrailEventHooks.pre_call) assert handler.event_hook == GuardrailEventHooks.pre_call - def test_default_on_defaults_to_true_when_none_passed(self, mock_env): - """`initialize_guardrail` always passes ``default_on=litellm_params.default_on`` - (which is ``None`` when the user omits ``default_on``). The logger must - coerce a None ``default_on`` to True, otherwise ``should_run_guardrail`` - (which checks ``self.default_on is True``) silently skips the guardrail.""" + def test_default_on_defaults_to_false_when_none_passed(self, mock_env): + """Follows the standard litellm pattern: omitted ``default_on`` resolves + to ``False`` (off by default). Users must explicitly set + ``default_on: true`` to enable the guardrail for all requests.""" with patch("asyncio.create_task", Mock()): handler = RubrikLogger(default_on=None) - assert handler.default_on is True + assert handler.default_on is False def test_explicit_default_on_false_preserved(self, mock_env): """A user explicitly setting ``default_on: false`` in their guardrail @@ -421,7 +422,7 @@ class TestBatchLogging: ) assert len(handler.log_queue) == 1 msgs = handler.log_queue[0]["messages"] - assert isinstance(msgs, list) + assert isinstance(msgs, tuple) assert msgs[0]["role"] == "system" assert msgs[1] == {"role": "user", "content": "hi"} @@ -444,7 +445,10 @@ class TestBatchLogging: ) assert handler.log_queue[0]["id"] == "litellm-call-123" - async def test_non_anthropic_id_unchanged(self, handler): + async def test_litellm_call_id_always_used_as_correlation_key(self, handler): + """The merged plugin always uses litellm_call_id as the log ID for all + providers (not just Anthropic) so that logs correlate with the + moderation (_blocking) and failure logs for the same request.""" kwargs = { "standard_logging_object": { "id": "chatcmpl-original", @@ -461,7 +465,7 @@ class TestBatchLogging: await handler.async_log_success_event( kwargs=kwargs, response_obj=None, start_time=None, end_time=None ) - assert handler.log_queue[0]["id"] == "chatcmpl-original" + assert handler.log_queue[0]["id"] == "litellm-call-123" async def test_payload_deep_copied_not_mutated(self, handler): """Verify the shared standard_logging_object is not mutated.""" @@ -536,7 +540,7 @@ class TestApplyGuardrail: tc2 = make_tool_call_dict("call_2", "get_time") inputs = make_inputs_with_tools([tc1, tc2]) - handler.tool_blocking_client = _echo_service() + handler.moderation_client = _echo_service() result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -548,7 +552,7 @@ class TestApplyGuardrail: tc2 = make_tool_call_dict("call_2", "drop_database") inputs = make_inputs_with_tools([tc1, tc2]) - handler.tool_blocking_client = _mock_service_response( + handler.moderation_client = _mock_service_response( { "choices": [ { @@ -594,7 +598,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client with pytest.raises(ModifyResponseException): await handler.apply_guardrail( @@ -607,7 +611,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -618,7 +622,7 @@ class TestApplyGuardrail: tc1 = make_tool_call_dict("call_1", "test_tool") inputs = make_inputs_with_tools([tc1]) - handler.tool_blocking_client = _mock_service_response({"choices": []}) + handler.moderation_client = _mock_service_response({"choices": []}) result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -641,7 +645,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -675,7 +679,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client logging_obj = Mock() logging_obj.model_call_details = { @@ -697,7 +701,10 @@ class TestApplyGuardrail: assert req["model"] == "gpt-4" assert req["messages"] == [{"role": "user", "content": "hi"}] - async def test_proxy_server_request_headers_stripped(self, handler): + async def test_proxy_server_request_not_forwarded(self, handler): + """proxy_server_request is intentionally NOT included in the request + envelope: in litellm >=1.83 its ``body`` carries a UserAPIKeyAuth + instance that breaks json.dumps, silently fail-opening the guardrail.""" tc = make_tool_call_dict("call_1", "test_tool") inputs = make_inputs_with_tools([tc]) @@ -712,7 +719,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client logging_obj = Mock() logging_obj.model_call_details = { @@ -739,8 +746,8 @@ class TestApplyGuardrail: logging_obj=logging_obj, ) - forwarded = captured_payload["request"]["proxy_server_request"] - assert forwarded == {"url": "/chat/completions", "method": "POST"} + # proxy_server_request is deliberately excluded from the forwarded envelope + assert "proxy_server_request" not in captured_payload["request"] # -- Anthropic format ---------------------------------------------------------- @@ -760,7 +767,7 @@ class TestApplyGuardrailAnthropicFormat: ) inputs = make_inputs_with_tools([tc], texts=["I'll check the weather."]) - handler.tool_blocking_client = _echo_service() + handler.moderation_client = _echo_service() result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -771,7 +778,7 @@ class TestApplyGuardrailAnthropicFormat: tc = make_tool_call_dict("toolu_123", "dangerous_tool", '{"arg": "value"}') inputs = make_inputs_with_tools([tc]) - handler.tool_blocking_client = _mock_service_response( + handler.moderation_client = _mock_service_response( { "choices": [ { @@ -790,21 +797,21 @@ class TestApplyGuardrailAnthropicFormat: inputs=inputs, request_data={}, input_type="response" ) - async def test_text_only_response_no_blocking(self, handler): + async def test_text_only_response_sent_to_moderation(self, handler): + """Text-only responses (no tool calls) are sent to the response + moderation service to check the assistant's text content.""" from litellm.types.utils import GenericGuardrailAPIInputs inputs = GenericGuardrailAPIInputs(texts=["Hello! I'm Claude."]) - mock_client = AsyncMock() - mock_client.post = AsyncMock() - handler.tool_blocking_client = mock_client + # Service allows the response (returns the content unchanged) + handler.moderation_client = _echo_service() result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" ) assert result is inputs - mock_client.post.assert_not_called() async def test_service_failure_preserves_tools(self, handler): tc = make_tool_call_dict("toolu_123", "get_weather", '{"location": "SF"}') @@ -812,7 +819,7 @@ class TestApplyGuardrailAnthropicFormat: mock_client = AsyncMock() mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -850,10 +857,13 @@ class TestNormalizeToolCalls: RubrikLogger._normalize_tool_calls(["not_a_tool_call"]) -# -- Extract blocked tools ----------------------------------------------------- +# -- Extract response block ---------------------------------------------------- -class TestExtractBlockedTools: +class TestExtractResponseBlock: + """Tests for _extract_response_block, which replaces the upstream + _extract_blocked_tools and handles both text blocks and tool blocks.""" + def test_all_allowed_returns_none(self): from litellm.types.utils import ChatCompletionMessageToolCall, Function @@ -870,7 +880,7 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + result = RubrikLogger._extract_response_block(service_resp, [tc], "") assert result is None def test_some_blocked_returns_explanation(self): @@ -896,13 +906,13 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + result = RubrikLogger._extract_response_block(service_resp, [tc1, tc2], "") assert result is not None - assert "blocked fn2" in result + assert "blocked fn2" in result.explanation def test_empty_choices_raises(self): - with pytest.raises(Exception, match="empty response"): - RubrikLogger._extract_blocked_tools({"choices": []}, []) + with pytest.raises(_MalformedToolBlockingResponseError): + RubrikLogger._extract_response_block({"choices": []}, [], "") def test_null_tool_calls_treated_as_all_blocked(self): from litellm.types.utils import ChatCompletionMessageToolCall, Function @@ -920,36 +930,55 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + result = RubrikLogger._extract_response_block(service_resp, [tc], "") assert result is not None - assert "blocked everything" in result + assert "blocked everything" in result.explanation - def test_duplicate_ids_block_when_only_one_returned(self): + def test_text_block_detected(self): + """When the service replaces the response text wholesale, it's a text block.""" from litellm.types.utils import ChatCompletionMessageToolCall, Function - tc1 = ChatCompletionMessageToolCall( - id="call_dup", - type="function", - function=Function(name="fn", arguments="{}"), - ) - tc2 = ChatCompletionMessageToolCall( - id="call_dup", - type="function", - function=Function(name="fn", arguments="{}"), - ) service_resp = { "choices": [ { "message": { - "tool_calls": [{"id": "call_dup"}], - "content": "blocked duplicate", + "tool_calls": [], + "content": "This content violates policy.", } } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + result = RubrikLogger._extract_response_block( + service_resp, [], "Original assistant text." + ) assert result is not None - assert "blocked duplicate" in result + assert "violates policy" in result.explanation + + def test_tool_block_with_appended_explanation(self): + """When the service appends an explanation to the original text, only the + appended part is returned as the explanation.""" + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="fn", arguments="{}") + ) + original_text = "Here is my response." + appended_explanation = "Tool call was blocked." + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [], + "content": original_text + "\n\n" + appended_explanation, + } + } + ] + } + result = RubrikLogger._extract_response_block( + service_resp, [tc], original_text + ) + assert result is not None + assert appended_explanation in result.explanation # -- Sanitize proxy server request ------------------------------------------- @@ -1010,3 +1039,796 @@ class TestResolveModel: {"response": response}, {"model": "fallback"} ) assert result == "unknown" + + +# -- Additional Initialization edge cases ------------------------------------ + + +class TestInitializationEdgeCases: + def test_batch_size_zero_uses_default(self): + """RUBRIK_BATCH_SIZE=0 must warn and fall back to the default.""" + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "0"}, + ): + h = RubrikLogger() + # Should use default, not 0 + assert h.batch_size > 0 + + def test_batch_size_negative_uses_default(self): + """RUBRIK_BATCH_SIZE=-1 must warn and fall back to the default.""" + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "-5"}, + ): + h = RubrikLogger() + assert h.batch_size > 0 + + +# -- aclose() ----------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestAclose: + async def test_aclose_cancels_task_does_not_close_shared_client(self, mock_env): + """aclose() cancels the periodic flush task but does NOT close the shared + moderation_client — closing a shared cached client would break other + RubrikLogger instances that share the same connection pool.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + + mock_task = Mock() + mock_task.cancel = Mock() + handler._periodic_flush_task = mock_task + + handler.moderation_client = AsyncMock() + handler.moderation_client.close = AsyncMock() + + await handler.aclose() + + mock_task.cancel.assert_called_once() + handler.moderation_client.close.assert_not_awaited() + + async def test_aclose_with_none_task_does_not_close_client(self, mock_env): + """aclose() with no flush task still does not close the shared client.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + + handler._periodic_flush_task = None + handler.moderation_client = AsyncMock() + handler.moderation_client.close = AsyncMock() + + await handler.aclose() + + handler.moderation_client.close.assert_not_awaited() + + +# -- apply_guardrail edge cases ----------------------------------------------- + + +@pytest.mark.asyncio +class TestApplyGuardrailEdgeCases: + async def test_unknown_input_type_returns_inputs_unchanged(self, handler): + """When input_type is not 'request' or 'response', inputs are returned as-is.""" + inputs = make_inputs_with_tools([make_tool_call_dict("call_1", "tool")]) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="unknown" + ) + assert result is inputs + + async def test_response_with_no_texts_and_no_tool_calls_returns_inputs(self, handler): + """_moderate_response early-returns when both texts and tool_calls are empty.""" + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs() + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_moderate_response_empty_call_details_emits_warning(self, handler): + """When logging_obj is present but model_call_details is empty, a warning is + logged and moderation proceeds (fail-open on HTTP error).""" + tc = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc]) + + logging_obj = Mock() + logging_obj.model_call_details = {} + + handler.moderation_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + assert result is inputs + + +# -- Prompt moderation -------------------------------------------------------- + + +@pytest.mark.asyncio +class TestPromptModeration: + async def test_prompt_moderation_passthrough(self, handler): + """Webhook returns {} (empty dict) → inputs returned unchanged.""" + inputs = {"structured_messages": [{"role": "user", "content": "Hello"}]} + + handler.moderation_client = _mock_service_response({}) + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_prompt_moderation_blocked_raises(self, handler): + """Webhook returns synthetic chat.completion → raises ModifyResponseException.""" + inputs = { + "structured_messages": [{"role": "user", "content": "Harmful prompt"}], + "model": "gpt-4", + } + + handler.moderation_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "This request violates our policy.", + } + } + ] + } + ) + + with pytest.raises(ModifyResponseException) as exc_info: + await handler.apply_guardrail( + inputs=inputs, request_data={"model": "gpt-4"}, input_type="request" + ) + assert "violates our policy" in exc_info.value.message + + async def test_prompt_moderation_no_messages_skips_moderation(self, handler): + """When structured_messages is absent/empty, moderation is skipped.""" + inputs = {"model": "gpt-4"} + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_prompt_moderation_stashes_logging_obj_on_block(self, handler): + """On a prompt block, _stash_block_context must set the blocked flag.""" + inputs = { + "structured_messages": [{"role": "user", "content": "bad prompt"}], + } + + handler.moderation_client = _mock_service_response( + { + "choices": [ + {"message": {"role": "assistant", "content": "Blocked."}} + ] + } + ) + + logging_obj = Mock() + logging_obj.model_call_details = {} + request_data: dict = {} + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=logging_obj, + ) + + assert logging_obj.model_call_details.get("_rubrik_blocked") is True + assert request_data.get("_rubrik_logging_obj") is logging_obj + + +# -- _stash_block_context ----------------------------------------------------- + + +class TestStashBlockContext: + def test_with_non_none_logging_obj_sets_flag_and_stashes(self): + """Sets _rubrik_blocked flag and stores logging_obj on request_data.""" + logging_obj = Mock() + logging_obj.model_call_details = {} + request_data: dict = {} + + RubrikLogger._stash_block_context(logging_obj, request_data) + + assert logging_obj.model_call_details["_rubrik_blocked"] is True + assert request_data["_rubrik_logging_obj"] is logging_obj + + def test_with_none_logging_obj_stores_none_on_request_data(self): + """When logging_obj is None, stores None on request_data (logged as error).""" + request_data: dict = {"litellm_call_id": "test-id"} + + RubrikLogger._stash_block_context(None, request_data) + + assert request_data["_rubrik_logging_obj"] is None + + +# -- _normalize_tool_calls duck-typed ----------------------------------------- + + +class TestNormalizeToolCallsDuckTyped: + def test_duck_typed_object_with_id_and_function_attrs(self): + """Objects that have .id and .function attrs but are not + ChatCompletionMessageToolCall are handled by the third branch.""" + from litellm.types.utils import Function + + tc = Mock() + tc.id = "call_duck" + tc.type = "function" + tc.function = Function(name="duck_tool", arguments='{"x": 1}') + # Make isinstance(..., ChatCompletionMessageToolCall) return False + # by using a plain Mock (not a ChatCompletionMessageToolCall subclass) + + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_duck" + assert result[0].function.name == "duck_tool" + + def test_duck_typed_without_type_defaults_to_function(self): + """getattr(tc, "type", None) falls back to "function" when absent.""" + from litellm.types.utils import Function + + tc = Mock(spec=["id", "function"]) # no .type attr + tc.id = "call_no_type" + tc.function = Function(name="fn", arguments="{}") + + result = RubrikLogger._normalize_tool_calls([tc]) + assert result[0].type == "function" + + +# -- _flatten_messages_for_moderation ----------------------------------------- + + +class TestFlattenMessagesForModeration: + def test_plain_string_content_preserved(self): + messages = [{"role": "user", "content": "Hello world"}] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["role"] == "user" + assert result[0]["content"] == "Hello world" + + def test_content_list_flattened_to_string(self): + """Content as a list of parts (e.g. Anthropic multi-part) is flattened.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello from parts"}, + ], + } + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["role"] == "user" + assert "Hello from parts" in result[0]["content"] + + def test_non_dict_messages_skipped(self): + """Non-dict entries in the messages list are silently skipped.""" + messages = [ + "raw string message", + {"role": "user", "content": "valid"}, + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["content"] == "valid" + + def test_none_messages_returns_empty(self): + result = RubrikLogger._flatten_messages_for_moderation(None) + assert result == () + + def test_multiple_messages_preserved_in_order(self): + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Question?"}, + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 2 + assert result[0]["role"] == "system" + assert result[1]["role"] == "user" + + +# -- _build_prompt_moderation_payload ----------------------------------------- + + +class TestBuildPromptModerationPayload: + def test_payload_includes_tools_when_present(self): + inputs = { + "model": "gpt-4", + "structured_messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "fn"}}], + } + payload = RubrikLogger._build_prompt_moderation_payload(inputs, {}) + assert payload["tools"] == [{"type": "function", "function": {"name": "fn"}}] + + def test_payload_includes_user_when_present(self): + inputs = { + "structured_messages": [{"role": "user", "content": "hi"}], + } + request_data = {"user": "alice"} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["user"] == "alice" + + def test_payload_uses_explicit_correlation_key(self): + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + request_data = { + "correlation_key": "corr-123", + "litellm_call_id": "litellm-456", + } + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["correlation_key"] == "corr-123" + + def test_payload_falls_back_to_litellm_call_id(self): + """When correlation_key is absent, litellm_call_id is used.""" + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + request_data = {"litellm_call_id": "litellm-789"} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["correlation_key"] == "litellm-789" + + def test_payload_omits_optional_fields_when_absent(self): + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, {}) + assert "tools" not in payload + assert "user" not in payload + assert "correlation_key" not in payload + + +# -- _extract_request_data tools preference ----------------------------------- + + +class TestExtractRequestDataToolsPreference: + def test_prefers_tools_from_request_data_over_optional_params(self): + """When 'tools' key exists in request_data, it wins over optional_params.""" + call_details = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-4", + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + }, + } + request_data = { + "tools": [{"type": "function", "function": {"name": "from_request"}}] + } + result = RubrikLogger._extract_request_data(call_details, request_data) + assert result["tools"] == [ + {"type": "function", "function": {"name": "from_request"}} + ] + + def test_falls_back_to_optional_params_when_not_in_request_data(self): + call_details = { + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + } + } + result = RubrikLogger._extract_request_data(call_details, {}) + assert result["tools"] == [ + {"type": "function", "function": {"name": "from_optional"}} + ] + + def test_explicit_empty_list_in_request_data_is_forwarded(self): + """An explicit empty tools list signals 'no tools' to the moderation service.""" + call_details = { + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + } + } + request_data = {"tools": []} + result = RubrikLogger._extract_request_data(call_details, request_data) + assert result["tools"] == [] + + +# -- _extract_prompt_refusal -------------------------------------------------- + + +class TestExtractPromptRefusal: + def test_passthrough_response_returns_none(self): + """Empty dict (passthrough) → None.""" + assert RubrikLogger._extract_prompt_refusal({}) is None + + def test_no_choices_returns_none(self): + assert RubrikLogger._extract_prompt_refusal({"choices": []}) is None + + def test_block_response_returns_content(self): + service_response = { + "choices": [{"message": {"content": "Request blocked by Rubrik."}}] + } + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by Rubrik." + + def test_empty_content_falls_back_to_default_message(self): + """When content is empty string or falsy, falls back to default refusal.""" + service_response = {"choices": [{"message": {"content": ""}}]} + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by policy." + + def test_none_content_falls_back_to_default_message(self): + service_response = {"choices": [{"message": {"content": None}}]} + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by policy." + + +# -- _prepend_system_prompt exception path ------------------------------------ + + +class TestPrependSystemPromptException: + def test_exception_during_unpack_is_caught_and_logged(self): + """When an exception is raised inside _prepend_system_prompt, it is swallowed.""" + + class ExplodingList(list): + def __iter__(self): + raise RuntimeError("iteration error!") + + payload = {"messages": ExplodingList()} + source = {"system": "You are an assistant."} + + # Must not raise + RubrikLogger._prepend_system_prompt(payload, source) + + +# -- _append_and_maybe_flush batch trigger ------------------------------------ + + +@pytest.mark.asyncio +class TestAppendAndMaybeFlush: + async def test_flush_triggered_when_queue_reaches_batch_size(self, handler): + """flush_queue is called when the queue length reaches batch_size.""" + handler.batch_size = 2 + handler.flush_queue = AsyncMock() + + await handler._append_and_maybe_flush({"msg": "a"}) + handler.flush_queue.assert_not_called() + + await handler._append_and_maybe_flush({"msg": "b"}) + handler.flush_queue.assert_called_once() + + async def test_no_flush_before_batch_size(self, handler): + handler.batch_size = 5 + handler.flush_queue = AsyncMock() + + for i in range(4): + await handler._append_and_maybe_flush({"msg": str(i)}) + + handler.flush_queue.assert_not_called() + + +# -- _enqueue_log_event exception handling ------------------------------------ + + +@pytest.mark.asyncio +class TestEnqueueLogEventExceptions: + async def test_exception_from_prepare_log_payload_is_caught(self, handler): + """Exceptions raised by _prepare_log_payload are caught and logged.""" + handler._prepare_log_payload = AsyncMock( + side_effect=RuntimeError("payload error") + ) + + # Must not raise + await handler._enqueue_log_event( + {"standard_logging_object": {"messages": [], "response": ""}}, "test" + ) + assert len(handler.log_queue) == 0 + + +# -- async_log_success_event skip when _rubrik_blocked ------------------------ + + +@pytest.mark.asyncio +class TestSuccessEventBlockedSkip: + async def test_skips_enqueue_when_rubrik_blocked_flag_set(self, handler): + """When kwargs['_rubrik_blocked'] is True, the event is not enqueued.""" + kwargs = { + "_rubrik_blocked": True, + "litellm_call_id": "blocked-call-123", + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 0 + + +# -- async_post_call_failure_hook --------------------------------------------- + + +@pytest.mark.asyncio +class TestPostCallFailureHook: + async def test_non_modify_exception_returns_immediately(self, handler): + """Non-ModifyResponseException causes a no-op.""" + await handler.async_post_call_failure_hook( + request_data={"litellm_call_id": "test"}, + original_exception=ValueError("unrelated error"), + user_api_key_dict=None, + ) + assert len(handler.log_queue) == 0 + + async def test_modify_exception_without_stashed_logging_obj_emits_warning( + self, handler + ): + """ModifyResponseException with no _rubrik_logging_obj → warning, no enqueue.""" + request_data = {"litellm_call_id": "test-123", "model": "gpt-4"} + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=None, + ) + assert len(handler.log_queue) == 0 + + async def test_modify_exception_with_valid_logging_obj_enqueues_payload( + self, handler + ): + """ModifyResponseException + stashed logging_obj → builds and enqueues.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-abc", + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": { + "id": "chatcmpl-original", + "model": "gpt-4", + "response": "original", + "messages": [{"role": "user", "content": "hi"}], + }, + "metadata": {}, + } + + request_data = {"_rubrik_logging_obj": logging_obj} + exc = ModifyResponseException( + message="blocked by policy", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + handler.batch_size = 10**6 # disable auto-flush + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=None, + ) + assert len(handler.log_queue) == 1 + assert "ModifyResponseException" in handler.log_queue[0]["response"] + + async def test_logging_obj_popped_from_request_data(self, handler): + """_rubrik_logging_obj must be popped from request_data so it is not + forwarded downstream.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-pop", + "model": "gpt-4", + "messages": [], + "standard_logging_object": { + "id": "chatcmpl-pop", + "model": "gpt-4", + "response": "text", + "messages": [], + }, + "metadata": {}, + } + + request_data = {"_rubrik_logging_obj": logging_obj} + exc = ModifyResponseException( + message="popped", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + handler.batch_size = 10**6 + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=None, + ) + assert "_rubrik_logging_obj" not in request_data + + async def test_build_and_enqueue_swallows_attribute_error_from_prepare_payload( + self, handler + ): + """When _prepare_block_failure_payload raises AttributeError/KeyError/TypeError, + the error is logged and the event is silently dropped (lines 806-812).""" + logging_obj = Mock() + # Make model_call_details.get() raise TypeError + logging_obj.model_call_details = None # .get() will raise AttributeError + + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + # Must not raise + await handler._build_and_enqueue_block_event(logging_obj, exc, None) + assert len(handler.log_queue) == 0 + + async def test_build_and_enqueue_swallows_flush_exception(self, handler): + """When _append_and_maybe_flush raises, the error is logged (lines 816-817).""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-flush-err", + "model": "gpt-4", + "messages": [], + "standard_logging_object": { + "id": "id-flush-err", + "model": "gpt-4", + "response": "text", + "messages": [], + }, + "metadata": {}, + } + + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + handler._append_and_maybe_flush = AsyncMock( + side_effect=RuntimeError("flush failed") + ) + + # Must not raise + await handler._build_and_enqueue_block_event(logging_obj, exc, None) + + +# -- _prepare_block_failure_payload and _build_fallback_payload --------------- + + +class TestPrepareBlockFailurePayload: + def test_uses_standard_logging_object_when_present(self, handler): + """When standard_logging_object is on model_call_details, it is used as base.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-slo", + "model": "gpt-4", + "standard_logging_object": { + "id": "chatcmpl-original", + "model": "gpt-4", + "response": "original response", + "messages": [{"role": "user", "content": "hi"}], + }, + "metadata": {}, + } + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc) + + assert "ModifyResponseException: blocked" in payload["response"] + assert payload["id"] == "call-slo" + + def test_uses_fallback_when_standard_logging_object_absent(self, handler): + """When standard_logging_object is absent, _build_fallback_payload is used.""" + from datetime import datetime + + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-fallback", + "model": "claude-3", + "messages": [{"role": "user", "content": "question"}], + "optional_params": {"temperature": 0.5}, + "metadata": {"user_api_key_hash": "hash-abc"}, + "start_time": datetime(2024, 6, 1), + } + exc = ModifyResponseException( + message="prompt blocked", + model="claude-3", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc) + + assert payload["id"] == "call-fallback" + assert payload["model"] == "claude-3" + assert payload["model_group"] == "claude-3" + assert "ModifyResponseException: prompt blocked" in payload["response"] + assert payload["metadata"]["user_api_key_hash"] == "hash-abc" + assert payload["status"] == "failure" + + def test_fallback_payload_without_start_time(self, handler): + """_build_fallback_payload handles missing start_time gracefully.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-notime", + "model": "gpt-4", + "messages": [], + "optional_params": {}, + "metadata": {}, + } + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc) + assert payload["startTime"] is None + + +# -- async_send_batch empty queue and flush_queue edge cases ------------------ + + +@pytest.mark.asyncio +class TestQueueEdgeCases: + async def test_async_send_batch_returns_early_on_empty_queue(self, handler): + """async_send_batch is a no-op when the queue is empty.""" + handler.async_httpx_client = AsyncMock() + await handler.async_send_batch() + handler.async_httpx_client.post.assert_not_called() + + async def test_flush_queue_returns_early_when_flush_lock_is_none(self, handler): + """flush_queue is a no-op when flush_lock is None.""" + handler.flush_lock = None + handler.log_queue = [{"msg": "a"}] + handler.async_httpx_client = AsyncMock() + + await handler.flush_queue() + handler.async_httpx_client.post.assert_not_called() + + async def test_flush_queue_returns_early_when_queue_empty_inside_lock(self, handler): + """flush_queue acquires the lock then no-ops when the queue is empty.""" + handler.log_queue = [] + handler.async_httpx_client = AsyncMock() + + await handler.flush_queue() + handler.async_httpx_client.post.assert_not_called() + + +# -- _post_json non-dict response --------------------------------------------- + + +@pytest.mark.asyncio +class TestPostJson: + async def test_raises_type_error_for_list_response(self, handler): + """When the service returns a JSON array instead of a dict, TypeError is raised.""" + mock_client = AsyncMock() + mock_resp = Mock() + mock_resp.json.return_value = ["not", "a", "dict"] + mock_resp.raise_for_status = Mock() + mock_client.post = AsyncMock(return_value=mock_resp) + handler.moderation_client = mock_client + + with pytest.raises(TypeError, match="non-dict JSON"): + await handler._post_json( + handler.prompt_moderation_endpoint, {}, "Test service" + ) + + async def test_raises_type_error_for_string_response(self, handler): + """A bare string response also raises TypeError.""" + mock_client = AsyncMock() + mock_resp = Mock() + mock_resp.json.return_value = "blocked" + mock_resp.raise_for_status = Mock() + mock_client.post = AsyncMock(return_value=mock_resp) + handler.moderation_client = mock_client + + with pytest.raises(TypeError, match="non-dict JSON"): + await handler._post_json( + handler.response_moderation_endpoint, {}, "Test service" + )