fix(guardrails): honor streaming settings on the Noma v2 guardrail (#46606)

* fix(guardrails): honor streaming settings on the Noma v2 guardrail

The Noma v2 initializer now passes streaming_buffer_until_moderated,
streaming_end_of_stream_only, streaming_sampling_rate and
streaming_buffer_release_on_scan from the guardrail config to
NomaV2Guardrail, so the unified streaming hook applies them. The
constructor validates them, so a custom guardrail class built on
NomaV2Guardrail gets the same checks. Unset keys keep the hook defaults,
and an invalid value is ignored with a warning naming the guardrail, so
the guardrail still loads.

Resolves LIT-9633

* test(guardrails): cover Noma v2 streaming settings across clients, failures and workers

Adds integration cells for the custom-subclass loader, string and invalid
values, unset and false keys, a Noma outage with and without block_failures,
an upstream rejection, n=2 choices, pre_call and non-streaming guardrails,
the OpenAI and Anthropic SDKs (sync and async), stored-guardrail info, a
PATCH while traffic flows, a burst while Noma is down, a slow Noma, and
replaced workers.

Resolves LIT-9633

* test(guardrails): pin the PATCH probes to one worker at a time

Each probe suspends the other worker and opens a fresh connection, so the
test checks that every worker applies the patched setting.

Resolves LIT-9633

* test(guardrails): check sampled scans by count, not window position

/v1/messages hands the streaming hook raw network reads, so where the
default sampling windows fall depends on timing; the base leg caught one
shifted window. Check the scan count and the final full-text scan instead.

Resolves LIT-9633

---------

Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>

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

@@ -38,6 +38,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
 def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardrail"):
     import litellm
 
+    extras: Final = litellm_params.model_extra or {}
     _noma_v2_callback: Final = NomaV2Guardrail(
         guardrail_name=guardrail.get("guardrail_name", ""),
         api_key=litellm_params.api_key,
@@ -49,6 +50,10 @@ def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardra
         event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
         timeout=litellm_params.timeout,
+        streaming_buffer_until_moderated=extras.get("streaming_buffer_until_moderated"),
+        streaming_end_of_stream_only=extras.get("streaming_end_of_stream_only"),
+        streaming_sampling_rate=extras.get("streaming_sampling_rate"),
+        streaming_buffer_release_on_scan=extras.get("streaming_buffer_release_on_scan"),
     )
     litellm.logging_callback_manager.add_litellm_callback(_noma_v2_callback)
 

litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py

@@ -26,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import (
 from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage
 from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs
 from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.proxy.guardrails.guardrail_hooks.noma import NomaV2StreamingParams
 from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
 
 if TYPE_CHECKING:
@@ -63,6 +64,10 @@ class NomaV2Guardrail(CustomGuardrail):
         monitor_mode: bool | None = None,
         block_failures: bool | None = None,
         gateway_name: str | None = None,
+        streaming_buffer_until_moderated: bool | None = None,
+        streaming_end_of_stream_only: bool | None = None,
+        streaming_sampling_rate: int | None = None,
+        streaming_buffer_release_on_scan: bool | None = None,
         **kwargs: Any,
     ) -> None:
         self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
@@ -86,6 +91,24 @@ class NomaV2Guardrail(CustomGuardrail):
         if self._requires_api_key(api_base=self.api_base) and not self.api_key:
             raise ValueError("Noma v2 guardrail requires api_key when using Noma SaaS endpoint")
 
+        streaming: Final = NomaV2StreamingParams.model_validate(
+            {
+                "streaming_buffer_until_moderated": streaming_buffer_until_moderated,
+                "streaming_end_of_stream_only": streaming_end_of_stream_only,
+                "streaming_sampling_rate": streaming_sampling_rate,
+                "streaming_buffer_release_on_scan": streaming_buffer_release_on_scan,
+            },
+            context={"guardrail_name": kwargs.get("guardrail_name")},
+        )
+        if streaming.streaming_buffer_until_moderated is not None:
+            self.streaming_buffer_until_moderated: bool = streaming.streaming_buffer_until_moderated
+        if streaming.streaming_end_of_stream_only is not None:
+            self.streaming_end_of_stream_only: bool = streaming.streaming_end_of_stream_only
+        if streaming.streaming_sampling_rate is not None:
+            self.streaming_sampling_rate: int = streaming.streaming_sampling_rate
+        if streaming.streaming_buffer_release_on_scan is not None:
+            self.streaming_buffer_release_on_scan: bool = streaming.streaming_buffer_release_on_scan
+
         kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
 
         base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs

litellm/types/proxy/guardrails/guardrail_hooks/noma.py

@@ -1,4 +1,16 @@
-from pydantic import Field
+from typing import Final
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationError,
+    ValidationInfo,
+    ValidatorFunctionWrapHandler,
+    field_validator,
+)
+
+from litellm._logging import verbose_proxy_logger
 
 from .base import GuardrailConfigModel
 
