Add Rubrik as officially-supported guardrail plugin

Adds tool blocking and batch logging integration with an external Rubrik
webhook service. The plugin validates LLM tool calls against a policy
service (fail-open on errors) and batch-logs all requests/responses.

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

docs/my-website/docs/proxy/guardrails/rubrik.md

@@ -0,0 +1,140 @@
+# Rubrik Guardrail
+
+Use Rubrik's tool blocking and logging integration to validate LLM tool calls against an external policy service and batch-log all LLM requests/responses.
+
+**Key features:**
+- **Tool blocking**: Validates tool calls against an external Rubrik service after LLM completion. Blocked tool calls trigger a policy violation response.
+- **Batch logging**: Logs all LLM requests and responses to Rubrik with configurable sampling and batching.
+- **Fail-open**: If the tool blocking service is unavailable, requests are allowed through unchanged.
+
+---
+
+## Quick Start
+
+### 1. Set Environment Variables
+
+```bash
+export RUBRIK_WEBHOOK_URL="https://your-rubrik-service.example.com"
+export RUBRIK_API_KEY="your-rubrik-api-key"        # optional
+```
+
+### 2. Configure `config.yaml`
+
+```yaml
+model_list:
+  - model_name: gpt-4
+    litellm_params:
+      model: openai/gpt-4
+      api_key: os.environ/OPENAI_API_KEY
+
+guardrails:
+  - guardrail_name: "rubrik"
+    litellm_params:
+      guardrail: rubrik
+      mode: "post_call"
+      api_key: os.environ/RUBRIK_API_KEY
+      api_base: os.environ/RUBRIK_WEBHOOK_URL
+      default_on: true
+```
+
+### 3. Launch the Proxy
+
+```bash
+litellm --config config.yaml --port 4000
+```
+
+### 4. Test It
+
+```bash
+curl -X POST http://localhost:4000/chat/completions \
+  -H "Content-Type: application/json" \
+  -H "Authorization: Bearer sk-1234" \
+  -d '{
+    "model": "gpt-4",
+    "messages": [{"role": "user", "content": "What is the weather in SF?"}],
+    "tools": [
+      {
+        "type": "function",
+        "function": {
+          "name": "get_weather",
+          "description": "Get the weather for a location",
+          "parameters": {
+            "type": "object",
+            "properties": {
+              "location": {"type": "string"}
+            },
+            "required": ["location"]
+          }
+        }
+      }
+    ]
+  }'
+```
+
+---
+
+## Configuration Reference
+
+### Environment Variables
+
+| Variable | Required | Default | Description |
+|----------|----------|---------|-------------|
+| `RUBRIK_WEBHOOK_URL` | Yes (or `api_base` in config) | — | Base URL of the Rubrik webhook service |
+| `RUBRIK_API_KEY` | No | — | Bearer token for authenticating with the Rubrik service |
+| `RUBRIK_SAMPLING_RATE` | No | `1.0` | Fraction of requests to log (0.0 to 1.0). Set to `0.5` to log ~50% of requests. |
+| `RUBRIK_BATCH_SIZE` | No | `512` | Number of log entries to buffer before flushing. Logs are also flushed on a periodic interval. |
+
+### YAML Config Parameters
+
+| Parameter | Description |
+|-----------|-------------|
+| `guardrail: rubrik` | Selects the Rubrik guardrail integration |
+| `mode: "post_call"` | Run after the LLM response is received |
+| `api_key` | Rubrik API key (can use `os.environ/RUBRIK_API_KEY`) |
+| `api_base` | Rubrik webhook base URL (can use `os.environ/RUBRIK_WEBHOOK_URL`) |
+| `default_on` | When `true`, the guardrail runs on all requests without needing per-request opt-in |
+
+---
+
+## How Tool Blocking Works
+
+1. After the LLM returns a response with tool calls, the Rubrik guardrail sends them to the blocking service at `{api_base}/v1/after_completion/openai/v1`.
+2. The service evaluates each tool call against configured policies and returns the set of **allowed** tool calls.
+3. If any tool calls are blocked, the proxy returns a `200` response with the policy violation explanation instead of the original LLM response.
+4. If the blocking service is unreachable or returns an error, the guardrail **fails open** — the original response is returned unchanged.
+
+### Request/Response format
+
+The guardrail sends a JSON envelope to the blocking service:
+
+```json
+{
+  "request": {
+    "messages": [...],
+    "model": "gpt-4",
+    "proxy_server_request": {...}
+  },
+  "response": {
+    "id": "chatcmpl-...",
+    "object": "chat.completion",
+    "choices": [{
+      "message": {
+        "role": "assistant",
+        "tool_calls": [...]
+      }
+    }]
+  }
+}
+```
+
+The service should return an OpenAI chat completion format response containing only the **allowed** tool calls and an optional `content` field with the blocking explanation.
+
+---
+
+## How Batch Logging Works
+
+All LLM requests (successes and failures) are queued and sent in batches to `{api_base}/v1/litellm/batch`.
+
+- Logs are flushed when the queue reaches `RUBRIK_BATCH_SIZE` (default 512) or on a periodic interval (default 5 seconds).
+- Use `RUBRIK_SAMPLING_RATE` to reduce logging volume in high-traffic deployments.
+- For Anthropic `/v1/messages` requests, the log ID is normalized to `litellm_call_id` for consistency across tool blocking and logging.

