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

[devin-ai-integration]

Mirror of #34019, authored by @seph-barker and credited to him; the diff is limited to the rubrik files, with the only delta being immutable annotations in rubrik.py (and the matching container-type assertions) so internal staging's lint gates pass. Everything below is #34019's original description.

Summary

This PR significantly extends the Rubrik guardrail integration with capabilities developed in the customer's production deployment.

New capabilities

Prompt moderation (pre_call)

• Sends the normalized prompt to /v1/before_prompt/openai/v1 before the LLM is called
• Flattens Anthropic content-parts to strings so block-content prompts are not silently passed through
• Forwards tools and correlation_key (falls back to litellm_call_id) so the webhook can write its forensic log
• Works for both OpenAI and Anthropic wire formats via litellm's normalization layer

Response text blocking (post_call)

• The after_completion moderation hook now covers assistant text in addition to tool calls
• Distinguishes text blocks (wholesale replacement) from tool-block explanations (appended) via startswith diffing

Streaming support

• streaming_end_of_stream_only = True: moderates once per stream, not once per chunk
• streaming_buffer_until_moderated = True: withholds chunks until end-of-stream moderation passes, then releases clean response or block message only (requires litellm >= feat(guardrails): buffer streamed responses until moderated (+ clean SSE on block) #31389)

Block failure logging

• async_post_call_failure_hook logs every ModifyResponseException block to the Rubrik batch endpoint
• Uses a best-effort fallback payload for prompt blocks (where no standard_logging_object exists, since the LLM was never called)
• Bypasses sampling — every security block is logged

Queue backpressure

• max_queue_size = 10_000 with drop-oldest eviction so a webhook outage cannot grow the retry queue unboundedly
• Warning throttled to once per 60 seconds with a cumulative drop count

Consistent flush

• flush_queue override snapshots once for both the HTTP send and queue drain, preventing duplicate delivery on concurrent flush calls
• _log_batch_to_rubrik now re-raises so flush_queue preserves undelivered events for the next retry

Correlation IDs

• _apply_correlation_id pins every log's id to litellm_call_id for all providers, so batch logs, moderation logs, and failure logs share the same correlation key
• _prepend_system_prompt prepends Anthropic system prompts to batch log messages so sessions are reconstructable

Miscellaneous

• Dedicated httpx.AsyncClient (moderation_client) for webhook calls with explicit pool limits, separate from the shared logging HTTP client
• get_supported_event_hooks classmethod advertising [pre_call, post_call]
• _sanitize_proxy_server_request kept as a public static method
• _MalformedToolBlockingResponseError for structurally invalid service responses; logged at CRITICAL
• Module-level rubrik_handler singleton removed (not appropriate for a library package)

Test plan

• All 66 tests pass after updating them to match the new API
• New test cases: TestExtractResponseBlock (text blocking, appended explanation), test_litellm_call_id_always_used_as_correlation_key, test_proxy_server_request_not_forwarded, test_text_only_response_sent_to_moderation
• Run: python -m pytest tests/test_litellm/integrations/test_rubrik.py -v

Examples

Prompt blocked:

{
    "id": "chatcmpl-f0488c26-b3fb-444b-a930-8d181c19eebb",
    "created": 1785778033,
    "model": "claude-sonnet-4-5",
    "object": "chat.completion",
    "system_fingerprint": null,
    "choices": [
        {
            "finish_reason": "content_filter",
            "index": 0,
            "message": {
                "content": "Prompt blocked by Rubrik Agent Cloud (Reference ID: 2fbbccb5-80b5-4853-ae06-ada58b5fa528)",
                "role": "assistant",
                "tool_calls": null,
                "function_call": null,
                "provider_specific_fields": null
            }
        }
    ],
    "usage": {
        "completion_tokens": 0,
        "prompt_tokens": 0,
        "total_tokens": 0,
        "completion_tokens_details": null,
        "prompt_tokens_details": null
    }
}

Response blocked

{
    "id": "chatcmpl-4d9dea55-df83-4c49-b2b3-f31708ac9830",
    "created": 1785778097,
    "model": "claude-sonnet-4-5",
    "object": "chat.completion",
    "system_fingerprint": null,
    "choices": [
        {
            "finish_reason": "content_filter",
            "index": 0,
            "message": {
                "content": "Response blocked by Rubrik Agent Cloud (Reference ID: cbc30b5e-582f-49c9-9b5c-67636d60b8d6)",
                "role": "assistant",
                "tool_calls": null,
                "function_call": null,
                "provider_specific_fields": null
            }
        }
    ],
    "usage": {
        "completion_tokens": 157,
        "prompt_tokens": 582,
        "total_tokens": 739,
        "completion_tokens_details": {
            "accepted_prediction_tokens": null,
            "audio_tokens": null,
            "reasoning_tokens": 0,
            "rejected_prediction_tokens": null,
            "text_tokens": 157,
            "image_tokens": null,
            "video_tokens": null
        },
        "prompt_tokens_details": {
            "audio_tokens": null,
            "cached_tokens": 0,
            "text_tokens": 582,
            "image_tokens": null,
            "video_tokens": null,
            "cache_creation_tokens": 0,
            "cache_creation_token_details": {
                "ephemeral_5m_input_tokens": 0,
                "ephemeral_1h_input_tokens": 0
            }
        },
        "cache_creation_input_tokens": 0,
        "cache_read_input_tokens": 0,
        "iterations": null,
        "inference_geo": null,
        "speed": null,
        "service_tier": null
    }
}

Streaming response blocked (does not include original, pre-moderated response)

/v1/chat/completions (OpenAI format)

{
    "error": {
        "message": "Response blocked by Rubrik Agent Cloud (Reference ID: 77947a70-f802-499e-a002-5d8d59579b7f)",
        "type": "None",
        "param": "None",
        "code": "500"
    }
}

/v1/messages (Anthropic format)

data: {"id": "msg_ff97513d-38f5-4287-8b26-c4809081aeb2", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Prompt blocked by Rubrik Agent Cloud (Reference ID: db91a3bb-ae9b-43c7-a579-d03f869a6f57)"}], "model": "claude-sonnet-4-5", "stop_reason": "end_turn", "usage": {"input_tokens": 0, "output_tokens": 0}}

Link to Devin session: https://app.devin.ai/sessions/14ad43fe8dda4ca1a74517189ae6a2c0
Requested by: @yucheng-berri

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[greptile-apps]

Greptile Summary

This PR expands the Rubrik integration from tool-call checks into prompt and response moderation, buffered streaming enforcement, correlated block logging, and bounded batch delivery.

• Adds normalized pre-call prompt moderation and post-call text/tool moderation.
• Buffers streaming output until end-of-stream moderation completes.
• Adds block-event payload construction, correlation IDs, queue backpressure, and retry-safe flushing.
• Updates Rubrik registration, tests, and static-analysis budgets.

Confidence Score: 4/5

The PR is not yet safe to merge because completion-style prompts can still reach the model without their text being sent for moderation.

The prior completion-moderation issue remains present: the fallback creates a local message from inputs["texts"], but payload construction rereads the unchanged structured_messages field and sends no prompt content. Buffered Anthropic stream blocks also continue to omit their Rubrik failure event because the stream translation consumes the block exception before the proxy failure hook runs.

Files Needing Attention: litellm/integrations/rubrik.py

Important Files Changed

Filename	Overview
litellm/integrations/rubrik.py	Adds the primary moderation, streaming, logging, queue, and correlation behavior; completion-style prompts remain absent from the webhook payload.
litellm/proxy/guardrails/guardrail_hooks/rubrik/init.py	Updates Rubrik guardrail initialization to pass the expanded configuration into the integration.
tests/test_litellm/integrations/test_rubrik.py	Substantially expands coverage for prompt/response moderation and logging, but does not exercise completion inputs containing only texts.
basedpyright-code-budget.json	Ratchets basedpyright diagnostic budgets downward to reflect the typing improvements.
ruff-strict-budget.json	Updates the strict lint budget for the changed Rubrik implementation.
type-discipline-budget.json	Ratchets the type-discipline budget to match the revised annotations.

_{Reviews (3): Last reviewed commit: "chore: credit the original author of the..." | Re-trigger Greptile}

[devin-ai-integration]

@greptileai this PR is a mirror of the already-merged #34019 onto litellm_internal_staging; the only delta is lint work (immutable annotations in rubrik.py plus ratcheted budget files), so both findings are pre-existing upstream behavior and are deferred to a follow-up PR. Please re-review with that scope in mind.

[codecov]

Codecov Report

❌ Patch coverage is 98.55072% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/integrations/rubrik.py	98.55%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!

[yucheng-berri]

@greptileai review latest head

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_rubrik_guardrail_internal_staging (19ba727) with litellm_internal_staging (7dab1ff)¹}

Footnotes

1. No successful run was found on litellm_internal_staging (c98d595) during the generation of this report, so 7dab1ff was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