@@ -55,3 +67,35 @@ class NomaV2GuardrailConfigModel(GuardrailConfigModel):
     @staticmethod
     def ui_friendly_name() -> str:
         return "Noma Security v2"
+
+
+class NomaV2StreamingParams(BaseModel):
+    """Streaming settings a Noma v2 guardrail reads from the top level of its litellm_params.
+
+    A key left unset, or set to a value that does not validate, stays None, so the unified
+    streaming hook keeps its own default for it and the guardrail still loads.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    streaming_buffer_until_moderated: bool | None = None
+    streaming_end_of_stream_only: bool | None = None
+    streaming_sampling_rate: int | None = Field(default=None, ge=1)
+    streaming_buffer_release_on_scan: bool | None = None
+
+    @field_validator("*", mode="wrap")
+    @classmethod
+    def _ignore_invalid_value(
+        cls, value: object, handler: ValidatorFunctionWrapHandler, info: ValidationInfo
+    ) -> object:
+        try:
+            validated: Final = handler(value)
+        except ValidationError:
+            verbose_proxy_logger.warning(
+                "Noma v2 guardrail %s: ignoring %s=%.100r, which is not a valid value; the default applies",
+                info.context.get("guardrail_name") if info.context else None,
+                info.field_name,
+                value,
+            )
+            return None
+        return validated

tests/unit/proxy/guardrails/guardrail_hooks/noma/test_noma_v2.py

@@ -1,17 +1,30 @@
 import json
+from collections.abc import Callable, Mapping
+from pathlib import Path
+from typing import TypeAlias
 
 import httpx
 import pytest
 import respx
 
 import litellm
+from litellm._logging import verbose_proxy_logger
+from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.proxy.guardrails.guardrail_hooks.noma import (
     NomaV2Guardrail,
     guardrail_initializer_registry,
 )
+from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
+from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
 from litellm.types.guardrails import LitellmParams
 
 _API_BASE = "https://noma.example.test"
+_STREAMING_KEYS = (
+    "streaming_buffer_until_moderated",
+    "streaming_end_of_stream_only",
+    "streaming_sampling_rate",
+    "streaming_buffer_release_on_scan",
+)
 
 
 @pytest.fixture(autouse=True)
@@ -97,3 +110,111 @@ async def test_positional_args_keep_their_meaning_after_gateway_name_was_added(r
         await guardrail.apply_guardrail(
             inputs={"texts": ["hello"]}, request_data={"metadata": {}}, input_type="request"
         )
+
+
+_Loader: TypeAlias = Callable[[Path, Mapping[str, object]], CustomGuardrail | None]
+_SUBCLASS_MODULE = (
+    "from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail\n\n\n"
+    "class SubclassedNoma(NomaV2Guardrail):\n"
+    "    pass\n"
+)
+
+
+def _initializer(guardrail_type: str, fixed: Mapping[str, object]) -> _Loader:
+    def load(_tmp_path: Path, params: Mapping[str, object]) -> CustomGuardrail | None:
+        litellm_params = LitellmParams(
+            guardrail=guardrail_type, mode="post_call", api_base=_API_BASE, **fixed, **params
+        )
+        return guardrail_initializer_registry[guardrail_type](litellm_params, {"guardrail_name": "noma-guard"})
+
+    return load
+
+
+def _custom_subclass(tmp_path: Path, params: Mapping[str, object]) -> CustomGuardrail | None:
+    (tmp_path / "noma_subclass.py").write_text(_SUBCLASS_MODULE)
+    guardrail_type = "noma_subclass.SubclassedNoma"
+    return InMemoryGuardrailHandler().initialize_custom_guardrail(
+        guardrail={"guardrail_name": "noma-guard"},
+        guardrail_type=guardrail_type,
+        litellm_params=LitellmParams(guardrail=guardrail_type, mode="post_call", api_base=_API_BASE, **params),
+        config_file_path=str(tmp_path / "config.yaml"),
+    )
+
+
+_LOADERS = pytest.mark.parametrize(
+    "load",
+    [_initializer("noma_v2", {}), _initializer("noma", {"use_v2": True}), _custom_subclass],
+    ids=["noma_v2", "noma-use_v2", "custom-subclass"],
+)
+
+
+def _streaming_flags(load: _Loader, tmp_path: Path, params: Mapping[str, object]) -> list[object]:
+    guardrail = load(tmp_path, params)
+    return [UnifiedLLMGuardrails().resolve_streaming_flag(guardrail, key, "hook-default") for key in _STREAMING_KEYS]
+
+
+@_LOADERS
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [
+        (
+            {
+                "streaming_buffer_until_moderated": True,
+                "streaming_end_of_stream_only": True,
+                "streaming_sampling_rate": 3,
+                "streaming_buffer_release_on_scan": True,
+            },
+            [True, True, 3, True],
+        ),
+        (
+            {
+                "streaming_buffer_until_moderated": False,
+                "streaming_end_of_stream_only": False,
+                "streaming_sampling_rate": 1,
+                "streaming_buffer_release_on_scan": False,
+            },
+            [False, False, 1, False],
+        ),
+        ({"streaming_sampling_rate": 7}, ["hook-default", "hook-default", 7, "hook-default"]),
+        (
+            {"streaming_buffer_until_moderated": "false", "streaming_sampling_rate": "2"},
+            [False, "hook-default", 2, "hook-default"],
+        ),
+        ({}, ["hook-default"] * 4),
+    ],
+    ids=["all-on", "all-off", "one-key", "string-values", "unset"],
+)
+def test_streaming_settings_in_guardrail_config_reach_the_streaming_hook(
+    load: _Loader, tmp_path: Path, configured: dict[str, object], expected: list[object]
+) -> None:
+    assert _streaming_flags(load, tmp_path, configured) == expected
+
+
+@_LOADERS
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [
+        ({"streaming_sampling_rate": 0}, ["hook-default"] * 4),
+        ({"streaming_sampling_rate": -1}, ["hook-default"] * 4),
+        ({"streaming_sampling_rate": "often"}, ["hook-default"] * 4),
+        ({"streaming_buffer_until_moderated": "sometimes"}, ["hook-default"] * 4),
+        (
+            {"streaming_buffer_until_moderated": True, "streaming_sampling_rate": 0},
+            [True, "hook-default", "hook-default", "hook-default"],
+        ),
+    ],
+    ids=["rate-zero", "rate-negative", "rate-word", "buffer-word", "valid-key-kept"],
+)
+def test_invalid_streaming_value_is_ignored_and_the_guardrail_still_loads(
+    load: _Loader,
+    tmp_path: Path,
+    configured: dict[str, object],
+    expected: list[object],
+    caplog: pytest.LogCaptureFixture,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(verbose_proxy_logger, "propagate", True)
+    assert _streaming_flags(load, tmp_path, configured) == expected
+    warnings = [record.getMessage() for record in caplog.records if "not a valid value" in record.getMessage()]
+    assert warnings, caplog.text
+    assert all(message.startswith("Noma v2 guardrail noma-guard: ignoring ") for message in warnings), warnings