docs/my-website/sidebars.js

@@ -90,6 +90,7 @@ const sidebars = {
           "proxy/guardrails/custom_code_guardrail",
           "proxy/guardrails/prompt_injection",
           "proxy/guardrails/tool_permission",
+          "proxy/guardrails/rubrik",
           "proxy/guardrails/zscaler_ai_guard",
           "proxy/guardrails/javelin"
         ].sort(),

litellm/integrations/rubrik.py

@@ -0,0 +1,459 @@
+"""Rubrik LiteLLM Plugin for tool blocking and batch logging."""
+
+import asyncio
+import copy
+import os
+import random
+import time
+import urllib.parse
+import uuid
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal, Optional
+
+import httpx
+from litellm._logging import verbose_logger
+from litellm.integrations.custom_batch_logger import CustomBatchLogger
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    ModifyResponseException,
+)
+from litellm.llms.custom_httpx.http_handler import (
+    get_async_httpx_client,
+    httpxSpecialProvider,
+)
+from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.utils import (
+    ChatCompletionMessageToolCall,
+    Function,
+    GenericGuardrailAPIInputs,
+    StandardLoggingPayload,
+)
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import (
+        Logging as LiteLLMLoggingObj,
+    )
+
+_ENDPOINT_ANTHROPIC_MESSAGES = "/messages"
+_WEBHOOK_PATH_TOOL_BLOCKING = "/v1/after_completion/openai/v1"
+_WEBHOOK_PATH_LOGGING_BATCH = "/v1/litellm/batch"
+
+
+@dataclass
+class BlockedToolsResult:
+    """Returned by _extract_blocked_tools when at least one tool was blocked."""
+
+    allowed_tools: list
+    explanation: str
+
+
+class RubrikLogger(CustomGuardrail, CustomBatchLogger):
+    def __init__(
+        self,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        **kwargs,
+    ):
+        self.flush_lock = asyncio.Lock()
+        kwargs.setdefault("guardrail_name", "rubrik")
+        kwargs.setdefault("event_hook", GuardrailEventHooks.post_call)
+        kwargs.setdefault("default_on", True)
+        super().__init__(
+            flush_lock=self.flush_lock,
+            **kwargs,
+        )
+
+        verbose_logger.debug("initializing rubrik logger")
+
+        self.sampling_rate = 1.0
+        rbrk_sampling_rate = os.getenv("RUBRIK_SAMPLING_RATE")
+        if rbrk_sampling_rate is not None:
+            try:
+                parsed_rate = float(rbrk_sampling_rate.strip())
+                self.sampling_rate = max(0.0, min(1.0, parsed_rate))
+                if parsed_rate != self.sampling_rate:
+                    verbose_logger.warning(
+                        f"RUBRIK_SAMPLING_RATE={parsed_rate} clamped to "
+                        f"{self.sampling_rate}"
+                    )
+            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."
+            )
+        _batch_size = os.getenv("RUBRIK_BATCH_SIZE")
+
+        if _batch_size:
+            try:
+                self.batch_size = int(_batch_size)
+            except ValueError:
+                verbose_logger.warning(
+                    f"Invalid RUBRIK_BATCH_SIZE: {_batch_size!r}, using default"
+                )
+
+        _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}"
+
+        self.async_httpx_client = get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.LoggingCallback
+        )
+
+        # Dedicated client to avoid connection pooling issues with LiteLLM's shared client
+        self.tool_blocking_client = httpx.AsyncClient(
+            timeout=httpx.Timeout(5.0, connect=2.0),
+            limits=httpx.Limits(max_connections=10, max_keepalive_connections=5),
+        )
+
+        self._headers: dict[str, str] = {"Content-Type": "application/json"}
+        if self.key:
+            self._headers["Authorization"] = f"Bearer {self.key}"
+
+        asyncio.create_task(self.periodic_flush())
+
+    async def aclose(self):
+        """Close the dedicated tool blocking HTTP client."""
+        await self.tool_blocking_client.aclose()
+
+    # -- Guardrail hook --------------------------------------------------------
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        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
+
+        tool_calls = inputs.get("tool_calls")
+        if not tool_calls:
+            return inputs
+
+        try:
+            return await self._check_tool_calls(
+                inputs, tool_calls, request_data, logging_obj
+            )
+        except ModifyResponseException:
+            raise
+        except Exception as e:
+            verbose_logger.error(
+                f"Tool blocking hook failed: {e}. "
+                "Returning original response unchanged.",
+                exc_info=True,
+            )
+            return inputs
+
+    async def _check_tool_calls(
+        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)
+
+        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
+        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)
+
+        service_response = await self._post_to_tool_blocking_service(
+            response_data, req_data
+        )
+        blocked = self._extract_blocked_tools(service_response, message_tool_calls)
+
+        if blocked:
+            model = self._resolve_model(request_data, call_details)
+            raise ModifyResponseException(
+                message=blocked.explanation,
+                model=model,
+                request_data=request_data,
+                guardrail_name=self.guardrail_name,
+            )
+
+        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__}: {tc!r}"
+                )
+        return result
+
+    @staticmethod
+    def _build_tool_call_payload(
+        tool_calls: list[ChatCompletionMessageToolCall],
+        request_id: str | None,
+    ) -> dict[str, Any]:
+        """Build a full OpenAI ChatCompletion-format dict for the blocking service."""
+        return {
+            "id": request_id or f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": "",
+            "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",
+                }
+            ],
+        }
+
+    @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 {}
+        return {
+            "messages": call_details.get("messages"),
+            "model": call_details.get("model"),
+            "proxy_server_request": litellm_params.get("proxy_server_request"),
+        }
+
+    @staticmethod
+    def _resolve_model(
+        request_data: dict[str, Any], call_details: dict[str, Any]
+    ) -> str:
+        """Get the model name for the ModifyResponseException."""
+        response = request_data.get("response")
+        if response and hasattr(response, "model"):
+            return response.model or "unknown"
+        return call_details.get("model", "unknown")
+
+    # -- Logging hooks ---------------------------------------------------------
+
+    async def _prepare_log_payload(
+        self, kwargs: dict, event_type: str
+    ) -> StandardLoggingPayload | None:
+        """Shared logic for success and failure logging."""
+        if random.random() > self.sampling_rate:
+            verbose_logger.debug(
+                f"Skipping Rubrik {event_type} logging "
+                f"(sampling_rate={self.sampling_rate})"
+            )
+            return None
+
+        # Deep-copy so mutations don't affect other callbacks sharing this object
+        standard_logging_payload: StandardLoggingPayload = copy.deepcopy(
+            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,
+                )
+
+        return standard_logging_payload
+
+    async def _enqueue_log_event(self, kwargs: dict, event_type: str):
+        try:
+            payload = await self._prepare_log_payload(kwargs, event_type)
+            if payload is None:
+                return
+
+            self.log_queue.append(payload)
+
+            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 async_log_success_event(
+        self, kwargs, response_obj, start_time, end_time
+    ):
+        await self._enqueue_log_event(kwargs, "success")
+
+    async def async_log_failure_event(
+        self, kwargs, response_obj, start_time, end_time
+    ):
+        await self._enqueue_log_event(kwargs, "failure")
+
+    # -- Batch logging ---------------------------------------------------------
+
+    async def _log_batch_to_rubrik(self, data):
+        try:
+            response = await self.async_httpx_client.post(
+                url=self.logging_endpoint,
+                json=data,
+                headers=self._headers,
+            )
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            verbose_logger.exception(
+                f"Rubrik HTTP Error: {e.response.status_code} - {e.response.text}"
+            )
+        except Exception:
+            verbose_logger.exception("Rubrik Layer Error")
+
+    async def async_send_batch(self):
+        """Handles sending batches of responses to Rubrik."""
+        if not self.log_queue:
+            return
+
+        await self._log_batch_to_rubrik(
+            data=self.log_queue,
+        )
+
+    # -- Tool blocking service -------------------------------------------------
+
+    async def _post_to_tool_blocking_service(
+        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.
+
+        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: "
+            f"{self.tool_blocking_endpoint}"
+        )
+        http_response = await self.tool_blocking_client.post(
+            self.tool_blocking_endpoint,
+            json=envelope,
+            headers=self._headers,
+        )
+        http_response.raise_for_status()
+        result: dict[str, Any] = http_response.json()
+        return result
+
+    @staticmethod
+    def _extract_blocked_tools(
+        service_response: dict[str, Any],
+        all_tool_calls: list[ChatCompletionMessageToolCall],
+    ) -> BlockedToolsResult | None:
+        """Determine whether any tool calls were blocked by the service.
+
+        Compares the service response (which contains only allowed tools) against
+        the full set of tool calls. Returns None if all tools are allowed, or a
+        BlockedToolsResult.
+
+        Expects service_response in OpenAI chat completion format:
+            {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
+        """
+        choices = service_response.get("choices", [])
+        if not choices:
+            raise Exception("Tool blocking service returned empty response")
+
+        message = choices[0].get("message", {})
+        returned_tool_calls = message.get("tool_calls", [])
+        blocking_explanation = message.get("content", "")
+
+        allowed_ids = {tc["id"] for tc in returned_tool_calls if tc.get("id")}
+        allowed_tools = [tc for tc in all_tool_calls if tc.id in allowed_ids]
+
+        if len(allowed_tools) == len(all_tool_calls):
+            return None
+
+        return BlockedToolsResult(
+            allowed_tools=allowed_tools,
+            explanation=f"\n\n{blocking_explanation}",
+        )

litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py

@@ -0,0 +1,35 @@
+"""Rubrik guardrail integration for LiteLLM."""
+
+from typing import TYPE_CHECKING
+
+from litellm.integrations.rubrik import RubrikLogger
+from litellm.types.guardrails import SupportedGuardrailIntegrations
+
+if TYPE_CHECKING:
+    from litellm.types.guardrails import Guardrail, LitellmParams
+
+
+def initialize_guardrail(
+    litellm_params: "LitellmParams", guardrail: "Guardrail"
+) -> RubrikLogger:
+    import litellm
+
+    rubrik_callback = RubrikLogger(
+        api_key=litellm_params.api_key,
+        api_base=litellm_params.api_base,
+        guardrail_name=guardrail.get("guardrail_name", ""),
+        event_hook=litellm_params.mode,
+        default_on=litellm_params.default_on,
+    )
+
+    litellm.logging_callback_manager.add_litellm_callback(rubrik_callback)
+    return rubrik_callback
+
+
+guardrail_initializer_registry = {
+    SupportedGuardrailIntegrations.RUBRIK.value: initialize_guardrail,
+}
+
+guardrail_class_registry = {
+    SupportedGuardrailIntegrations.RUBRIK.value: RubrikLogger,
+}

litellm/types/guardrails.py

