From 44ee768307f00317727fed588546f217a358c815 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 17 Feb 2026 17:21:40 -0800 Subject: [PATCH] feat: add claude_code guardrails + provider-agnostic block_hosted_tools Adds three new guardrail hooks for Claude Code proxy deployments: **block_hosted_tools** (provider-agnostic, new top-level hook) - Blocks platform-executed hosted tools from Anthropic, OpenAI, and Gemini - Per-provider YAML configs (anthropic.yaml, openai.yaml, gemini.yaml) define tool type prefixes (Anthropic versioned types), exact type strings (OpenAI), and native top-level keys (Gemini googleSearch, codeExecution, etc.) - Registered as guardrail type: "block_hosted_tools" **claude_code_prompt_cache** - Auto-injects cache_control: {type: ephemeral} into system messages - Only applies to Anthropic API models (detected via get_llm_provider) - Reuses AnthropicCacheControlHook._safe_insert_cache_control_in_message **claude_code_block_expensive_flags** - Blocks speed=fast (~6x pricing), inference_geo, thinking.type=enabled - Config-driven via expensive_api_flags.yaml; inherits Anthropic hosted tool prefixes from block_hosted_tools/anthropic.yaml via inherit_from Policy templates updated: - claude-code-pii-protection: uses block_hosted_tools + litellm_content_filter - claude-code-cost-optimization: uses prompt_cache + block_expensive_flags Co-Authored-By: Claude Sonnet 4.6 --- litellm/policy_templates_backup.json | 119 ++++++++++++ .../block_hosted_tools/__init__.py | 43 +++++ .../block_hosted_tools/anthropic.yaml | 24 +++ .../block_hosted_tools/gemini.yaml | 35 ++++ .../block_hosted_tools/guardrail.py | 177 ++++++++++++++++++ .../block_hosted_tools/openai.yaml | 19 ++ .../guardrail_hooks/claude_code/__init__.py | 77 ++++++++ .../claude_code/block_expensive_flags.py | 175 +++++++++++++++++ .../claude_code/expensive_api_flags.yaml | 35 ++++ .../claude_code/prompt_cache.py | 91 +++++++++ litellm/types/guardrails.py | 3 + policy_templates.json | 119 ++++++++++++ 12 files changed, 917 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/anthropic.yaml create mode 100644 litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/gemini.yaml create mode 100644 litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/guardrail.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/openai.yaml create mode 100644 litellm/proxy/guardrails/guardrail_hooks/claude_code/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/claude_code/block_expensive_flags.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/claude_code/expensive_api_flags.yaml create mode 100644 litellm/proxy/guardrails/guardrail_hooks/claude_code/prompt_cache.py diff --git a/litellm/policy_templates_backup.json b/litellm/policy_templates_backup.json index a0ffd6acd30..ec6c0533fef 100644 --- a/litellm/policy_templates_backup.json +++ b/litellm/policy_templates_backup.json @@ -674,5 +674,124 @@ ], "guardrails_remove": [] } + }, + { + "id": "claude-code-pii-protection", + "title": "Claude Code PII Protection", + "description": "Prevents data leakage in Claude Code sessions by blocking all Anthropic-hosted tools (bash, web_search, web_fetch, code_execution, memory, text_editor) and masking PII and credentials found in request messages.", + "icon": "ShieldCheckIcon", + "iconColor": "text-indigo-500", + "iconBg": "bg-indigo-50", + "guardrails": [ + "claude-code-block-hosted-tools", + "claude-code-pii-credentials", + "claude-code-pii-personal-data" + ], + "complexity": "Medium", + "guardrailDefinitions": [ + { + "guardrail_name": "claude-code-block-hosted-tools", + "litellm_params": { + "guardrail": "block_hosted_tools", + "mode": "pre_call" + }, + "guardrail_info": { + "description": "Blocks platform-hosted tools from Anthropic (bash, web_search, web_fetch, code_execution, memory, text_editor, tool_search), OpenAI (code_interpreter, file_search, web_search_preview), and Gemini (googleSearch, code_execution, url_context). Only custom tool schemas are permitted." + } + }, + { + "guardrail_name": "claude-code-pii-credentials", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "patterns": [ + {"pattern_type": "prebuilt", "pattern_name": "aws_access_key", "action": "BLOCK"}, + {"pattern_type": "prebuilt", "pattern_name": "aws_secret_key", "action": "BLOCK"}, + {"pattern_type": "prebuilt", "pattern_name": "github_token", "action": "BLOCK"}, + {"pattern_type": "prebuilt", "pattern_name": "slack_token", "action": "BLOCK"}, + {"pattern_type": "prebuilt", "pattern_name": "generic_api_key", "action": "BLOCK"} + ], + "pattern_redaction_format": "[{pattern_name}_REDACTED]" + }, + "guardrail_info": { + "description": "Blocks requests containing credentials or API keys (AWS, GitHub, Slack) that could be exfiltrated through the LLM context." + } + }, + { + "guardrail_name": "claude-code-pii-personal-data", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "patterns": [ + {"pattern_type": "prebuilt", "pattern_name": "us_ssn", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "email", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "us_phone", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "credit_card", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "visa", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "mastercard", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "iban", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "ipv4", "action": "MASK"} + ], + "pattern_redaction_format": "[{pattern_name}_REDACTED]" + }, + "guardrail_info": { + "description": "Masks SSNs, email addresses, phone numbers, credit card numbers, IBANs, and IP addresses before they enter the LLM context." + } + } + ], + "templateData": { + "policy_name": "claude-code-pii-protection", + "description": "Prevents data leakage in Claude Code sessions. Blocks Anthropic-hosted tools and masks PII and credentials in request messages.", + "guardrails_add": [ + "claude-code-block-hosted-tools", + "claude-code-pii-credentials", + "claude-code-pii-personal-data" + ], + "guardrails_remove": [] + } + }, + { + "id": "claude-code-cost-optimization", + "title": "Claude Code Cost Optimization", + "description": "Reduces Claude Code API spend by blocking expensive inference modes (fast/turbo, inference_geo, and extended thinking) and automatically injecting prompt caching headers into system messages to maximize cache hit rates.", + "icon": "CurrencyDollarIcon", + "iconColor": "text-green-500", + "iconBg": "bg-green-50", + "guardrails": [ + "claude-code-inject-prompt-cache", + "claude-code-block-expensive-flags" + ], + "complexity": "Low", + "guardrailDefinitions": [ + { + "guardrail_name": "claude-code-inject-prompt-cache", + "litellm_params": { + "guardrail": "claude_code_prompt_cache", + "mode": "pre_call" + }, + "guardrail_info": { + "description": "Automatically adds cache_control: {type: ephemeral} to system messages so Anthropic can cache the system prompt prefix. Only applies to Anthropic API models. Reduces cost on repeated calls that share the same system prompt." + } + }, + { + "guardrail_name": "claude-code-block-expensive-flags", + "litellm_params": { + "guardrail": "claude_code_block_expensive_flags", + "mode": "pre_call" + }, + "guardrail_info": { + "description": "Blocks expensive Anthropic API flags: fast/turbo inference (speed=fast, ~6x pricing), inference geo-routing (inference_geo, 1.1x pricing), and extended thinking (thinking.type=enabled). Also blocks Anthropic-hosted tools inherited from the hosted tools policy." + } + } + ], + "templateData": { + "policy_name": "claude-code-cost-optimization", + "description": "Reduces Claude Code API spend by blocking fast mode, inference_geo, extended thinking, and Anthropic-hosted tools, while auto-injecting prompt caching into system messages.", + "guardrails_add": [ + "claude-code-inject-prompt-cache", + "claude-code-block-expensive-flags" + ], + "guardrails_remove": [] + } } ] diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/__init__.py new file mode 100644 index 00000000000..b7651048da3 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/__init__.py @@ -0,0 +1,43 @@ +""" +Block Hosted Tools guardrail for LiteLLM. + +Blocks platform-executed server-side tools from Anthropic, OpenAI, and Gemini. +Provider tool lists are maintained in per-provider YAML files (anthropic.yaml, +openai.yaml, gemini.yaml) alongside this module. +""" + +from typing import TYPE_CHECKING + +import litellm +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .guardrail import BlockHostedToolsGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> BlockHostedToolsGuardrail: + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("BlockHostedToolsGuardrail requires a guardrail_name") + instance = BlockHostedToolsGuardrail( + guardrail_name=guardrail_name, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(instance) + return instance + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.BLOCK_HOSTED_TOOLS.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.BLOCK_HOSTED_TOOLS.value: BlockHostedToolsGuardrail, +} + +__all__ = ["BlockHostedToolsGuardrail", "initialize_guardrail"] diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/anthropic.yaml b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/anthropic.yaml new file mode 100644 index 00000000000..423e9f33453 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/anthropic.yaml @@ -0,0 +1,24 @@ +# Anthropic server-side hosted tools. +# Tools are matched by prefix against the tool 'type' field, because Anthropic +# appends a version date suffix (e.g. bash_20250124, web_search_20250305). +# +# Reference: https://docs.anthropic.com/en/docs/agents-and-tools/computer-use + +provider: "anthropic" +description: "Anthropic server-side hosted tools executed on Anthropic infrastructure" + +tool_type_prefixes: + - prefix: "bash_" + description: "Bash shell execution tool" + - prefix: "web_search_" + description: "Web search tool" + - prefix: "web_fetch_" + description: "Web fetch / URL retrieval tool" + - prefix: "code_execution_" + description: "Code execution sandbox tool" + - prefix: "memory_" + description: "Memory read/write tool" + - prefix: "text_editor_" + description: "Text editor tool" + - prefix: "tool_search_" + description: "Tool search / discovery tool" diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/gemini.yaml b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/gemini.yaml new file mode 100644 index 00000000000..dee17f9489c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/gemini.yaml @@ -0,0 +1,35 @@ +# Google Gemini / Vertex AI platform-hosted tools executed on Google infrastructure. +# Gemini tools appear in two formats in the API: +# - OpenAI-compat: tool dict has 'type' set to the value below +# - Gemini native: tool dict has a top-level key matching the tool name +# Both formats are checked. +# +# Reference: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/function-calling + +provider: "gemini" +description: "Google Gemini / Vertex AI platform-hosted tools" + +# Matched against the 'type' field (OpenAI-compat format) +tool_type_exact: + - type: "code_execution" + description: "Code execution sandbox (OpenAI-compat format)" + - type: "url_context" + description: "URL fetching and context grounding (OpenAI-compat format)" + +# Matched as top-level dict keys (Gemini native format) +# e.g. {"googleSearch": {}}, {"codeExecution": {}} +tool_top_level_keys: + - key: "googleSearch" + description: "Google Search grounding" + - key: "googleSearchRetrieval" + description: "Google Search retrieval with data store integration" + - key: "enterpriseWebSearch" + description: "Enterprise web search" + - key: "url_context" + description: "URL context grounding" + - key: "code_execution" + description: "Code execution sandbox" + - key: "googleMaps" + description: "Google Maps integration" + - key: "computerUse" + description: "Computer use / automated screen interaction" diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/guardrail.py new file mode 100644 index 00000000000..b9557051f00 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/guardrail.py @@ -0,0 +1,177 @@ +""" +Block Hosted Tools Guardrail + +Provider-agnostic guardrail that blocks platform-executed hosted tools from +Anthropic, OpenAI, and Google Gemini. Tool definitions are loaded from +per-provider YAML files in this directory so the list can be extended without +touching Python code. + +Detection strategy per provider: + Anthropic — prefix match on tool 'type' field (versioned: bash_20250124, etc.) + OpenAI — exact match on tool 'type' field (code_interpreter, file_search, …) + Gemini — exact match on tool 'type' field OR presence of a top-level key + (googleSearch, codeExecution, …) in the tool dict +""" + +import os +from typing import TYPE_CHECKING, Dict, FrozenSet, List, Literal, Optional, Tuple + +import yaml +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_DIR = os.path.dirname(__file__) + + +# --------------------------------------------------------------------------- +# Config loading — runs once at import time +# --------------------------------------------------------------------------- + +def _load_provider_configs() -> Dict[str, dict]: + """Load all *.yaml files in this directory, keyed by provider name.""" + configs: Dict[str, dict] = {} + for fname in os.listdir(_DIR): + if not fname.endswith(".yaml"): + continue + with open(os.path.join(_DIR, fname)) as f: + cfg = yaml.safe_load(f) + provider = cfg.get("provider") + if provider: + configs[provider] = cfg + return configs + + +def _build_match_sets( + configs: Dict[str, dict], +) -> Tuple[tuple, FrozenSet[str], FrozenSet[str]]: + """ + Derive three matching structures from all provider configs: + + Returns: + type_prefixes — tuple of prefix strings (Anthropic-style versioned types) + exact_types — frozenset of exact type strings (OpenAI / Gemini-compat) + top_level_keys — frozenset of Gemini native top-level dict keys + """ + prefixes: List[str] = [] + exact: List[str] = [] + keys: List[str] = [] + + for cfg in configs.values(): + for entry in cfg.get("tool_type_prefixes", []): + if isinstance(entry, dict) and "prefix" in entry: + prefixes.append(entry["prefix"]) + for entry in cfg.get("tool_type_exact", []): + if isinstance(entry, dict) and "type" in entry: + exact.append(entry["type"]) + for entry in cfg.get("tool_top_level_keys", []): + if isinstance(entry, dict) and "key" in entry: + keys.append(entry["key"]) + + return tuple(prefixes), frozenset(exact), frozenset(keys) + + +_PROVIDER_CONFIGS = _load_provider_configs() +_TYPE_PREFIXES, _EXACT_TYPES, _TOP_LEVEL_KEYS = _build_match_sets(_PROVIDER_CONFIGS) + + +# --------------------------------------------------------------------------- +# Tool-matching helpers +# --------------------------------------------------------------------------- + +def _match_tool(tool: dict) -> Optional[str]: + """ + Return a human-readable description if the tool is a known hosted tool, + or None if it is a user-defined tool (pass through). + """ + if not isinstance(tool, dict): + return None + + raw_type: Optional[str] = tool.get("type") + + # Skip generic OpenAI function wrapper — never a hosted tool + if raw_type == "function": + return None + + if raw_type: + # Prefix match (Anthropic versioned types) + if raw_type.startswith(_TYPE_PREFIXES): + name = tool.get("name") or raw_type + return f"{name} ({raw_type})" + # Exact match (OpenAI / Gemini-compat) + if raw_type in _EXACT_TYPES: + name = tool.get("name") or raw_type + return f"{name} ({raw_type})" + + # Gemini native top-level keys ({"googleSearch": {}, ...}) + for key in _TOP_LEVEL_KEYS: + if key in tool: + return f"{key}" + + return None + + +# --------------------------------------------------------------------------- +# Guardrail class +# --------------------------------------------------------------------------- + +class BlockHostedToolsGuardrail(CustomGuardrail): + """ + Guardrail that blocks platform-hosted tools from Anthropic, OpenAI, and Gemini. + + Raises HTTP 403 if any tool in the request matches a known hosted tool + from any supported provider. Provider tool lists are defined in YAML + files alongside this module and loaded at import time. + """ + + def __init__(self, **kwargs): + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [GuardrailEventHooks.pre_call] + super().__init__(**kwargs) + verbose_proxy_logger.debug( + f"BlockHostedToolsGuardrail initialized " + f"(providers: {list(_PROVIDER_CONFIGS.keys())})" + ) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + tools: List[dict] = list(inputs.get("tools") or []) # type: ignore[assignment] + blocked: List[str] = [] + + for tool in tools: + desc = _match_tool(tool) + if desc: + blocked.append(desc) + + if blocked: + raise HTTPException( + status_code=403, + detail={ + "error": ( + "the following platform-hosted tools are disabled by " + "your organization's policy: " + ", ".join(blocked) + ), + "guardrail": self.guardrail_name, + "blocked_tools": blocked, + }, + ) + + return inputs diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/openai.yaml b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/openai.yaml new file mode 100644 index 00000000000..8b4a50ebaca --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_hosted_tools/openai.yaml @@ -0,0 +1,19 @@ +# OpenAI platform-hosted tools executed on OpenAI's infrastructure. +# These are matched by exact 'type' field value (no version suffix). +# +# Reference: https://platform.openai.com/docs/guides/tools + +provider: "openai" +description: "OpenAI platform-hosted tools executed on OpenAI infrastructure" + +tool_type_exact: + - type: "code_interpreter" + description: "Python code execution and file output sandbox" + - type: "file_search" + description: "Semantic file search across vector stores" + - type: "web_search_preview" + description: "Bing-powered live web search" + - type: "computer_use_preview" + description: "Computer use / automated screen interaction" + - type: "shell" + description: "Shell command execution in containers" diff --git a/litellm/proxy/guardrails/guardrail_hooks/claude_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/claude_code/__init__.py new file mode 100644 index 00000000000..d10dda9b5f3 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/claude_code/__init__.py @@ -0,0 +1,77 @@ +""" +Claude Code guardrail integrations for LiteLLM. + +Two focused policy-enforcement guardrails for Claude Code deployments: + +1. claude_code_prompt_cache — auto-inject Anthropic prompt-caching headers +2. claude_code_block_expensive_flags — block expensive API flags (fast mode, etc.) + +Hosted tool blocking is handled by the provider-agnostic 'block_hosted_tools' +guardrail (see guardrail_hooks/block_hosted_tools/). +""" + +from typing import TYPE_CHECKING + +import litellm +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .block_expensive_flags import ClaudeCodeBlockExpensiveFlagsGuardrail +from .prompt_cache import ClaudeCodePromptCacheGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +# ------------------------------------------------------------------ # +# Per-guardrail initializer functions # +# ------------------------------------------------------------------ # + + +def _init_prompt_cache( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> ClaudeCodePromptCacheGuardrail: + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("ClaudeCodePromptCacheGuardrail requires a guardrail_name") + instance = ClaudeCodePromptCacheGuardrail( + guardrail_name=guardrail_name, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(instance) + return instance + + +def _init_block_expensive_flags( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> ClaudeCodeBlockExpensiveFlagsGuardrail: + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("ClaudeCodeBlockExpensiveFlagsGuardrail requires a guardrail_name") + instance = ClaudeCodeBlockExpensiveFlagsGuardrail( + guardrail_name=guardrail_name, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(instance) + return instance + + +# ------------------------------------------------------------------ # +# Registries consumed by the guardrail loader # +# ------------------------------------------------------------------ # + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.CLAUDE_CODE_PROMPT_CACHE.value: _init_prompt_cache, + SupportedGuardrailIntegrations.CLAUDE_CODE_BLOCK_EXPENSIVE_FLAGS.value: _init_block_expensive_flags, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.CLAUDE_CODE_PROMPT_CACHE.value: ClaudeCodePromptCacheGuardrail, + SupportedGuardrailIntegrations.CLAUDE_CODE_BLOCK_EXPENSIVE_FLAGS.value: ClaudeCodeBlockExpensiveFlagsGuardrail, +} + +__all__ = [ + "ClaudeCodePromptCacheGuardrail", + "ClaudeCodeBlockExpensiveFlagsGuardrail", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/claude_code/block_expensive_flags.py b/litellm/proxy/guardrails/guardrail_hooks/claude_code/block_expensive_flags.py new file mode 100644 index 00000000000..9d0c1fdf1be --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/claude_code/block_expensive_flags.py @@ -0,0 +1,175 @@ +""" +Claude Code - Block Expensive API Flags Guardrail + +Blocks Anthropic API parameters that trigger feature-specific pricing surcharges +(fast mode, inference_geo, extended thinking). Also inherits the hosted tool +type prefixes from hosted_tool_types.yaml so hosted tools are blocked here too. + +Blocked params are driven by expensive_api_flags.yaml which references +hosted_tool_types.yaml via `inherit_from`, following the same pattern as +harmful_child_safety.yaml inherits from harm_toxic_abuse.json. +""" + +import os +from typing import TYPE_CHECKING, Any, List, Literal, Optional + +import yaml +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_DIR = os.path.dirname(__file__) +_FLAGS_YAML = os.path.join(_DIR, "expensive_api_flags.yaml") +_TOOLS_YAML = os.path.join(_DIR, "hosted_tool_types.yaml") + + +def _load_config() -> dict: + """Load expensive_api_flags.yaml, merging any inherited config.""" + with open(_FLAGS_YAML) as f: + config: dict = yaml.safe_load(f) + + inherited_prefixes: List[str] = [] + inherit_from = config.get("inherit_from") + if inherit_from: + inherit_path = os.path.join(_DIR, inherit_from) + with open(inherit_path) as f: + inherited: dict = yaml.safe_load(f) + inherited_prefixes = [ + entry["prefix"] + for entry in inherited.get("tool_type_prefixes", []) + if isinstance(entry, dict) and "prefix" in entry + ] + + config["_inherited_tool_type_prefixes"] = inherited_prefixes + return config + + +_CONFIG: dict = _load_config() +_BLOCKED_PARAMS: List[dict] = _CONFIG.get("blocked_params", []) +_INHERITED_TOOL_TYPE_PREFIXES: tuple = tuple(_CONFIG.get("_inherited_tool_type_prefixes", [])) + + +def _tool_type(tool: dict) -> Optional[str]: + t = tool.get("type") + if t and t != "function": + return t + return None + + +def _is_hosted_tool(tool: dict) -> bool: + if not _INHERITED_TOOL_TYPE_PREFIXES: + return False + t = _tool_type(tool) + if not t: + return False + return t.startswith(_INHERITED_TOOL_TYPE_PREFIXES) + + +def _check_param( + request_data: dict, param_cfg: dict +) -> Optional[str]: + """ + Return an error message if the param in request_data matches a blocked value. + Returns None when the param is not blocked. + """ + param = param_cfg.get("param") + if not param: + return None + + value: Any = request_data.get(param) + if value is None: + return None + + nested_key = param_cfg.get("nested_key") + if nested_key: + # e.g. thinking.type — value must be a dict + if not isinstance(value, dict): + return None + value = value.get(nested_key) + if value is None: + return None + + blocked_values: List[str] = param_cfg.get("blocked_values", []) + if "*" in blocked_values or str(value) in blocked_values: + return param_cfg.get("error_message", f"{param} is disabled by your organization's policy") + + return None + + +class ClaudeCodeBlockExpensiveFlagsGuardrail(CustomGuardrail): + """ + Guardrail that blocks expensive Anthropic API flags. + + Checks request_data for feature-specific pricing flags (fast mode, + inference_geo, extended thinking) and Anthropic-hosted tools inherited + from hosted_tool_types.yaml. Raises HTTP 403 on the first violation. + """ + + def __init__(self, **kwargs): + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [GuardrailEventHooks.pre_call] + super().__init__(**kwargs) + verbose_proxy_logger.debug("ClaudeCodeBlockExpensiveFlagsGuardrail initialized") + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + # ------------------------------------------------------------------ # + # 1. Check request-level blocked params (speed, inference_geo, etc.) # + # ------------------------------------------------------------------ # + for param_cfg in _BLOCKED_PARAMS: + error_msg = _check_param(request_data, param_cfg) + if error_msg: + param = param_cfg.get("param", "unknown") + raise HTTPException( + status_code=403, + detail={ + "error": error_msg, + "guardrail": self.guardrail_name, + "blocked_param": param, + }, + ) + + # ------------------------------------------------------------------ # + # 2. Check for inherited hosted tools (from hosted_tool_types.yaml) # + # ------------------------------------------------------------------ # + if _INHERITED_TOOL_TYPE_PREFIXES: + tools: List[dict] = list(inputs.get("tools") or []) # type: ignore[assignment] + blocked_tools: List[str] = [] + for tool in tools: + if _is_hosted_tool(tool): + tool_type = _tool_type(tool) or "unknown" + tool_name = tool.get("name") or tool_type + blocked_tools.append(f"{tool_name} ({tool_type})") + + if blocked_tools: + raise HTTPException( + status_code=403, + detail={ + "error": ( + "the following Anthropic-hosted tools are disabled by " + "your organization's policy: " + ", ".join(blocked_tools) + ), + "guardrail": self.guardrail_name, + "blocked_tools": blocked_tools, + }, + ) + + return inputs diff --git a/litellm/proxy/guardrails/guardrail_hooks/claude_code/expensive_api_flags.yaml b/litellm/proxy/guardrails/guardrail_hooks/claude_code/expensive_api_flags.yaml new file mode 100644 index 00000000000..c0708e091ff --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/claude_code/expensive_api_flags.yaml @@ -0,0 +1,35 @@ +# Expensive Anthropic API flags. +# Blocks request parameters that trigger feature-specific pricing surcharges. +# Also inherits hosted tool type prefixes so hosted tools are blocked as well. +# +# Pricing reference: https://platform.claude.com/docs/en/about-claude/pricing#feature-specific-pricing + +category_name: "claude_code_expensive_flags" +description: "Blocks expensive Anthropic API features that incur pricing surcharges" +default_action: "BLOCK" + +# Inherit hosted tool type prefixes — hosted tools are also expensive. +inherit_from: "../block_hosted_tools/anthropic.yaml" + +# Request-level parameters to block. +# Each entry maps a top-level request_data key to one or more blocked values. +# Use "*" as the value to block the flag entirely regardless of its value. +blocked_params: + - param: "speed" + blocked_values: + - "fast" + description: "Fast/turbo inference mode (~6x pricing surcharge)" + error_message: "fast inference mode (speed=fast) is disabled by your organization's policy" + + - param: "inference_geo" + blocked_values: + - "*" + description: "Inference geo-routing (1.1x pricing surcharge)" + error_message: "inference geo-routing (inference_geo) is disabled by your organization's policy" + + - param: "thinking" + blocked_values: + - "enabled" + nested_key: "type" + description: "Extended thinking / budget-token thinking mode" + error_message: "extended thinking mode (thinking.type=enabled) is disabled by your organization's policy" diff --git a/litellm/proxy/guardrails/guardrail_hooks/claude_code/prompt_cache.py b/litellm/proxy/guardrails/guardrail_hooks/claude_code/prompt_cache.py new file mode 100644 index 00000000000..53e058f873d --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/claude_code/prompt_cache.py @@ -0,0 +1,91 @@ +""" +Claude Code - Prompt Cache Injection Guardrail + +Automatically injects cache_control: {type: ephemeral} into system messages +so Anthropic can cache the prompt prefix, reducing costs on repeated calls. + +Only runs when the request targets an Anthropic API model. Uses the existing +AnthropicCacheControlHook._safe_insert_cache_control_in_message utility so +the injection logic is not duplicated. +""" + +from typing import TYPE_CHECKING, Literal, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import ChatCompletionCachedContent +from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.utils import get_llm_provider + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +def _is_anthropic_model(model: Optional[str]) -> bool: + """Return True when the model resolves to the Anthropic provider.""" + if not model: + return False + try: + _, custom_llm_provider, _, _ = get_llm_provider(model=model) + return custom_llm_provider == "anthropic" + except Exception: + return False + + +class ClaudeCodePromptCacheGuardrail(CustomGuardrail): + """ + Guardrail that injects Anthropic prompt-caching headers into system messages. + + Targets only Anthropic API models — the provider is detected from the + model string so this guardrail is safe to apply globally without + restricting other providers. + """ + + def __init__(self, **kwargs): + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [GuardrailEventHooks.pre_call] + super().__init__(**kwargs) + verbose_proxy_logger.debug("ClaudeCodePromptCacheGuardrail initialized") + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + model: Optional[str] = request_data.get("model") or inputs.get("model") # type: ignore[assignment] + if not _is_anthropic_model(model): + verbose_proxy_logger.debug( + f"ClaudeCodePromptCacheGuardrail: skipping non-Anthropic model '{model}'" + ) + return inputs + + messages = request_data.get("messages") or [] + if not messages: + return inputs + + control = ChatCompletionCachedContent(type="ephemeral") + modified = [] + for msg in messages: + if msg.get("role") == "system": + msg = AnthropicCacheControlHook._safe_insert_cache_control_in_message( + message=msg, # type: ignore[arg-type] + control=control, + ) + modified.append(msg) + + request_data["messages"] = modified + verbose_proxy_logger.debug( + "ClaudeCodePromptCacheGuardrail: injected cache_control into system messages" + ) + return inputs diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 74ccb34ca6e..f893f2246e0 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -69,6 +69,9 @@ class SupportedGuardrailIntegrations(Enum): GENERIC_GUARDRAIL_API = "generic_guardrail_api" QUALIFIRE = "qualifire" CUSTOM_CODE = "custom_code" + BLOCK_HOSTED_TOOLS = "block_hosted_tools" + CLAUDE_CODE_PROMPT_CACHE = "claude_code_prompt_cache" + CLAUDE_CODE_BLOCK_EXPENSIVE_FLAGS = "claude_code_block_expensive_flags" class Role(Enum): diff --git a/policy_templates.json b/policy_templates.json index a0ffd6acd30..ec6c0533fef 100644 --- a/policy_templates.json +++ b/policy_templates.json @@ -674,5 +674,124 @@ ], "guardrails_remove": [] } + }, + { + "id": "claude-code-pii-protection", + "title": "Claude Code PII Protection", + "description": "Prevents data leakage in Claude Code sessions by blocking all Anthropic-hosted tools (bash, web_search, web_fetch, code_execution, memory, text_editor) and masking PII and credentials found in request messages.", + "icon": "ShieldCheckIcon", + "iconColor": "text-indigo-500", + "iconBg": "bg-indigo-50", + "guardrails": [ + "claude-code-block-hosted-tools", + "claude-code-pii-credentials", + "claude-code-pii-personal-data" + ], + "complexity": "Medium", + "guardrailDefinitions": [ + { + "guardrail_name": "claude-code-block-hosted-tools", + "litellm_params": { + "guardrail": "block_hosted_tools", + "mode": "pre_call" + }, + "guardrail_info": { + "description": "Blocks platform-hosted tools from Anthropic (bash, web_search, web_fetch, code_execution, memory, text_editor, tool_search), OpenAI (code_interpreter, file_search, web_search_preview), and Gemini (googleSearch, code_execution, url_context). Only custom tool schemas are permitted." + } + }, + { + "guardrail_name": "claude-code-pii-credentials", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "patterns": [ + {"pattern_type": "prebuilt", "pattern_name": "aws_access_key", "action": "BLOCK"}, + {"pattern_type": "prebuilt", "pattern_name": "aws_secret_key", "action": "BLOCK"}, + {"pattern_type": "prebuilt", "pattern_name": "github_token", "action": "BLOCK"}, + {"pattern_type": "prebuilt", "pattern_name": "slack_token", "action": "BLOCK"}, + {"pattern_type": "prebuilt", "pattern_name": "generic_api_key", "action": "BLOCK"} + ], + "pattern_redaction_format": "[{pattern_name}_REDACTED]" + }, + "guardrail_info": { + "description": "Blocks requests containing credentials or API keys (AWS, GitHub, Slack) that could be exfiltrated through the LLM context." + } + }, + { + "guardrail_name": "claude-code-pii-personal-data", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "patterns": [ + {"pattern_type": "prebuilt", "pattern_name": "us_ssn", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "email", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "us_phone", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "credit_card", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "visa", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "mastercard", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "iban", "action": "MASK"}, + {"pattern_type": "prebuilt", "pattern_name": "ipv4", "action": "MASK"} + ], + "pattern_redaction_format": "[{pattern_name}_REDACTED]" + }, + "guardrail_info": { + "description": "Masks SSNs, email addresses, phone numbers, credit card numbers, IBANs, and IP addresses before they enter the LLM context." + } + } + ], + "templateData": { + "policy_name": "claude-code-pii-protection", + "description": "Prevents data leakage in Claude Code sessions. Blocks Anthropic-hosted tools and masks PII and credentials in request messages.", + "guardrails_add": [ + "claude-code-block-hosted-tools", + "claude-code-pii-credentials", + "claude-code-pii-personal-data" + ], + "guardrails_remove": [] + } + }, + { + "id": "claude-code-cost-optimization", + "title": "Claude Code Cost Optimization", + "description": "Reduces Claude Code API spend by blocking expensive inference modes (fast/turbo, inference_geo, and extended thinking) and automatically injecting prompt caching headers into system messages to maximize cache hit rates.", + "icon": "CurrencyDollarIcon", + "iconColor": "text-green-500", + "iconBg": "bg-green-50", + "guardrails": [ + "claude-code-inject-prompt-cache", + "claude-code-block-expensive-flags" + ], + "complexity": "Low", + "guardrailDefinitions": [ + { + "guardrail_name": "claude-code-inject-prompt-cache", + "litellm_params": { + "guardrail": "claude_code_prompt_cache", + "mode": "pre_call" + }, + "guardrail_info": { + "description": "Automatically adds cache_control: {type: ephemeral} to system messages so Anthropic can cache the system prompt prefix. Only applies to Anthropic API models. Reduces cost on repeated calls that share the same system prompt." + } + }, + { + "guardrail_name": "claude-code-block-expensive-flags", + "litellm_params": { + "guardrail": "claude_code_block_expensive_flags", + "mode": "pre_call" + }, + "guardrail_info": { + "description": "Blocks expensive Anthropic API flags: fast/turbo inference (speed=fast, ~6x pricing), inference geo-routing (inference_geo, 1.1x pricing), and extended thinking (thinking.type=enabled). Also blocks Anthropic-hosted tools inherited from the hosted tools policy." + } + } + ], + "templateData": { + "policy_name": "claude-code-cost-optimization", + "description": "Reduces Claude Code API spend by blocking fast mode, inference_geo, extended thinking, and Anthropic-hosted tools, while auto-injecting prompt caching into system messages.", + "guardrails_add": [ + "claude-code-inject-prompt-cache", + "claude-code-block-expensive-flags" + ], + "guardrails_remove": [] + } } ]