feat(guardrails/rubrik): prompt moderation, response-text blocking, streaming buffer, failure logging (#35722)

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* style: ruff format rubrik.py after Counter import addition

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>
Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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(