@@ -84,6 +84,7 @@ class SupportedGuardrailIntegrations(Enum):
     BLOCK_CODE_EXECUTION = "block_code_execution"
     AKTO = "akto"
     MCP_JWT_SIGNER = "mcp_jwt_signer"
+    RUBRIK = "rubrik"
 
 
 class Role(Enum):

tests/test_litellm/integrations/rubrik_test_helpers.py

@@ -0,0 +1,23 @@
+"""Shared helpers for Rubrik plugin tests."""
+
+from typing import Any, Dict
+
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+
+def make_tool_call_dict(
+    tc_id: str, name: str, arguments: str = "{}"
+) -> Dict[str, Any]:
+    """Create a tool call dict matching the ChatCompletionMessageToolCall schema."""
+    return {
+        "id": tc_id,
+        "type": "function",
+        "function": {"name": name, "arguments": arguments},
+    }
+
+
+def make_inputs_with_tools(
+    tool_calls: list, texts: list | None = None
+) -> GenericGuardrailAPIInputs:
+    """Create GenericGuardrailAPIInputs with tool_calls."""
+    return GenericGuardrailAPIInputs(texts=texts or [], tool_calls=tool_calls)

tests/test_litellm/integrations/test_rubrik.py

@@ -0,0 +1,732 @@
+"""
+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.
+"""
+
+import os
+from typing import Any, Dict
+from unittest.mock import AsyncMock, Mock, patch
+
+import httpx
+import pytest
+
+from litellm.integrations.custom_guardrail import ModifyResponseException
+from litellm.integrations.rubrik import RubrikLogger
+
+from tests.test_litellm.integrations.rubrik_test_helpers import (
+    make_inputs_with_tools,
+    make_tool_call_dict,
+)
+
+
+@pytest.fixture
+def mock_env():
+    """Set up environment variables for testing."""
+    with patch.dict(
+        os.environ,
+        {
+            "RUBRIK_WEBHOOK_URL": "http://localhost:8080",
+            "RUBRIK_API_KEY": "test-api-key",
+        },
+    ):
+        yield
+
+
+@pytest.fixture
+def handler(mock_env):
+    """Create a RubrikLogger instance for testing."""
+    with patch("asyncio.create_task", Mock()):
+        return RubrikLogger()
+
+
+# -- Initialization -----------------------------------------------------------
+
+
+class TestInitialization:
+    def test_init_success(self, mock_env):
+        with patch("asyncio.create_task", Mock()):
+            handler = RubrikLogger()
+            assert (
+                handler.tool_blocking_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, httpx.AsyncClient)
+
+    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
+                == "http://ctor-host:9090/v1/after_completion/openai/v1"
+            )
+
+    def test_init_without_url(self):
+        with patch.dict(os.environ, {}, clear=True):
+            with pytest.raises(ValueError, match="Rubrik webhook URL not configured"):
+                RubrikLogger()
+
+    def test_init_without_api_key(self):
+        with patch.dict(
+            os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080"}, clear=True
+        ):
+            with patch("asyncio.create_task", Mock()):
+                assert RubrikLogger().key is None
+
+    def test_trailing_slash_removed(self):
+        with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080/"}):
+            with patch("asyncio.create_task", Mock()):
+                assert (
+                    RubrikLogger().tool_blocking_endpoint
+                    == "http://localhost:8080/v1/after_completion/openai/v1"
+                )
+
+    def test_v1_suffix_stripped_as_substring_not_charset(self):
+        with patch("asyncio.create_task", Mock()):
+            with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v1"}):
+                assert (
+                    RubrikLogger().tool_blocking_endpoint
+                    == "http://host/v1/after_completion/openai/v1"
+                )
+
+            with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v11"}):
+                assert (
+                    RubrikLogger().tool_blocking_endpoint
+                    == "http://host/v11/v1/after_completion/openai/v1"
+                )
+
+    def test_sampling_rate_fractional(self):
+        with patch("asyncio.create_task", Mock()):
+            with patch.dict(
+                os.environ,
+                {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "0.5"},
+            ):
+                assert RubrikLogger().sampling_rate == 0.5
+
+    def test_sampling_rate_invalid_ignored(self):
+        with patch("asyncio.create_task", Mock()):
+            with patch.dict(
+                os.environ,
+                {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "abc"},
+            ):
+                assert RubrikLogger().sampling_rate == 1.0
+
+    def test_sampling_rate_clamped(self):
+        with patch("asyncio.create_task", Mock()):
+            with patch.dict(
+                os.environ,
+                {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "2.0"},
+            ):
+                assert RubrikLogger().sampling_rate == 1.0
+            with patch.dict(
+                os.environ,
+                {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "-0.5"},
+            ):
+                assert RubrikLogger().sampling_rate == 0.0
+
+    def test_batch_size_invalid_ignored(self):
+        with patch("asyncio.create_task", Mock()):
+            with patch.dict(
+                os.environ,
+                {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "abc"},
+            ):
+                # Should use default without crashing
+                assert isinstance(RubrikLogger().batch_size, int)
+
+    def test_batch_size_valid(self):
+        with patch("asyncio.create_task", Mock()):
+            with patch.dict(
+                os.environ,
+                {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "256"},
+            ):
+                assert RubrikLogger().batch_size == 256
+
+    def test_headers_with_api_key(self, handler):
+        assert handler._headers["Authorization"] == "Bearer test-api-key"
+        assert handler._headers["Content-Type"] == "application/json"
+
+    def test_headers_without_api_key(self):
+        with patch.dict(
+            os.environ, {"RUBRIK_WEBHOOK_URL": "http://host"}, clear=True
+        ):
+            with patch("asyncio.create_task", Mock()):
+                h = RubrikLogger()
+                assert "Authorization" not in h._headers
+
+
+# -- Batch Logging ------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+class TestBatchLogging:
+    async def test_log_success_event_appends_to_queue(self, handler):
+        kwargs = {
+            "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) == 1
+
+    async def test_log_failure_event_appends_to_queue(self, handler):
+        kwargs = {
+            "standard_logging_object": {
+                "messages": [{"role": "user", "content": "hi"}],
+                "response": "error",
+            },
+        }
+        await handler.async_log_failure_event(
+            kwargs=kwargs, response_obj=None, start_time=None, end_time=None
+        )
+        assert len(handler.log_queue) == 1
+
+    async def test_log_success_event_sampling_skips(self, handler):
+        handler.sampling_rate = 0.0
+        kwargs = {
+            "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 def test_flush_queue_sends_batch(self, handler):
+        handler.log_queue = [{"msg": "a"}, {"msg": "b"}]
+        mock_response = Mock()
+        mock_response.status_code = 200
+        handler.async_httpx_client = AsyncMock()
+        handler.async_httpx_client.post = AsyncMock(return_value=mock_response)
+        await handler.flush_queue()
+        handler.async_httpx_client.post.assert_called_once()
+        assert len(handler.log_queue) == 0
+
+    async def test_log_batch_error_does_not_crash(self, handler):
+        handler.log_queue = [{"msg": "a"}]
+        mock_response = Mock()
+        mock_response.status_code = 500
+        mock_response.text = "Internal Server Error"
+        mock_response.raise_for_status = Mock(
+            side_effect=httpx.HTTPStatusError(
+                "err", request=Mock(), response=mock_response
+            )
+        )
+        handler.async_httpx_client = AsyncMock()
+        handler.async_httpx_client.post = AsyncMock(return_value=mock_response)
+        await handler.flush_queue()
+        assert len(handler.log_queue) == 0
+
+    async def test_system_prompt_prepended_to_messages(self, handler):
+        kwargs = {
+            "standard_logging_object": {
+                "messages": [{"role": "user", "content": "hi"}],
+                "response": "hello",
+            },
+            "system": "You are a helpful assistant.",
+        }
+        await handler.async_log_success_event(
+            kwargs=kwargs, response_obj=None, start_time=None, end_time=None
+        )
+        assert len(handler.log_queue) == 1
+        msgs = handler.log_queue[0]["messages"]
+        assert msgs[0]["role"] == "system"
+        assert msgs[0]["content"] == "You are a helpful assistant."
+
+    async def test_system_prompt_with_dict_messages(self, handler):
+        kwargs = {
+            "standard_logging_object": {
+                "messages": {"role": "user", "content": "hi"},
+                "response": "hello",
+            },
+            "system": "Be concise.",
+        }
+        await handler.async_log_success_event(
+            kwargs=kwargs, response_obj=None, start_time=None, end_time=None
+        )
+        assert len(handler.log_queue) == 1
+        msgs = handler.log_queue[0]["messages"]
+        assert isinstance(msgs, list)
+        assert msgs[0]["role"] == "system"
+        assert msgs[1] == {"role": "user", "content": "hi"}
+
+    async def test_anthropic_id_normalization(self, handler):
+        kwargs = {
+            "standard_logging_object": {
+                "id": "chatcmpl-original",
+                "messages": [{"role": "user", "content": "hi"}],
+                "response": "hello",
+            },
+            "litellm_params": {
+                "proxy_server_request": {
+                    "url": "http://proxy/v1/messages",
+                },
+            },
+            "litellm_call_id": "litellm-call-123",
+        }
+        await handler.async_log_success_event(
+            kwargs=kwargs, response_obj=None, start_time=None, end_time=None
+        )
+        assert handler.log_queue[0]["id"] == "litellm-call-123"
+
+    async def test_non_anthropic_id_unchanged(self, handler):
+        kwargs = {
+            "standard_logging_object": {
+                "id": "chatcmpl-original",
+                "messages": [{"role": "user", "content": "hi"}],
+                "response": "hello",
+            },
+            "litellm_params": {
+                "proxy_server_request": {
+                    "url": "http://proxy/v1/chat/completions",
+                },
+            },
+            "litellm_call_id": "litellm-call-123",
+        }
+        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"
+
+    async def test_payload_deep_copied_not_mutated(self, handler):
+        """Verify the shared standard_logging_object is not mutated."""
+        original_payload = {
+            "id": "original-id",
+            "messages": [{"role": "user", "content": "hi"}],
+            "response": "hello",
+        }
+        kwargs = {
+            "standard_logging_object": original_payload,
+            "system": "System prompt.",
+        }
+        await handler.async_log_success_event(
+            kwargs=kwargs, response_obj=None, start_time=None, end_time=None
+        )
+        # Original payload should NOT have been mutated
+        assert original_payload["id"] == "original-id"
+        assert len(original_payload["messages"]) == 1
+
+
+# -- Tool Blocking (apply_guardrail) ------------------------------------------
+
+
+def _mock_service_response(response_json):
+    """Create a mock tool blocking client that returns the given JSON."""
+
+    async def mock_post(*_args, **kwargs):
+        mock_resp = Mock()
+        mock_resp.json.return_value = response_json
+        mock_resp.raise_for_status = Mock()
+        return mock_resp
+
+    mock_client = AsyncMock()
+    mock_client.post = mock_post
+    return mock_client
+
+
+def _echo_service():
+    """Create a mock tool blocking client that echoes the payload back."""
+
+    async def mock_post(*_args, **kwargs):
+        mock_resp = Mock()
+        mock_resp.json.return_value = kwargs.get("json", {}).get("response", {})
+        mock_resp.raise_for_status = Mock()
+        return mock_resp
+
+    mock_client = AsyncMock()
+    mock_client.post = mock_post
+    return mock_client
+
+
+@pytest.mark.asyncio
+class TestApplyGuardrail:
+    async def test_skips_requests(self, handler):
+        inputs = make_inputs_with_tools(
+            [make_tool_call_dict("call_1", "test_tool")]
+        )
+        result = await handler.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="request"
+        )
+        assert result is inputs
+
+    async def test_no_tool_calls(self, handler):
+        from litellm.types.utils import GenericGuardrailAPIInputs
+
+        inputs = GenericGuardrailAPIInputs(texts=["hello"])
+        result = await handler.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="response"
+        )
+        assert result is inputs
+
+    async def test_all_allowed(self, handler):
+        tc1 = make_tool_call_dict("call_1", "get_weather")
+        tc2 = make_tool_call_dict("call_2", "get_time")
+        inputs = make_inputs_with_tools([tc1, tc2])
+
+        handler.tool_blocking_client = _echo_service()
+
+        result = await handler.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="response"
+        )
+        assert result is inputs
+
+    async def test_all_blocked(self, handler):
+        tc1 = make_tool_call_dict("call_1", "delete_table")
+        tc2 = make_tool_call_dict("call_2", "drop_database")
+        inputs = make_inputs_with_tools([tc1, tc2])
+
+        handler.tool_blocking_client = _mock_service_response(
+            {
+                "choices": [
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": "Tool blocked by policy",
+                            "tool_calls": [],
+                        }
+                    }
+                ],
+            }
+        )
+
+        with pytest.raises(ModifyResponseException) as exc_info:
+            await handler.apply_guardrail(
+                inputs=inputs, request_data={}, input_type="response"
+            )
+        assert "Tool blocked by policy" in exc_info.value.message
+
+    async def test_partial_blocking(self, handler):
+        tc_blocked = make_tool_call_dict("call_A", "blocked_tool")
+        tc_allowed = make_tool_call_dict("call_B", "allowed_tool")
+        inputs = make_inputs_with_tools([tc_blocked, tc_allowed])
+
+        async def mock_post(*_args, **kwargs):
+            payload = kwargs.get("json", {}).get("response", {})
+            all_tcs = payload["choices"][0]["message"]["tool_calls"]
+            allowed = [tc for tc in all_tcs if tc.get("id") == "call_B"]
+            mock_resp = Mock()
+            mock_resp.json.return_value = {
+                "choices": [
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": "blocked",
+                            "tool_calls": allowed,
+                        }
+                    }
+                ],
+            }
+            mock_resp.raise_for_status = Mock()
+            return mock_resp
+
+        mock_client = AsyncMock()
+        mock_client.post = mock_post
+        handler.tool_blocking_client = mock_client
+
+        with pytest.raises(ModifyResponseException):
+            await handler.apply_guardrail(
+                inputs=inputs, request_data={}, input_type="response"
+            )
+
+    async def test_service_failure_fail_open(self, handler):
+        tc1 = make_tool_call_dict("call_1", "test_tool")
+        inputs = make_inputs_with_tools([tc1])
+
+        mock_client = AsyncMock()
+        mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout"))
+        handler.tool_blocking_client = mock_client
+
+        result = await handler.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="response"
+        )
+        assert result is inputs
+
+    async def test_service_empty_choices_fail_open(self, handler):
+        tc1 = make_tool_call_dict("call_1", "test_tool")
+        inputs = make_inputs_with_tools([tc1])
+
+        handler.tool_blocking_client = _mock_service_response({"choices": []})
+
+        result = await handler.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="response"
+        )
+        assert result is inputs
+
+    async def test_blocking_service_payload_format(self, handler):
+        tc1 = make_tool_call_dict("call_1", "get_weather", '{"location": "SF"}')
+        tc2 = make_tool_call_dict(
+            "call_2", "send_email", '{"to": "user@example.com"}'
+        )
+        inputs = make_inputs_with_tools([tc1, tc2])
+
+        captured_payload: Dict[str, Any] = {}
+
+        async def mock_post(*_args, **kwargs):
+            captured_payload.update(kwargs.get("json", {}))
+            mock_resp = Mock()
+            mock_resp.json.return_value = captured_payload.get("response", {})
+            mock_resp.raise_for_status = Mock()
+            return mock_resp
+
+        mock_client = AsyncMock()
+        mock_client.post = mock_post
+        handler.tool_blocking_client = mock_client
+
+        await handler.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="response"
+        )
+
+        # Verify envelope structure
+        assert "request" in captured_payload
+        assert "response" in captured_payload
+
+        response_data = captured_payload["response"]
+        message = response_data["choices"][0]["message"]
+        assert message["role"] == "assistant"
+        assert len(message["tool_calls"]) == 2
+        assert message["tool_calls"][0]["id"] == "call_1"
+        assert message["tool_calls"][0]["function"]["name"] == "get_weather"
+        assert message["tool_calls"][1]["id"] == "call_2"
+        assert message["tool_calls"][1]["function"]["name"] == "send_email"
+
+    async def test_request_data_included_in_envelope(self, handler):
+        tc = make_tool_call_dict("call_1", "test_tool")
+        inputs = make_inputs_with_tools([tc])
+
+        captured_payload: Dict[str, Any] = {}
+
+        async def mock_post(*_args, **kwargs):
+            captured_payload.update(kwargs.get("json", {}))
+            mock_resp = Mock()
+            mock_resp.json.return_value = captured_payload.get("response", {})
+            mock_resp.raise_for_status = Mock()
+            return mock_resp
+
+        mock_client = AsyncMock()
+        mock_client.post = mock_post
+        handler.tool_blocking_client = mock_client
+
+        logging_obj = Mock()
+        logging_obj.model_call_details = {
+            "messages": [{"role": "user", "content": "hi"}],
+            "model": "gpt-4",
+            "litellm_params": {
+                "proxy_server_request": {"url": "/chat/completions"},
+            },
+        }
+
+        await handler.apply_guardrail(
+            inputs=inputs,
+            request_data={},
+            input_type="response",
+            logging_obj=logging_obj,
+        )
+
+        req = captured_payload["request"]
+        assert req["model"] == "gpt-4"
+        assert req["messages"] == [{"role": "user", "content": "hi"}]
+
+
+# -- Anthropic format ----------------------------------------------------------
+
+
+@pytest.mark.asyncio
+class TestApplyGuardrailAnthropicFormat:
+    """Verify blocking works correctly regardless of original provider format.
+
+    The framework converts Anthropic tool_use blocks to OpenAI-format
+    tool_calls before calling apply_guardrail.
+    """
+
+    async def test_single_tool_allowed(self, handler):
+        tc = make_tool_call_dict(
+            "toolu_123", "get_weather", '{"location": "Portland, OR"}'
+        )
+        inputs = make_inputs_with_tools([tc], texts=["I'll check the weather."])
+
+        handler.tool_blocking_client = _echo_service()
+
+        result = await handler.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="response"
+        )
+        assert result is inputs
+
+    async def test_single_tool_blocked(self, handler):
+        tc = make_tool_call_dict("toolu_123", "dangerous_tool", '{"arg": "value"}')
+        inputs = make_inputs_with_tools([tc])
+
+        handler.tool_blocking_client = _mock_service_response(
+            {
+                "choices": [
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": "blocked",
+                            "tool_calls": [],
+                        }
+                    }
+                ],
+            }
+        )
+
+        with pytest.raises(ModifyResponseException):
+            await handler.apply_guardrail(
+                inputs=inputs, request_data={}, input_type="response"
+            )
+
+    async def test_text_only_response_no_blocking(self, handler):
+        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
+
+        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"}')
+        inputs = make_inputs_with_tools([tc])
+
+        mock_client = AsyncMock()
+        mock_client.post = AsyncMock(
+            side_effect=httpx.TimeoutException("Timeout")
+        )
+        handler.tool_blocking_client = mock_client
+
+        result = await handler.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="response"
+        )
+        assert result is inputs
+
+
+# -- Normalize tool calls ------------------------------------------------------
+
+
+class TestNormalizeToolCalls:
+    def test_dict_input(self):
+        tc = make_tool_call_dict("call_1", "test", '{"a": 1}')
+        result = RubrikLogger._normalize_tool_calls([tc])
+        assert len(result) == 1
+        assert result[0].id == "call_1"
+        assert result[0].function.name == "test"
+        assert result[0].function.arguments == '{"a": 1}'
+
+    def test_typed_object_input(self):
+        from litellm.types.utils import ChatCompletionMessageToolCall, Function
+
+        tc = ChatCompletionMessageToolCall(
+            id="call_2",
+            type="function",
+            function=Function(name="fn", arguments="{}"),
+        )
+        result = RubrikLogger._normalize_tool_calls([tc])
+        assert len(result) == 1
+        assert result[0].id == "call_2"
+        assert result[0].function.name == "fn"
+
+    def test_unsupported_type_raises(self):
+        with pytest.raises(TypeError, match="Cannot normalize"):
+            RubrikLogger._normalize_tool_calls(["not_a_tool_call"])
+
+
+# -- Extract blocked tools -----------------------------------------------------
+
+
+class TestExtractBlockedTools:
+    def test_all_allowed_returns_none(self):
+        from litellm.types.utils import ChatCompletionMessageToolCall, Function
+
+        tc = ChatCompletionMessageToolCall(
+            id="call_1", type="function", function=Function(name="fn", arguments="{}")
+        )
+        service_resp = {
+            "choices": [
+                {
+                    "message": {
+                        "tool_calls": [{"id": "call_1"}],
+                        "content": "",
+                    }
+                }
+            ]
+        }
+        result = RubrikLogger._extract_blocked_tools(service_resp, [tc])
+        assert result is None
+
+    def test_some_blocked_returns_result(self):
+        from litellm.types.utils import ChatCompletionMessageToolCall, Function
+
+        tc1 = ChatCompletionMessageToolCall(
+            id="call_1",
+            type="function",
+            function=Function(name="fn1", arguments="{}"),
+        )
+        tc2 = ChatCompletionMessageToolCall(
+            id="call_2",
+            type="function",
+            function=Function(name="fn2", arguments="{}"),
+        )
+        service_resp = {
+            "choices": [
+                {
+                    "message": {
+                        "tool_calls": [{"id": "call_1"}],
+                        "content": "blocked fn2",
+                    }
+                }
+            ]
+        }
+        result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2])
+        assert result is not None
+        assert len(result.allowed_tools) == 1
+        assert result.allowed_tools[0].id == "call_1"
+        assert "blocked fn2" in result.explanation
+
+    def test_empty_choices_raises(self):
+        with pytest.raises(Exception, match="empty response"):
+            RubrikLogger._extract_blocked_tools({"choices": []}, [])
+
+
+# -- Resolve model -------------------------------------------------------------
+
+
+class TestResolveModel:
+    def test_model_from_response(self):
+        from unittest.mock import Mock
+
+        response = Mock()
+        response.model = "gpt-4"
+        result = RubrikLogger._resolve_model({"response": response}, {})
+        assert result == "gpt-4"
+
+    def test_model_from_call_details(self):
+        result = RubrikLogger._resolve_model({}, {"model": "claude-3"})
+        assert result == "claude-3"
+
+    def test_fallback_to_unknown(self):
+        result = RubrikLogger._resolve_model({}, {})
+        assert result == "unknown"
+
+    def test_empty_model_on_response_returns_unknown(self):
+        from unittest.mock import Mock
+
+        response = Mock()
+        response.model = ""
+        result = RubrikLogger._resolve_model({"response": response}, {"model": "fallback"})
+        assert result == "unknown"