fix(proxy): dispatch prompt injection llm api check

litellm/proxy/utils.py

@@ -376,6 +376,7 @@ class _CallbackCapabilities:
     has_iterator_override: bool = False
     has_streaming_chunk_override: bool = False
     has_guardrail: bool = False
+    has_during_call_hook: bool = False
     has_pre_call_override: bool = False
     # Tuple[(resolved_callback, "override" | "apply_guardrail"), ...]
     # Ordered the same as ``litellm.callbacks``; used to build the streaming
@@ -387,6 +388,14 @@ class _CallbackCapabilities:
     resolved_callbacks: Tuple[Any, ...] = field(default_factory=tuple)
 
 
+def _has_native_during_call_hook(callback: CustomLogger) -> bool:
+    base_hook = CustomLogger.async_moderation_hook
+    callback_hook = getattr(type(callback), "async_moderation_hook", base_hook)
+    return getattr(callback_hook, "__func__", callback_hook) is not getattr(
+        base_hook, "__func__", base_hook
+    )
+
+
 class ProxyLogging:
     """
     Logging/Custom Handlers for proxy.
@@ -1685,6 +1694,7 @@ class ProxyLogging:
         has_iterator_override = False
         has_streaming_chunk_override = False
         has_guardrail = False
+        has_during_call_hook = False
         has_pre_call_override = False
         iterator_overrides: List[Tuple[Any, str]] = []  # (callback, kind)
         resolved_callbacks: List[Any] = []
@@ -1706,6 +1716,9 @@ class ProxyLogging:
                 continue
             if isinstance(resolved, CustomGuardrail):
                 has_guardrail = True
+                has_during_call_hook = True
+            elif _has_native_during_call_hook(resolved):
+                has_during_call_hook = True
             # Use the same leaf-class ``__dict__`` check as the other hook
             # capabilities: only callbacks that actually override the hook
             # contribute to the flag. Setting this for every ``CustomLogger``
@@ -1746,6 +1759,7 @@ class ProxyLogging:
             or any(kind == "apply_guardrail" for _, kind in iterator_overrides),
             has_streaming_chunk_override=has_streaming_chunk_override,
             has_guardrail=has_guardrail,
+            has_during_call_hook=has_during_call_hook,
             has_pre_call_override=has_pre_call_override,
             iterator_overrides=tuple(iterator_overrides),
             resolved_callbacks=tuple(resolved_callbacks),
@@ -1794,7 +1808,7 @@ class ProxyLogging:
 
     @staticmethod
     def has_during_call_guardrails() -> bool:
-        return ProxyLogging._callback_capabilities().has_guardrail
+        return ProxyLogging._callback_capabilities().has_during_call_hook
 
     async def during_call_hook(
         self,
@@ -1803,18 +1817,19 @@ class ProxyLogging:
         call_type: CallTypesLiteral,
     ):
         """
