From 5862f45a6cbeab0b50ce9350f9e43930b8d6a29d Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Wed, 20 May 2026 19:15:38 +0000
Subject: [PATCH] fix(rubrik): collapse BlockedToolsResult dead-code into
 Optional[str]
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The `allowed_tools` field on `BlockedToolsResult` was computed in
`_extract_blocked_tools` but never read by the only caller — when any
tool was blocked the integration unconditionally raised
`ModifyResponseException` to reject the full response, never doing
partial filtering. Drop the dataclass and return the blocking
explanation directly as `Optional[str]` so there's no misleading shape
hinting at unused partial-filter capability.

Co-authored-by: Greptile <greptile-apps[bot]@users.noreply.github.com>
---
 litellm/integrations/rubrik.py                | 34 +++++++------------
 .../test_litellm/integrations/test_rubrik.py  |  9 ++---
 2 files changed, 15 insertions(+), 28 deletions(-)

diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py
index c9e22661267..da6c38b8c23 100644
--- a/litellm/integrations/rubrik.py
+++ b/litellm/integrations/rubrik.py
@@ -6,7 +6,6 @@ import random
 import time
 import urllib.parse
 import uuid
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal, Optional
 
 import httpx
@@ -39,14 +38,6 @@ _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):
     preserve_events_added_during_flush = True
 
@@ -210,12 +201,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
         service_response = await self._post_to_tool_blocking_service(
             response_data, req_data
         )
-        blocked = self._extract_blocked_tools(service_response, message_tool_calls)
+        blocked_explanation = self._extract_blocked_tools(
+            service_response, message_tool_calls
+        )
 
-        if blocked:
+        if blocked_explanation is not None:
             model = self._resolve_model(request_data, call_details)
             raise ModifyResponseException(
-                message=blocked.explanation,
+                message=blocked_explanation,
                 model=model,
                 request_data=request_data,
                 guardrail_name=self.guardrail_name,
@@ -504,12 +497,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
     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.
+    ) -> Optional[str]:
+        """Return the blocking explanation if any tool calls were blocked.
 
         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.
+        the full set of tool calls. Returns ``None`` if all tools are allowed, or
+        the explanation string (prefixed with newlines) otherwise.
 
         Expects service_response in OpenAI chat completion format:
             {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
@@ -523,13 +516,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
         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]
+        allowed_count = sum(1 for tc in all_tool_calls if tc.id in allowed_ids)
 
-        if len(allowed_tools) == len(all_tool_calls):
+        if allowed_count == len(all_tool_calls):
             return None
 
         explanation = blocking_explanation or "Tool call blocked by policy."
-        return BlockedToolsResult(
-            allowed_tools=allowed_tools,
-            explanation=f"\n\n{explanation}",
-        )
+        return f"\n\n{explanation}"
diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py
index 8800fb945e9..b681a8ed1ff 100644
--- a/tests/test_litellm/integrations/test_rubrik.py
+++ b/tests/test_litellm/integrations/test_rubrik.py
@@ -830,7 +830,7 @@ class TestExtractBlockedTools:
         result = RubrikLogger._extract_blocked_tools(service_resp, [tc])
         assert result is None
 
-    def test_some_blocked_returns_result(self):
+    def test_some_blocked_returns_explanation(self):
         from litellm.types.utils import ChatCompletionMessageToolCall, Function
 
         tc1 = ChatCompletionMessageToolCall(
@@ -855,9 +855,7 @@ class TestExtractBlockedTools:
         }
         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
+        assert "blocked fn2" in result
 
     def test_empty_choices_raises(self):
         with pytest.raises(Exception, match="empty response"):
@@ -881,8 +879,7 @@ class TestExtractBlockedTools:
         }
         result = RubrikLogger._extract_blocked_tools(service_resp, [tc])
         assert result is not None
-        assert result.allowed_tools == []
-        assert "blocked everything" in result.explanation
+        assert "blocked everything" in result
 
 
 # -- Sanitize proxy server request -------------------------------------------