-        Runs the CustomGuardrail's async_moderation_hook() in parallel
+        Runs during-call async_moderation_hook() callbacks in parallel
         """
-        # Fast path: skip the entire guardrail scan when no CustomGuardrail
-        # callbacks are registered. Saves per-request iteration over
+        # Fast path: skip the entire guardrail scan when no during-call
+        # moderation callbacks are registered. Saves per-request iteration over
         # ``litellm.callbacks`` plus an ``asyncio.gather([])`` round trip on
-        # deployments with no guardrails configured.
-        if not ProxyLogging._callback_capabilities().has_guardrail:
+        # deployments with no during-call checks configured.
+        caps = ProxyLogging._callback_capabilities()
+        if not caps.has_during_call_hook:
             return data
         # Step 1: Collect all guardrail tasks to run in parallel
         guardrail_tasks = []
 
-        for callback in litellm.callbacks:
+        for callback in caps.resolved_callbacks:
             if isinstance(callback, CustomGuardrail):
                 ################################################################
                 # Check if guardrail should be run for GuardrailEventHooks.during_call hook
@@ -1871,6 +1886,18 @@ class ProxyLogging:
                         ),
                     )
                 guardrail_tasks.append(guardrail_task)
+            elif _has_native_during_call_hook(callback):
+                if call_type == CallTypes.call_mcp_tool.value:
+                    continue
+                guardrail_task = self._run_guardrail_task_with_enrichment(
+                    callback,
+                    callback.async_moderation_hook(
+                        data=data,
+                        user_api_key_dict=user_api_key_dict,  # type: ignore
+                        call_type=call_type,  # type: ignore
+                    ),
+                )
+                guardrail_tasks.append(guardrail_task)
 
         # Step 2: Run all guardrail tasks in parallel
         if guardrail_tasks:

tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py

@@ -2,13 +2,17 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
 import litellm
+from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.proxy._types import LiteLLMPromptInjectionParams
+from litellm.proxy.hooks.prompt_injection_detection import (
+    _OPTIONAL_PromptInjectionDetection,
+)
 from litellm.proxy.utils import ProxyLogging
 from litellm.types.guardrails import GuardrailEventHooks
 
@@ -31,8 +35,27 @@ def _make_guardrail(name="g1", should_run=True, response=None):
     return cb
 
 
+class _NativeDuringCallLogger(CustomLogger):
+    async def async_moderation_hook(self, data, user_api_key_dict, call_type):
+        return data
+
+
+def test_callback_capabilities_detect_native_during_call_hook(monkeypatch):
+    callback = _NativeDuringCallLogger()
+    monkeypatch.setattr(litellm, "callbacks", [callback])
+
+    caps = ProxyLogging._callback_capabilities()
+
+    assert caps.has_during_call_hook is True
+    assert caps.has_guardrail is False
+    assert caps.resolved_callbacks == (callback,)
+    assert ProxyLogging.has_during_call_guardrails() is True
+
+
 @pytest.mark.asyncio
-async def test_during_call_hook_no_guardrail_fast_path_returns_data(proxy_logging, make_user_api_key_auth, mock_callbacks_disabled):
+async def test_during_call_hook_no_guardrail_fast_path_returns_data(
+    proxy_logging, make_user_api_key_auth, mock_callbacks_disabled
+):
     data = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1}
     out = await proxy_logging.during_call_hook(
         data=data,
@@ -43,7 +66,9 @@ async def test_during_call_hook_no_guardrail_fast_path_returns_data(proxy_loggin
 
 
 @pytest.mark.asyncio
-async def test_during_call_hook_runs_guardrails_in_parallel(proxy_logging, make_user_api_key_auth, monkeypatch):
+async def test_during_call_hook_runs_guardrails_in_parallel(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
     g1 = _make_guardrail("a")
     g2 = _make_guardrail("b")
     monkeypatch.setattr(litellm, "callbacks", [g1, g2])
@@ -62,7 +87,72 @@ async def test_during_call_hook_runs_guardrails_in_parallel(proxy_logging, make_
 
 
 @pytest.mark.asyncio
-async def test_during_call_hook_guardrail_skipped_when_should_not_run(proxy_logging, make_user_api_key_auth, monkeypatch):
+async def test_during_call_hook_runs_prompt_injection_llm_api_check(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    prompt_injection_detection = _OPTIONAL_PromptInjectionDetection(
+        prompt_injection_params=LiteLLMPromptInjectionParams(
+            llm_api_check=True,
+            llm_api_name="prompt-check",
+            llm_api_system_prompt="Return SAFE or UNSAFE",
+            llm_api_fail_call_string="UNSAFE",
+        )
+    )
+    router = MagicMock()
+    router.model_names = ["prompt-check"]
+    router.acompletion = AsyncMock(return_value=None)
+    prompt_injection_detection.update_environment(router=router)
+    monkeypatch.setattr(litellm, "callbacks", [prompt_injection_detection])
+
+    data = {
+        "model": "m",
+        "messages": [{"role": "user", "content": "hello"}],
+        "temperature": 0.1,
+    }
+    out = await proxy_logging.during_call_hook(
+        data=data,
+        user_api_key_dict=make_user_api_key_auth(),
+        call_type="completion",
+    )
+
+    assert out is data
+    router.acompletion.assert_awaited_once()
+    assert router.acompletion.await_args.kwargs["model"] == "prompt-check"
+
+
+@pytest.mark.asyncio
+async def test_during_call_hook_skips_prompt_injection_for_mcp_calls(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    prompt_injection_detection = _OPTIONAL_PromptInjectionDetection(
+        prompt_injection_params=LiteLLMPromptInjectionParams(
+            llm_api_check=True,
+            llm_api_name="prompt-check",
+            llm_api_system_prompt="Return SAFE or UNSAFE",
+            llm_api_fail_call_string="UNSAFE",
+        )
+    )
+    router = MagicMock()
+    router.model_names = ["prompt-check"]
+    router.acompletion = AsyncMock(return_value=None)
+    prompt_injection_detection.update_environment(router=router)
+    monkeypatch.setattr(litellm, "callbacks", [prompt_injection_detection])
+
+    data = {"model": "m", "name": "tool", "arguments": {"query": "hello"}}
+    out = await proxy_logging.during_call_hook(
+        data=data,
+        user_api_key_dict=make_user_api_key_auth(),
+        call_type="call_mcp_tool",
+    )
+
+    assert out is data
+    router.acompletion.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_during_call_hook_guardrail_skipped_when_should_not_run(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
     g = _make_guardrail("g", should_run=False)
     monkeypatch.setattr(litellm, "callbacks", [g])
     await proxy_logging.during_call_hook(
@@ -74,7 +164,9 @@ async def test_during_call_hook_guardrail_skipped_when_should_not_run(proxy_logg
 
 
 @pytest.mark.asyncio
-async def test_during_call_hook_guardrail_error_raises(proxy_logging, make_user_api_key_auth, monkeypatch):
+async def test_during_call_hook_guardrail_error_raises(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
     g = _make_guardrail("bad")
     g.async_moderation_hook = AsyncMock(side_effect=RuntimeError("blocked"))
     monkeypatch.setattr(litellm, "callbacks", [g])