Merge 495bd070ad into f52e5057f9

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

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+from litellm.types.guardrails import (
+    GuardrailEventHooks,
+    Mode,
+    SupportedGuardrailIntegrations,
+)
+
+from .needlepath import NeedlepathGuardrail
+
+if TYPE_CHECKING:
+    from litellm.types.guardrails import Guardrail, LitellmParams
+
+
+def _coerce_event_hook(
+    mode: str | list[str] | Mode,  # mutable-ok: the type of LitellmParams.mode
+) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode:  # mutable-ok: dispatch needs a list
+    if isinstance(mode, Mode):
+        return mode
+    if isinstance(mode, list):
+        # A real `list` is required here, not just any Sequence: CustomGuardrail's
+        # dispatch (should_run_guardrail) does `isinstance(self.event_hook, list)`
+        # to decide whether this hook fires for a given event. A tuple would fail
+        # that check silently and the guardrail would never run for a multi-hook
+        # config.
+        return [GuardrailEventHooks(item) for item in mode]
+    return GuardrailEventHooks(mode)
+
+
+def _get_optional_value(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> object:
+    """Read a knob from ``optional_params`` first, then from the top level.
+
+    Both spellings are accepted because both appear in the wild: the nested
+    ``optional_params`` block is the documented form, and several deployments
+    set guardrail knobs flat alongside ``api_key``.
+    """
+    if optional_params is not None:
+        value: Final[object] = getattr(optional_params, attribute_name, None)  # pyright: ignore[reportAny]  # knob name is dynamic
+        if value is not None:
+            return value
+    flat: Final[object] = getattr(litellm_params, attribute_name, None)  # pyright: ignore[reportAny]  # flat knobs are pydantic extras
+    return flat
+
+
+def _bool_knob(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> bool | None:
+    value: Final = _get_optional_value(litellm_params, optional_params, attribute_name)
+    return value if isinstance(value, bool) else None
+
+
+def _int_knob(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> int | None:
+    value: Final = _get_optional_value(litellm_params, optional_params, attribute_name)
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, int):
+        return value
+    if isinstance(value, str):
+        try:
+            return int(value)
+        except ValueError:
+            return None
+    return None
+
+
+def _str_knob(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> str | None:
+    value: Final = _get_optional_value(litellm_params, optional_params, attribute_name)
+    return value if isinstance(value, str) else None
+
+
+def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> NeedlepathGuardrail:
+    import litellm
+
+    optional_params: Final = getattr(litellm_params, "optional_params", None)
+
+    _callback: Final = NeedlepathGuardrail(
+        api_base=litellm_params.api_base,
+        api_key=litellm_params.api_key,
+        select_tool_outputs=_bool_knob(litellm_params, optional_params, "select_tool_outputs"),
+        select_history=_bool_knob(litellm_params, optional_params, "select_history"),
+        select_system=_bool_knob(litellm_params, optional_params, "select_system"),
+        min_chars_to_select=_int_knob(litellm_params, optional_params, "min_chars_to_select"),
+        max_context_tokens=_int_knob(litellm_params, optional_params, "max_context_tokens"),
+        operating_point=_str_knob(litellm_params, optional_params, "operating_point"),
+        guardrail_name=guardrail["guardrail_name"],
+        event_hook=_coerce_event_hook(litellm_params.mode),
+        default_on=litellm_params.default_on or False,
+    )
+    litellm.logging_callback_manager.add_litellm_callback(  # pyright: ignore[reportUnknownMemberType]  # callback manager is untyped
+        _callback
+    )
+    return _callback
+
+
+# guardrail_registry.py's module discovery does `isinstance(registry, dict)` before
+# merging this into the central registry; a Mapping/MappingProxyType would fail that
+# check silently and this guardrail would never be discovered at proxy startup.
+guardrail_initializer_registry: Final = {
+    SupportedGuardrailIntegrations.NEEDLEPATH.value: initialize_guardrail,
+}
+
+# Same reason as guardrail_initializer_registry above: discovery's isinstance(dict) check.
+guardrail_class_registry: Final = {
+    SupportedGuardrailIntegrations.NEEDLEPATH.value: NeedlepathGuardrail,
+}

litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py

@@ -0,0 +1,711 @@
+"""Needlepath guardrail: query-conditioned extractive context selection.
+
+Bulky message content (tool outputs by default) is sent to the Needlepath
+selection service before the request reaches the LLM. The service returns the
+spans of that content which carry the answer to a query, and the guardrail
+writes those spans back over the message they came from. Nothing is
+paraphrased or rewritten: the returned block is made of extracts of the text
+that was submitted.
+
+Selection is per message and query-conditioned. The query for a tool output is
+the intent of the tool call that produced it (``name`` plus ``arguments``,
+found through ``tool_call_id``); anything else uses the last user message. Each
+message is selected independently, so one message's outcome never changes
+another's.
+
+**This guardrail is unconditionally fail-open.** Every path that does not
+produce a usable selection returns the caller's messages untouched. A proxy
+that silently blanks a tool output is far worse than a proxy that does nothing,
+so there is no configuration in which a selection failure becomes a request
+failure. See ``_selected_text`` for the full list of declines.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import ipaddress
+import time
+from collections.abc import Mapping, Sequence
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final, Literal
+from urllib.parse import urlparse
+
+import httpx
+from httpx import Response as HttpxResponse
+from pydantic import TypeAdapter, ValidationError
+
+import litellm
+from litellm._logging import verbose_proxy_logger
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    log_guardrail_information,  # pyright: ignore[reportUnknownVariableType]  # decorator is untyped in custom_guardrail
+)
+from litellm.llms.custom_httpx.http_handler import (
+    get_async_httpx_client,  # pyright: ignore[reportUnknownVariableType]  # helper is untyped in http_handler
+    httpxSpecialProvider,
+)
+from litellm.proxy.guardrails.guardrail_hooks.content_text import (
+    content_to_text,
+    is_all_text_parts,
+    merge_rewritten_text_parts,
+)
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.guardrails import GuardrailEventHooks, Mode
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import (
+        Logging as LiteLLMLoggingObj,
+    )
+    from litellm.types.proxy.guardrails.guardrail_hooks.needlepath import (
+        NeedlepathGuardrailConfigModel,
+    )
+
+DEFAULT_API_BASE: Final = "https://api.nextmoca.com"
+SELECT_PATH: Final = "/v1/context/select"
+# Immutable engine label. Pinned rather than inherited from the service default
+# so an upgrade on the service side cannot change what this proxy sends without
+# an operator changing this config.
+DEFAULT_OPERATING_POINT: Final = "np-2026-10-r8"
+DEFAULT_MAX_CONTEXT_TOKENS: Final = 4000
+DEFAULT_MIN_CHARS_TO_SELECT: Final = 500
+# The shared client's read timeout is measured in minutes, which is far too long
+# to hold an inbound LLM request behind an optional optimisation. A stall past
+# this bound is a decline, and the original message is forwarded.
+_SELECT_TIMEOUT_SECONDS: Final = 30.0
+# A single proxy request can carry arbitrarily many eligible messages, and each
+# one becomes an outbound selection call. These two bounds keep a pathological
+# request (say, thousands of minimally qualifying tool outputs) from
+# monopolising the proxy's shared HTTP pool or burning selection quota: at most
+# _MAX_TARGETS_PER_REQUEST messages are selected per request -- the largest
+# first, where selection pays off most -- and at most
+# _MAX_CONCURRENT_SELECTIONS calls are in flight at once. Messages past the
+# cap are forwarded untouched, the same fail-open outcome as any decline.
+_MAX_TARGETS_PER_REQUEST: Final = 16
+_MAX_CONCURRENT_SELECTIONS: Final = 4
+# The service reports a deliberate no-op through the gate. Any reason under this
+# prefix means "the engine chose not to select"; the original content is what
+# the caller should send.
+_STANDDOWN_PREFIX: Final = "standdown:"
+# Security screening outcomes (top-level ``outcome``). A ``sanitized_fallback``
+# is the caller's full context with every flagged span replaced by a marker; it
+# is applied, because forwarding the original would put the flagged text back.
+# ``blocked`` delivers none of the caller's content and is a decline, the same
+# handling the published Needlepath clients give it.
+_OUTCOME_SANITIZED_FALLBACK: Final = "sanitized_fallback"
+_OUTCOME_BLOCKED: Final = "blocked"
+# ``format_metrics.security_action`` when screening removed flagged spans from
+# what it delivered.
+_SECURITY_ACTION_QUARANTINE: Final = "quarantine"
+# Request and response JSON is untrusted. These adapters turn a value of
+# unknown shape into a typed dict or list, or into None when it is not one.
+_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
+_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
+_BLOCKED_METADATA_HOSTS: Final = frozenset(
+    {
+        "metadata.google.internal",
+        "metadata.goog",
+        "metadata.azure.com",
+        "metadata.azure.internal",
+    }
+)
+_BLOCKED_METADATA_IPS: Final = frozenset(
+    ipaddress.ip_address(ip) for ip in ("169.254.169.254", "fd00:ec2::254", "100.100.100.200", "168.63.129.16")
+)
+# Record kinds the service publishes. A system prompt has no dedicated kind, so
+# it is submitted as external_data rather than invented as a new one: an
+# unrecognised kind is a 400 for the whole call.
+_KIND_TOOL_RESULT: Final = "tool_result"
+_KIND_USER_INPUT: Final = "user_input"
+_KIND_EXTERNAL_DATA: Final = "external_data"
+
+
+def _as_str_object_dict(value: object) -> Mapping[str, object] | None:
+    try:
+        return _STR_OBJECT_DICT_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def _as_object_list(value: object) -> Sequence[object] | None:
+    try:
+        return _OBJECT_LIST_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def _parse_ip_literal(host: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None:
+    """Parse ``host`` as an IP literal, including the alternate spellings the
+    socket layer accepts (single-integer IPv4, IPv4-mapped IPv6), so a blocked
+    address cannot be smuggled past a plain string comparison."""
+    try:
+        addr = ipaddress.ip_address(host)  # rebind-ok: reassigned in the except fallback below
+    except ValueError:
+        try:
+            addr = ipaddress.ip_address(int(host, 0))  # rebind-ok: fallback re-parse as an int literal
+        except (TypeError, ValueError):
+            return None
+    if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None:
+        return addr.ipv4_mapped
+    return addr
+
+
+def _validate_api_base(url: str) -> str:
+    """Return ``url`` if it passes basic outbound-target checks, else raise.
+
+    Defense in depth against a mistyped or hostile ``api_base``: non-http(s)
+    schemes and cloud-metadata hosts/IPs are refused. Private ranges stay
+    allowed so on-prem deployments work. This is not a complete SSRF control:
+    there is no DNS resolution here and the shared client follows redirects.
+    ``api_base`` is operator config, so that is an accepted limit.
+    """
+    parsed: Final = urlparse(url)
+    if parsed.scheme not in ("http", "https"):
+        raise ValueError(f"Needlepath guardrail api_base must be http or https, got scheme={parsed.scheme!r}")
+    host: Final = (parsed.hostname or "").lower()
+    if not host:
+        raise ValueError("Needlepath guardrail api_base has no host")
+    ip_literal: Final = _parse_ip_literal(host)
+    if host in _BLOCKED_METADATA_HOSTS or (ip_literal is not None and ip_literal in _BLOCKED_METADATA_IPS):
+        raise ValueError(f"Needlepath guardrail api_base {host!r} is a blocked cloud-metadata host")
+    return url
+
+
+def _write_text_back(content: object, new_text: str) -> object:
+    """Put ``new_text`` into a ``content`` value without changing its shape.
+
+    A string is replaced directly. An all-text part list collapses to a single
+    part carrying the last declared cache_control breakpoint. Anything else is
+    returned untouched: breakpoints are positional, so one selected block cannot
+    be written across a non-text part without moving text past it.
+    """
+    if isinstance(content, str):
+        return new_text
+    parts: Final = _as_object_list(content)
+    if parts is not None and is_all_text_parts(parts):
+        return merge_rewritten_text_parts(parts, new_text)
+    return content
+
+
+def _render_tool_intent(fn: Mapping[str, object]) -> str:
+    """A tool call rendered as the query its output should be selected against."""
+    name: Final = str(fn.get("name") or "").strip()
+    raw_args: Final = fn.get("arguments")
+    arg_text: Final = "" if raw_args is None else str(raw_args).strip()
+    if name and arg_text:
+        return f"{name}: {arg_text}"
+    return name or arg_text
+
+
+def _query_for_target(messages: Sequence[Mapping[str, object]], target_idx: int, fallback: str) -> str:
+    """The query ``messages[target_idx]`` should be selected against.
+
+    A tool or function result is selected against the intent of the call that
+    produced it, located by ``tool_call_id`` on an earlier assistant message.
+    Everything else, and any tool result whose call cannot be found, uses the
+    last user message.
+    """
+    msg: Final = messages[target_idx]
+    if msg.get("role") not in ("tool", "function"):
+        return fallback
+
+    tool_call_id: Final = msg.get("tool_call_id")
+    fn_name: Final = msg.get("name")
+    for idx in range(target_idx - 1, -1, -1):
+        previous = messages[idx]
+        if previous.get("role") != "assistant":
+            continue
+        for raw_call in _as_object_list(previous.get("tool_calls")) or ():
+            call = _as_str_object_dict(raw_call)
+            if call is None or not tool_call_id or call.get("id") != tool_call_id:
+                continue
+            fn = _as_str_object_dict(call.get("function"))
+            intent = _render_tool_intent(fn if fn is not None else MappingProxyType({}))
+            if intent:
+                return intent
+        # Legacy function_call turns carry no id, so require a name match.
+        # Without it an older, unrelated call would supply the wrong intent.
+        legacy = _as_str_object_dict(previous.get("function_call"))
+        if legacy is not None and fn_name and legacy.get("name") == fn_name:
+            intent = _render_tool_intent(legacy)
+            if intent:
+                return intent
+    return fallback
+
+
+def _title_for(messages: Sequence[Mapping[str, object]], target_idx: int) -> str | None:
+    """The tool name behind a message, used as the record title."""
+    msg: Final = messages[target_idx]
+    name: Final = msg.get("name")
+    if isinstance(name, str) and name.strip():
+        return name.strip()[:120]
+    tool_call_id: Final = msg.get("tool_call_id")
+    if not tool_call_id:
+        return None
+    for idx in range(target_idx - 1, -1, -1):
+        previous = messages[idx]
+        if previous.get("role") != "assistant":
+            continue
+        for raw_call in _as_object_list(previous.get("tool_calls")) or ():
+            call = _as_str_object_dict(raw_call)
+            if call is None or call.get("id") != tool_call_id:
+                continue
+            fn = _as_str_object_dict(call.get("function"))
+            fn_name = fn.get("name") if fn is not None else None
+            if isinstance(fn_name, str):
+                return fn_name[:120]
+    return None
+
+
+def _record_kind(role: object) -> str:
+    if role in ("tool", "function"):
+        return _KIND_TOOL_RESULT
+    if role == "user":
+        return _KIND_USER_INPUT
+    return _KIND_EXTERNAL_DATA
+
+
+def _safe_int(value: object) -> int | None:
+    """Read an integer counter from an untrusted body without raising.
+
+    A field that is missing or not a number is reported as ``None`` so the
+    caller can treat it as "unknown" rather than as zero, which is a decline.
+    """
+    if isinstance(value, bool) or not isinstance(value, (int, float, str)):
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def _screening_removed_content(body: Mapping[str, object]) -> bool:
+    """True when security screening removed flagged spans from the delivery.
+
+    Such a block must be applied even where an ordinary selection would be
+    declined: the declined path forwards the original content, flagged spans
+    included.
+    """
+    format_metrics: Final = _as_str_object_dict(body.get("format_metrics"))
+    return format_metrics is not None and format_metrics.get("security_action") == _SECURITY_ACTION_QUARANTINE
+
+
+def _safe_response_text(response: object, limit: int = 500) -> str:
+    """Read a response body for a log line without letting the read itself raise.
+
+    A corrupt ``Content-Encoding`` makes httpx's ``.text`` raise, which would
+    turn an already-handled decline into an unhandled 500.
+    """
+    try:
+        text: Final = getattr(response, "text", "")
+    except httpx.DecodingError:
+        return "<undecodable response body>"
+    return (text or "")[:limit]
+
+
+class NeedlepathGuardrail(CustomGuardrail):
+    """Select the spans of a message that answer the current query.
+
+    Every knob is optional and the defaults are deliberately narrow: tool
+    outputs only, nothing under ``min_chars_to_select`` characters, and a pinned
+    operating point.
+    """
+
+    def __init__(
+        self,
+        api_base: str | None = None,
+        api_key: str | None = None,
+        select_tool_outputs: bool | None = None,
+        select_history: bool | None = None,
+        select_system: bool | None = None,
+        min_chars_to_select: int | None = None,
+        max_context_tokens: int | None = None,
+        operating_point: str | None = None,
+        guardrail_name: str | None = None,
+        # A list, not a Sequence: CustomGuardrail's dispatch checks isinstance(list).
+        event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None,  # mutable-ok: see above
+        default_on: bool = False,
+    ) -> None:
+        raw_api_base: Final = (api_base or get_secret_str("NEEDLEPATH_API_BASE") or DEFAULT_API_BASE).rstrip("/")
+        self.needlepath_api_base = _validate_api_base(raw_api_base)
+        self.needlepath_api_key = api_key or get_secret_str("NEEDLEPATH_API_KEY")
+        if not self.needlepath_api_key:
+            raise ValueError(
+                "Needlepath guardrail requires an API key. Set `api_key` in the "
+                "guardrail config or the NEEDLEPATH_API_KEY env var."
+            )
+        self.select_tool_outputs = True if select_tool_outputs is None else select_tool_outputs
+        self.select_history = False if select_history is None else select_history
+        self.select_system = False if select_system is None else select_system
+        self.min_chars_to_select = (
+            DEFAULT_MIN_CHARS_TO_SELECT if min_chars_to_select is None else int(min_chars_to_select)
+        )
+        self.max_context_tokens = DEFAULT_MAX_CONTEXT_TOKENS if max_context_tokens is None else int(max_context_tokens)
+        # Pinned, not inherited. The labels are immutable, so a pinned one makes
+        # what this guardrail sends reproducible across service releases.
+        self.operating_point = operating_point or DEFAULT_OPERATING_POINT
+        self.async_handler = get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.GuardrailCallback,
+        )
+        super().__init__(  # pyright: ignore[reportUnknownMemberType]  # CustomGuardrail.__init__ is untyped
+            guardrail_name=guardrail_name,
+            event_hook=event_hook,
+            default_on=default_on,
+        )
+
+    def _request_headers(
+        self,
+    ) -> dict[str, str]:  # mutable-ok: the shared handler's `headers=` is typed as a dict; a fresh one per call
+        return {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.needlepath_api_key or ''}",
+        }
+
+    def _decline(self, reason: str, detail: Mapping[str, object] | None = None) -> None:
+        """Record a decline. The caller then forwards the original content.
+
+        Declining is the normal outcome for a stand-down or a service problem,
+        so it is logged at debug and never raised. Details can include upstream
+        response bytes and stay in the proxy's own logs.
+        """
+        # The dict fallback below is never a value a caller could mutate: `or {}` only
+        # feeds the debug-log %s formatting, and swapping it for a frozen mapping would
+        # change the logged repr from `{}` to `mappingproxy({})`.
+        verbose_proxy_logger.debug("Needlepath: declined (%s) detail=%s", reason, detail or {})
+
+    def _select_targets(self, messages: Sequence[Mapping[str, object]], query_idx: int | None) -> tuple[int, ...]:
+        """Indices of the messages whose text content is eligible for selection."""
+
+        def _eligible(idx: int, msg: Mapping[str, object]) -> bool:
+            # The message carrying the query is never rewritten: selecting the
+            # question against itself is meaningless and it is what the rest of
+            # the selection is conditioned on.
+            if idx == query_idx:
+                return False
+            role: Final = msg.get("role")
+            if role in ("tool", "function"):
+                if not self.select_tool_outputs:
+                    return False
+            elif role == "system":
+                if not self.select_system:
+                    return False
+            elif role == "user":
+                if not self.select_history:
+                    return False
+            else:
+                return False
+            content: Final = msg.get("content")
+            parts: Final = _as_object_list(content)
+            if parts is not None and not is_all_text_parts(parts):
+                return False
+            return len(content_to_text(content)) >= self.min_chars_to_select
+
+        return tuple(idx for idx, msg in enumerate(messages) if _eligible(idx, msg))
+
+    @staticmethod
+    def _last_user_message(messages: Sequence[Mapping[str, object]]) -> tuple[str, int | None]:
+        for idx in range(len(messages) - 1, -1, -1):
+            if messages[idx].get("role") == "user":
+                return content_to_text(messages[idx].get("content")), idx
+        return "", None
+
+    async def _post_select(self, payload: Mapping[str, object]) -> HttpxResponse | None:
+        """POST one selection request. Returns None on any transport failure."""
+        try:
+            return await self.async_handler.post(  # pyright: ignore[reportUnknownMemberType]  # AsyncHTTPHandler.post is untyped
+                url=f"{self.needlepath_api_base}{SELECT_PATH}",
+                json=dict(payload),
+                headers=self._request_headers(),
+                timeout=_SELECT_TIMEOUT_SECONDS,
+            )
+        except asyncio.CancelledError:
+            raise
+        except httpx.HTTPStatusError as e:
+            # The shared handler calls raise_for_status(), so 402, 403, 429 and
+            # every other non-2xx arrives here carrying the upstream body and
+            # our Authorization header. It is a decline and nothing about it
+            # reaches the client.
+            response: Final = getattr(e, "response", None)
+            self._decline(
+                "http_status",
+                {
+                    "status_code": getattr(response, "status_code", None),
+                    "body": _safe_response_text(response),
+                },
+            )
+            return None
+        except (httpx.RequestError, litellm.Timeout) as e:
+            # Every request-side httpx failure, timeouts included, is a
+            # RequestError. Catching the whole class is what keeps a network
+            # problem from escaping as a 500.
+            self._decline("transport", {"detail": str(e)})
+            return None
+        except Exception as e:  # noqa: BLE001  # fail-open: any client error is a decline
+            # Anything else the shared client raises is still a decline. A
+            # selection problem must never become a failed LLM request.
+            self._decline("unexpected_error", {"error_type": type(e).__name__, "detail": str(e)})
+            return None
+
+    async def _selected_text(self, text: str, query: str, title: str | None, source: object, kind: str) -> str | None:
+        """The selected block for one message, or None to keep the original.
+
+        None is returned, and the original content is forwarded, when:
+
+        * the service could not be reached, timed out, or answered non-2xx
+          (402, 403, 429 and everything else);
+        * the body is not JSON, or is JSON of an unexpected shape;
+        * ``outcome`` is ``blocked``: security screening left nothing
+          deliverable;
+        * the gate reports a stand-down (any ``gate.reason`` under
+          ``standdown:``), which is the service saying the full content is what
+          the caller should send;
+        * ``records_selected`` is 0, ``tokens_after`` is 0, or
+          ``rendered_context`` is missing or blank;
+        * the returned block is not shorter than the text it would replace, so a
+          rewrite could never grow a message.
+
+        A block from which security screening removed flagged spans is applied
+        whenever it is non-blank: ``outcome: sanitized_fallback``, or
+        ``format_metrics.security_action: quarantine``. The stand-down and
+        size checks above do not apply to it, because declining it would
+        forward the original content with the flagged spans still in it.
+        """
+        # `record` and `payload` are both real, mutable dicts by necessity: `record`
+        # gains optional keys after construction (below), and both are handed straight
+        # to httpx's `json=`, which needs a genuine dict/list -- a MappingProxyType
+        # is not JSON-serializable (json.dumps raises TypeError on one), so the
+        # freezing wrappers used elsewhere in this file cannot be used here.
+        record: Final[dict[str, object]] = {  # mutable-ok: see comment above
+            "id": "m0",
+            "kind": kind,
+            "text": text,
+        }
+        if title:
+            record["title"] = title
+        if isinstance(source, str) and source:
+            record["source"] = source
+        payload: Final[dict[str, object]] = {  # mutable-ok: see comment above
+            "records": (record,),
+            "task": {"prompt": query},
+            "budget": {
+                "max_context_tokens": self.max_context_tokens,
+                "operating_point": self.operating_point,
+            },
+            "render": True,
+            "render_format": "plain",
+        }
+
+        response: Final = await self._post_select(payload)
+        if response is None:
+            return None
+        if not 200 <= response.status_code < 300:
+            self._decline(
+                "http_status",
+                {
+                    "status_code": response.status_code,
+                    "body": _safe_response_text(response),
+                },
+            )
+            return None
+
+        try:
+            raw_body: Final[object] = response.json()  # pyright: ignore[reportAny]  # httpx Response.json() is untyped
+        except (ValueError, httpx.DecodingError, RecursionError):
+            # RecursionError: a deeply nested body overflows the JSON parser.
+            # It is a decline like any other malformed answer, not a 500.
+            self._decline("unreadable_body", {"body": _safe_response_text(response)})
+            return None
+        body: Final = _as_str_object_dict(raw_body)
+        if body is None:
+            self._decline("unexpected_shape", {"body": _safe_response_text(response)})
+            return None
+
+        outcome: Final = body.get("outcome")
+        rendered: Final = body.get("rendered_context")
+        has_rendered: Final = isinstance(rendered, str) and bool(rendered.strip())
+        if outcome == _OUTCOME_BLOCKED:
+            # Screening found nothing deliverable. The original content is
+            # forwarded, as the published clients do; the outcome is in the
+            # service's receipt for operators who want to alarm on it.
+            self._decline("blocked", MappingProxyType({}))
+            return None
+        if outcome == _OUTCOME_SANITIZED_FALLBACK:
+            # The engine stood down, so this carries no selected records and
+            # its gate reports a stand-down, and it can be about as long as the
+            # original. It is still the block to send: the original is the same
+            # content with the flagged spans still in it.
+            if not isinstance(rendered, str) or not has_rendered:
+                self._decline("empty_rendered_context", MappingProxyType({}))
+                return None
+            return rendered
+        screened: Final = _screening_removed_content(body)
+
+        gate: Final = _as_str_object_dict(body.get("gate"))
+        reason: Final = gate.get("reason") if gate is not None else None
+        if not screened and isinstance(reason, str) and reason.startswith(_STANDDOWN_PREFIX):
+            self._decline("gate_standdown", {"reason": reason})
+            return None
+
+        records_selected: Final = _safe_int(body.get("records_selected"))
+        tokens_after: Final = _safe_int(body.get("tokens_after"))
+        if records_selected == 0 or tokens_after == 0:
+            self._decline(
+                "empty_selection",
+                {
+                    "records_selected": records_selected,
+                    "tokens_after": tokens_after,
+                },
+            )
+            return None
+
+        if not isinstance(rendered, str) or not has_rendered:
+            self._decline("empty_rendered_context", MappingProxyType({}))
+            return None
+        if not screened and len(rendered) >= len(text):
+            # A selection that is not smaller is not a selection worth making,
+            # and applying it could only add tokens. A screened block is exempt:
+            # declining it would forward the flagged spans it removed.
+            self._decline("no_reduction", {"before": len(text), "after": len(rendered)})
+            return None
+        return rendered
+
+    @log_guardrail_information
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict[str, object],  # mutable-ok: base-class signature; logging info is written into it
+        input_type: Literal["request", "response"],
+        logging_obj: LiteLLMLoggingObj | None = None,
+    ) -> GenericGuardrailAPIInputs:
+        if input_type != "request":
+            return inputs
+
+        structured_messages: Final = _as_object_list(inputs.get("structured_messages"))
+        if not structured_messages:
+            return inputs
+        # The static type says every entry is an AllMessageValues TypedDict, but this
+        # is the proxy's client-facing boundary: the list arrives from arbitrary
+        # request JSON, so the runtime check is the guarantee the annotation only
+        # promises. Any entry that is not a dict leaves the request untouched.
+        messages: Final = tuple(m for raw in structured_messages if (m := _as_str_object_dict(raw)) is not None)
+        if len(messages) != len(structured_messages):
+            return inputs
+
+        fallback_query, query_idx = self._last_user_message(messages)
+        # One (idx, query) pair per eligible message whose query is non-blank.
+        # Selection is conditioned on a query; without one there is nothing to
+        # select against, so that message is left exactly as it arrived.
+        eligible_pairs: Final = tuple(
+            (idx, query)
+            for idx in self._select_targets(messages, query_idx)
+            if (query := _query_for_target(messages, idx, fallback_query)).strip()
+        )
+        if not eligible_pairs:
+            verbose_proxy_logger.debug("Needlepath: no messages eligible for selection")
+            return inputs
+
+        # Cap the fan-out: the largest messages are kept because that is where
+        # selection saves the most, and the choice must not depend on message
+        # order in the request. Everything past the cap is forwarded untouched.
+        selected_pairs: Final = (
+            eligible_pairs
+            if len(eligible_pairs) <= _MAX_TARGETS_PER_REQUEST
+            else tuple(
+                sorted(
+                    eligible_pairs,
+                    key=lambda pair: len(content_to_text(messages[pair[0]].get("content"))),
+                    reverse=True,
+                )[:_MAX_TARGETS_PER_REQUEST]
+            )
+        )
+        if len(selected_pairs) < len(eligible_pairs):
+            verbose_proxy_logger.debug(
+                "Needlepath: %d eligible messages, selecting only the %d largest",
+                len(eligible_pairs),
+                _MAX_TARGETS_PER_REQUEST,
+            )
+
+        semaphore: Final = asyncio.Semaphore(_MAX_CONCURRENT_SELECTIONS)
+
+        async def _bounded_selected_text(idx: int, query: str) -> str | None:
+            # The semaphore bounds how many selection calls this one proxy
+            # request holds open at a time; see _MAX_CONCURRENT_SELECTIONS.
+            async with semaphore:
+                return await self._selected_text(
+                    text=content_to_text(messages[idx].get("content")),
+                    query=query,
+                    title=_title_for(messages, idx),
+                    source=messages[idx].get("tool_call_id"),
+                    kind=_record_kind(messages[idx].get("role")),
+                )
+
+        start_time: Final = time.monotonic()
+        # One request per message: each carries its own query, and the service
+        # renders one block per call. Running them concurrently (up to the
+        # semaphore's bound) keeps the added latency near one round trip
+        # rather than one per message.
+        selections: Final = await asyncio.gather(*(_bounded_selected_text(idx, query) for idx, query in selected_pairs))
+        end_time: Final = time.monotonic()
+
+        # Needs in-place index assignment below to build the edited copy. It starts
+        # from the caller's own message objects, so only replaced indices differ
+        # from the original list -- see the identity note below.
+        selected_messages: Final[list[object]] = list(structured_messages)  # mutable-ok: see comment above
+        messages_selected = 0  # rebind-ok: loop accumulator, incremented per selection below
+        chars_before = 0  # rebind-ok: loop accumulator, incremented per selection below
+        chars_after = 0  # rebind-ok: loop accumulator, incremented per selection below
+        for (idx, _query), selection in zip(selected_pairs, selections):
+            original = messages[idx]
+            original_text = content_to_text(original.get("content"))
+            if selection is None:
+                continue
+            messages_selected += 1
+            chars_before += len(original_text)
+            chars_after += len(selection)
+            selected_messages[idx] = {
+                **original,
+                "content": _write_text_back(original.get("content"), selection),
+            }
+
+        if messages_selected == 0:
+            # Nothing was replaced, so hand back the exact inputs object. The
+            # handlers detect a guardrail edit by identity, and a fresh list
+            # would force a write-back of an untouched request (on Anthropic
+            # that reconversion strips cache_control from thinking blocks).
+            verbose_proxy_logger.debug("Needlepath: no selection applied; request unchanged")
+            return inputs
+
+        # Serialized as guardrail_json_response for standard logging (DataDog,
+        # Langfuse, ...); must stay a plain dict for json.dumps compatibility.
+        stats: Final[dict[str, object]] = {  # mutable-ok: see comment above
+            "messages_selected": messages_selected,
+            "messages_considered": len(selected_pairs),
+            "chars_before": chars_before,
+            "chars_after": chars_after,
+            "operating_point": self.operating_point,
+        }
+        verbose_proxy_logger.debug(
+            "Needlepath: selected %s of %s message(s), %s -> %s chars",
+            messages_selected,
+            len(selected_pairs),
+            chars_before,
+            chars_after,
+        )
+        self.add_standard_logging_guardrail_information_to_request_data(  # pyright: ignore[reportUnknownMemberType]  # untyped base helper
+            guardrail_json_response=stats,
+            request_data=request_data,
+            guardrail_status="success",
+            guardrail_provider="needlepath",
+            start_time=start_time,
+            end_time=end_time,
+            duration=end_time - start_time,
+        )
+        return {**inputs, "structured_messages": selected_messages}  # pyright: ignore[reportReturnType]  # plain dicts satisfy AllMessageValues at runtime
+
+    @staticmethod
+    def get_config_model() -> type[NeedlepathGuardrailConfigModel] | None:
+        from litellm.types.proxy.guardrails.guardrail_hooks.needlepath import (
+            NeedlepathGuardrailConfigModel,
+        )
+
+        return NeedlepathGuardrailConfigModel

litellm/types/guardrails.py

@@ -47,6 +47,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ibm import (
 from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import (
     ContentFilterCategoryConfig,
 )
+from litellm.types.proxy.guardrails.guardrail_hooks.needlepath import (
+    NeedlepathGuardrailConfigModel,
+)
 from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import (
     OvalixGuardrailConfigModel,
 )
@@ -163,6 +166,7 @@ class SupportedGuardrailIntegrations(Enum):
     AIRIA = "airia"
     WINGBACK = "wingback"
     DECISION_MODEL = "decision_model"
+    NEEDLEPATH = "needlepath"
 
 
 class Role(Enum):
@@ -1397,6 +1401,7 @@ class LitellmParams(  # pyright: ignore[reportIncompatibleVariableOverride]  # o
     HeadroomGuardrailConfigModel,
     CompresrGuardrailConfigModel,
     TypeSafeGuardrailConfigModel,
+    NeedlepathGuardrailConfigModel,
     RepelloAIGuardrailConfigModel,
     LassoGuardrailConfigModel,
     DeepKeepGuardrailConfigModel,

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

@@ -0,0 +1,56 @@
+from pydantic import BaseModel, ConfigDict, Field
+
+from .base import GuardrailConfigModel
+
+
+class NeedlepathGuardrailOptionalParams(BaseModel):
+    """Optional tuning knobs for the Needlepath guardrail."""
+
+    model_config = ConfigDict(frozen=True)
+
+    select_tool_outputs: bool | None = Field(
+        default=None,
+        description="Select over tool/function result messages (search hits, RAG chunks, API dumps). Defaults to True.",
+    )
+    select_history: bool | None = Field(
+        default=None,
+        description="Also select over prior (non-last) user messages. Defaults to False.",
+    )
+    select_system: bool | None = Field(
+        default=None,
+        description="Also select over system messages. Defaults to False.",
+    )
+    min_chars_to_select: int | None = Field(
+        default=None,
+        description="Skip messages whose text is shorter than this many characters. Defaults to 500.",
+    )
+    max_context_tokens: int | None = Field(
+        default=None,
+        description="Token budget requested for the selected block of a single message. Defaults to 4000.",
+    )
+    operating_point: str | None = Field(
+        default=None,
+        description=(
+            "Immutable engine label sent with every request. Defaults to 'np-2026-10-r8'. "
+            "Pinned rather than inherited from the service default, so what this guardrail "
+            "sends does not change underneath a deployment."
+        ),
+    )
+
+
+class NeedlepathGuardrailConfigModel(GuardrailConfigModel[NeedlepathGuardrailOptionalParams]):
+    api_key: str | None = Field(
+        default=None,
+        description="Needlepath API key. Falls back to the NEEDLEPATH_API_KEY env var.",
+    )
+    api_base: str | None = Field(
+        default=None,
+        description=(
+            "Base URL of the Needlepath API. Falls back to the NEEDLEPATH_API_BASE env var, "
+            "then https://api.nextmoca.com."
+        ),
+    )
+
+    @staticmethod
+    def ui_friendly_name() -> str:
+        return "Needlepath (context selection)"

tests/unit/proxy/guardrails/guardrail_hooks/test_needlepath.py

@@ -0,0 +1,683 @@
+"""
+Unit tests for the Needlepath guardrail.
+
+Tests cover:
+- apply_guardrail replaces an eligible message with the selected block, query-
+  conditioned on the intent of the tool call that produced it
+- target selection: tool outputs by default, system/history opt-in, min-chars
+  threshold, the query message itself is never rewritten
+- the fail-open contract, one test per decline: empty selection
+  (records_selected 0, tokens_after 0, blank or missing rendered_context), a
+  gate stand-down, a ``blocked`` security outcome, HTTP 402/403/429 and other
+  non-2xx, timeouts and transport errors, malformed JSON, and an unexpected
+  response schema. Every one asserts the messages come back byte-identical.
+- security screening: a ``sanitized_fallback``, or an engaged block that
+  screening quarantined, is applied even where an ordinary selection would be
+  declined, because declining would forward the flagged spans.
+- /v1/responses: a function_call_output is selected and written back through
+  LiteLLM's Responses translation.
+"""
+
+import asyncio
+import copy
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+import pytest
+
+from litellm.proxy.guardrails.guardrail_hooks.needlepath.needlepath import (
+    _MAX_CONCURRENT_SELECTIONS,
+    _MAX_TARGETS_PER_REQUEST,
+    DEFAULT_MAX_CONTEXT_TOKENS,
+    DEFAULT_MIN_CHARS_TO_SELECT,
+    DEFAULT_OPERATING_POINT,
+    NeedlepathGuardrail,
+)
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+FAKE_API_BASE = "https://needlepath.example.com"
+FAKE_API_KEY = "np_live_0123456789abcdefghijklmnopqr"
+
+TOOL_OUTPUT = "Result row: quarterly filing extract. " * 40  # > 500 chars
+SELECTED_BLOCK = "Q3 revenue was 4.2 billion."
+USER_QUESTION = "What was Q3 revenue?"
+
+AGENT_MESSAGES = [
+    {"role": "system", "content": "You are a filings analyst. " * 40},
+    {"role": "user", "content": USER_QUESTION},
+    {
+        "role": "assistant",
+        "content": None,
+        "tool_calls": [
+            {
+                "id": "call_1",
+                "type": "function",
+                "function": {"name": "edgar_search", "arguments": '{"cik": "0000320193"}'},
+            }
+        ],
+    },
+    {"role": "tool", "tool_call_id": "call_1", "content": TOOL_OUTPUT},
+]
+
+TOOL_INDEX = 3
+
+
+def _make_guardrail(**kwargs) -> NeedlepathGuardrail:
+    defaults = dict(
+        api_base=FAKE_API_BASE,
+        api_key=FAKE_API_KEY,
+        guardrail_name="needlepath-selection",
+        default_on=True,
+    )
+    defaults.update(kwargs)
+    return NeedlepathGuardrail(**defaults)
+
+
+def _select_response(
+    rendered_context: object = SELECTED_BLOCK,
+    records_selected: object = 1,
+    tokens_after: object = 120,
+    gate: object = None,
+    status: int = 200,
+    omit_rendered_context: bool = False,
+    outcome: object = "engaged",
+    format_metrics: object = None,
+) -> MagicMock:
+    body = {
+        "request_id": "req-1",
+        "outcome": outcome,
+        "fallback_used": False,
+        "records_selected": records_selected,
+        "records_submitted": 1,
+        "records_available": 1,
+        "tokens_before": 900,
+        "tokens_after": tokens_after,
+        "tokens_saved": 780,
+        "policy_version": "np-2026-10-r8",
+        "gate": {"engaged": True, "reason": "engage:needle"} if gate is None else gate,
+    }
+    if not omit_rendered_context:
+        body["rendered_context"] = rendered_context
+    if format_metrics is not None:
+        body["format_metrics"] = format_metrics
+    mock = MagicMock()
+    mock.status_code = status
+    mock.json.return_value = body
+    mock.text = ""
+    return mock
+
+
+def _apply_inputs(messages: list) -> GenericGuardrailAPIInputs:
+    return GenericGuardrailAPIInputs(structured_messages=copy.deepcopy(messages))
+
+
+@pytest.fixture
+def guardrail() -> NeedlepathGuardrail:
+    return _make_guardrail()
+
+
+async def _run(guardrail: NeedlepathGuardrail, post_mock, messages=None) -> GenericGuardrailAPIInputs:
+    inputs = _apply_inputs(AGENT_MESSAGES if messages is None else messages)
+    with patch.object(guardrail.async_handler, "post", post_mock):
+        return await guardrail.apply_guardrail(
+            inputs=inputs,
+            request_data={"model": "gpt-4o"},
+            input_type="request",
+        )
+
+
+async def _run_expect_unchanged(guardrail: NeedlepathGuardrail, post_mock, messages=None) -> None:
+    """A decline hands back the exact inputs object, with every message byte-identical."""
+    msgs = AGENT_MESSAGES if messages is None else messages
+    inputs = _apply_inputs(msgs)
+    with patch.object(guardrail.async_handler, "post", post_mock):
+        result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request")
+    assert result is inputs
+    assert result["structured_messages"] == msgs
+
+
+def _assert_byte_identical(result: GenericGuardrailAPIInputs, messages=None) -> None:
+    """Every message came back exactly as it went in."""
+    assert result["structured_messages"] == (AGENT_MESSAGES if messages is None else messages)
+
+
+# ── init ──────────────────────────────────────────────────────────────
+
+
+def test_init_raises_without_api_key(monkeypatch):
+    monkeypatch.delenv("NEEDLEPATH_API_KEY", raising=False)
+    with pytest.raises(ValueError, match="API key"):
+        NeedlepathGuardrail(guardrail_name="needlepath-selection")
+
+
+def test_init_defaults():
+    g = _make_guardrail()
+    assert g.needlepath_api_base == FAKE_API_BASE
+    assert g.select_tool_outputs is True
+    assert g.select_history is False
+    assert g.select_system is False
+    assert g.min_chars_to_select == DEFAULT_MIN_CHARS_TO_SELECT == 500
+    assert g.max_context_tokens == DEFAULT_MAX_CONTEXT_TOKENS == 4000
+    # Pinned, not inherited from the service default.
+    assert g.operating_point == DEFAULT_OPERATING_POINT == "np-2026-10-r8"
+
+
+def test_initialize_guardrail_reads_nested_then_flat_knobs():
+    from litellm.proxy.guardrails.guardrail_hooks.needlepath import (
+        guardrail_class_registry,
+        guardrail_initializer_registry,
+        initialize_guardrail,
+    )
+    from litellm.types.guardrails import LitellmParams, SupportedGuardrailIntegrations
+
+    litellm_params = LitellmParams(
+        guardrail="needlepath",
+        mode="pre_call",
+        api_key=FAKE_API_KEY,
+        api_base=FAKE_API_BASE,
+        optional_params={
+            "select_history": True,
+            "min_chars_to_select": 1200,
+            "operating_point": "np-2026-10-r6",
+        },
+        max_context_tokens=2000,
+    )
+    callback = initialize_guardrail(litellm_params, {"guardrail_name": "needlepath-selection"})
+    assert isinstance(callback, NeedlepathGuardrail)
+    assert callback.select_history is True
+    assert callback.select_tool_outputs is True
+    assert callback.min_chars_to_select == 1200
+    assert callback.max_context_tokens == 2000
+    assert callback.operating_point == "np-2026-10-r6"
+    assert guardrail_initializer_registry[SupportedGuardrailIntegrations.NEEDLEPATH.value] is initialize_guardrail
+    assert guardrail_class_registry[SupportedGuardrailIntegrations.NEEDLEPATH.value] is NeedlepathGuardrail
+
+
+def test_initialize_guardrail_defaults_to_r8():
+    from litellm.proxy.guardrails.guardrail_hooks.needlepath import initialize_guardrail
+    from litellm.types.guardrails import LitellmParams
+
+    litellm_params = LitellmParams(guardrail="needlepath", mode="pre_call", api_key=FAKE_API_KEY)
+    callback = initialize_guardrail(litellm_params, {"guardrail_name": "needlepath-selection"})
+    assert callback.operating_point == "np-2026-10-r8"
+
+
+def test_init_rejects_non_http_api_base():
+    with pytest.raises(ValueError, match="http or https"):
+        _make_guardrail(api_base="file:///etc/passwd")
+
+
+def test_init_rejects_cloud_metadata_api_base():
+    with pytest.raises(ValueError, match="cloud-metadata"):
+        _make_guardrail(api_base="http://169.254.169.254")
+
+
+# ── selection core ────────────────────────────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_replaces_tool_output_with_rendered_context(
+    guardrail: NeedlepathGuardrail,
+):
+    post = AsyncMock(return_value=_select_response())
+    result = await _run(guardrail, post)
+
+    _, call_kwargs = post.call_args
+    assert call_kwargs["url"] == f"{FAKE_API_BASE}/v1/context/select"
+    assert call_kwargs["headers"]["Authorization"] == f"Bearer {FAKE_API_KEY}"
+    payload = call_kwargs["json"]
+    assert len(payload["records"]) == 1
+    record = payload["records"][0]
+    assert record["text"] == TOOL_OUTPUT
+    assert record["kind"] == "tool_result"
+    # Title from the tool name, source from the tool_call_id.
+    assert record["title"] == "edgar_search"
+    assert record["source"] == "call_1"
+    # The query is the tool call's intent, not the user question.
+    assert payload["task"]["prompt"] == 'edgar_search: {"cik": "0000320193"}'
+    assert payload["budget"] == {
+        "max_context_tokens": 4000,
+        "operating_point": "np-2026-10-r8",
+    }
+    # The label travels in budget.operating_point; a top-level one is ignored.
+    assert "operating_point" not in payload
+    assert payload["render"] is True
+    assert payload["render_format"] == "plain"
+
+    out = result["structured_messages"]
+    assert out[TOOL_INDEX]["content"] == SELECTED_BLOCK
+    # Only that message changed; every other one is byte-identical.
+    assert out[0] == AGENT_MESSAGES[0]
+    assert out[1] == AGENT_MESSAGES[1]
+    assert out[2] == AGENT_MESSAGES[2]
+    # The rewritten message keeps its envelope.
+    assert out[TOOL_INDEX]["role"] == "tool"
+    assert out[TOOL_INDEX]["tool_call_id"] == "call_1"
+
+
+@pytest.mark.asyncio
+async def test_tool_output_without_matching_call_uses_user_question(
+    guardrail: NeedlepathGuardrail,
+):
+    messages = [
+        {"role": "user", "content": USER_QUESTION},
+        {"role": "tool", "tool_call_id": "call_missing", "content": TOOL_OUTPUT},
+    ]
+    post = AsyncMock(return_value=_select_response())
+    await _run(guardrail, post, messages)
+
+    _, call_kwargs = post.call_args
+    assert call_kwargs["json"]["task"]["prompt"] == USER_QUESTION
+
+
+@pytest.mark.asyncio
+async def test_system_and_history_not_selected_by_default(
+    guardrail: NeedlepathGuardrail,
+):
+    post = AsyncMock(return_value=_select_response())
+    await _run(guardrail, post)
+    # Only the tool output was submitted, even though the system message is
+    # comfortably over the character threshold.
+    assert post.call_count == 1
+    assert post.call_args[1]["json"]["records"][0]["text"] == TOOL_OUTPUT
+
+
+@pytest.mark.asyncio
+async def test_opt_in_system_selects_system_message():
+    guardrail = _make_guardrail(select_system=True)
+    post = AsyncMock(return_value=_select_response())
+    await _run(guardrail, post)
+
+    assert post.call_count == 2
+    submitted = {call[1]["json"]["records"][0]["kind"] for call in post.call_args_list}
+    assert submitted == {"tool_result", "external_data"}
+
+
+@pytest.mark.asyncio
+async def test_short_messages_skipped(guardrail: NeedlepathGuardrail):
+    messages = [
+        {"role": "user", "content": USER_QUESTION},
+        {"role": "tool", "tool_call_id": "call_1", "content": "short"},
+    ]
+    post = AsyncMock(return_value=_select_response())
+    result = await _run(guardrail, post, messages)
+
+    post.assert_not_called()
+    assert result["structured_messages"] == messages
+
+
+@pytest.mark.asyncio
+async def test_response_input_type_passthrough(guardrail: NeedlepathGuardrail):
+    inputs = _apply_inputs(AGENT_MESSAGES)
+    result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response")
+    assert result is inputs
+
+
+# ── fail-open contract ────────────────────────────────────────────────
+#
+# One test per decline. Each asserts the exact inputs object comes back
+# (handlers detect a guardrail edit by identity) and that the messages are
+# byte-identical to what arrived.
+
+
+@pytest.mark.asyncio
+async def test_fail_open_zero_records_selected(guardrail: NeedlepathGuardrail):
+    post = AsyncMock(return_value=_select_response(records_selected=0))
+    inputs = _apply_inputs(AGENT_MESSAGES)
+    with patch.object(guardrail.async_handler, "post", post):
+        result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request")
+    assert result is inputs
+    _assert_byte_identical(result)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_zero_tokens_after(guardrail: NeedlepathGuardrail):
+    post = AsyncMock(return_value=_select_response(tokens_after=0))
+    inputs = _apply_inputs(AGENT_MESSAGES)
+    with patch.object(guardrail.async_handler, "post", post):
+        result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request")
+    assert result is inputs
+    _assert_byte_identical(result)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_blank_rendered_context(guardrail: NeedlepathGuardrail):
+    post = AsyncMock(return_value=_select_response(rendered_context="   "))
+    await _run_expect_unchanged(guardrail, post)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_missing_rendered_context(guardrail: NeedlepathGuardrail):
+    post = AsyncMock(return_value=_select_response(omit_rendered_context=True))
+    await _run_expect_unchanged(guardrail, post)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "reason",
+    [
+        "standdown:flat_gap",
+        "standdown:high_drift",
+        "standdown:insufficient_candidates",
+        "standdown:answerability_low",
+        "standdown:gate_error:ValueError",
+    ],
+)
+async def test_fail_open_gate_standdown(guardrail: NeedlepathGuardrail, reason: str):
+    """A stand-down is the service saying the full content is what to send.
+
+    The reason string is an open enum, so this matches on the ``standdown:``
+    prefix rather than on the individual values.
+    """
+    post = AsyncMock(return_value=_select_response(gate={"engaged": False, "reason": reason}))
+    inputs = _apply_inputs(AGENT_MESSAGES)
+    with patch.object(guardrail.async_handler, "post", post):
+        result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request")
+    assert result is inputs
+    _assert_byte_identical(result)
+    # The service was called and answered 200; the decline is ours.
+    assert post.call_count == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("status_code", [400, 402, 403, 429, 500, 503])
+async def test_fail_open_http_error_status(guardrail: NeedlepathGuardrail, status_code: int):
+    """The shared handler raises for status, so every non-2xx arrives as an
+    HTTPStatusError carrying the upstream body. None of it reaches the client."""
+    request = httpx.Request("POST", f"{FAKE_API_BASE}/v1/context/select")
+    response = httpx.Response(status_code, request=request, text="upstream detail")
+    post = AsyncMock(side_effect=httpx.HTTPStatusError("boom", request=request, response=response))
+    await _run_expect_unchanged(guardrail, post)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_non_2xx_without_raise_for_status(guardrail: NeedlepathGuardrail):
+    """A handler configured not to raise still yields a non-2xx response object."""
+    post = AsyncMock(return_value=_select_response(status=429))
+    await _run_expect_unchanged(guardrail, post)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_timeout(guardrail: NeedlepathGuardrail):
+    post = AsyncMock(side_effect=httpx.ReadTimeout("timed out"))
+    await _run_expect_unchanged(guardrail, post)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_transport_error(guardrail: NeedlepathGuardrail):
+    post = AsyncMock(side_effect=httpx.ConnectError("connection refused"))
+    await _run_expect_unchanged(guardrail, post)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_unexpected_client_exception(guardrail: NeedlepathGuardrail):
+    """An exception outside httpx's hierarchy is still a decline, not a 500."""
+    await _run_expect_unchanged(guardrail, AsyncMock(side_effect=RuntimeError("client closed")))
+
+
+@pytest.mark.asyncio
+async def test_cancellation_is_not_swallowed(guardrail: NeedlepathGuardrail):
+    with pytest.raises(asyncio.CancelledError):
+        await _run(guardrail, AsyncMock(side_effect=asyncio.CancelledError()))
+
+
+@pytest.mark.asyncio
+async def test_fail_open_malformed_json(guardrail: NeedlepathGuardrail):
+    mock = MagicMock()
+    mock.status_code = 200
+    mock.json.side_effect = ValueError("not json")
+    mock.text = "<html>gateway error</html>"
+    await _run_expect_unchanged(guardrail, AsyncMock(return_value=mock))
+
+
+@pytest.mark.asyncio
+async def test_fail_open_unexpected_schema(guardrail: NeedlepathGuardrail):
+    """A 200 whose body is well-formed JSON but not the documented object."""
+    mock = MagicMock()
+    mock.status_code = 200
+    mock.json.return_value = ["not", "an", "object"]
+    mock.text = ""
+    await _run_expect_unchanged(guardrail, AsyncMock(return_value=mock))
+
+
+@pytest.mark.asyncio
+async def test_fail_open_selection_not_smaller(guardrail: NeedlepathGuardrail):
+    """A block that is not shorter than the original could only add tokens."""
+    post = AsyncMock(return_value=_select_response(rendered_context=TOOL_OUTPUT + " and more"))
+    await _run_expect_unchanged(guardrail, post)
+
+
+# ── security screening outcomes ───────────────────────────────────────
+# Shapes follow the documented r8 responses: a sanitized_fallback carries no
+# selected records, a stand-down gate reason, and the full context with each
+# flagged span replaced by a marker.
+
+QUARANTINE_MARKER = "[QUARANTINED: instruction-like candidate text removed]"
+SANITIZED_PACKET = TOOL_OUTPUT[:-60] + QUARANTINE_MARKER
+QUARANTINE_METRICS = {
+    "security_mode": "enforce",
+    "security_action": "quarantine",
+    "security_findings_count": 1,
+}
+
+
+def _sanitized_fallback_response(rendered_context: object = SANITIZED_PACKET) -> MagicMock:
+    return _select_response(
+        rendered_context=rendered_context,
+        records_selected=0,
+        tokens_after=210,
+        gate={"engaged": False, "reason": "standdown:source_support_missing"},
+        outcome="sanitized_fallback",
+        format_metrics={**QUARANTINE_METRICS, "security_delivery_outcome": "sanitized_fallback"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_sanitized_fallback_is_applied(guardrail: NeedlepathGuardrail):
+    """Declining would forward the original, flagged span included."""
+    post = AsyncMock(return_value=_sanitized_fallback_response())
+    result = await _run(guardrail, post)
+    out = result["structured_messages"]
+    assert out[TOOL_INDEX]["content"] == SANITIZED_PACKET
+    assert out[TOOL_INDEX]["tool_call_id"] == "call_1"
+    assert out[:TOOL_INDEX] == AGENT_MESSAGES[:TOOL_INDEX]
+
+
+@pytest.mark.asyncio
+async def test_sanitized_fallback_is_applied_when_not_shorter(guardrail: NeedlepathGuardrail):
+    """The marker can make the packet longer than the span it replaced."""
+    longer = TOOL_OUTPUT + QUARANTINE_MARKER
+    post = AsyncMock(return_value=_sanitized_fallback_response(rendered_context=longer))
+    result = await _run(guardrail, post)
+    assert result["structured_messages"][TOOL_INDEX]["content"] == longer
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("rendered", ["", "   ", None])
+async def test_sanitized_fallback_without_a_packet_fails_open(guardrail: NeedlepathGuardrail, rendered):
+    post = AsyncMock(return_value=_sanitized_fallback_response(rendered_context=rendered))
+    inputs = _apply_inputs(AGENT_MESSAGES)
+    with patch.object(guardrail.async_handler, "post", post):
+        result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request")
+    assert result is inputs
+    _assert_byte_identical(result)
+
+
+@pytest.mark.asyncio
+async def test_engaged_quarantined_block_is_applied_despite_standdown_and_size(
+    guardrail: NeedlepathGuardrail,
+):
+    """An engaged selection with a span quarantined is the safe delivery."""
+    longer = TOOL_OUTPUT + QUARANTINE_MARKER
+    post = AsyncMock(
+        return_value=_select_response(
+            rendered_context=longer,
+            gate={"engaged": False, "reason": "standdown:source_support_missing"},
+            format_metrics=QUARANTINE_METRICS,
+        )
+    )
+    result = await _run(guardrail, post)
+    assert result["structured_messages"][TOOL_INDEX]["content"] == longer
+
+
+@pytest.mark.asyncio
+async def test_screening_keep_does_not_bypass_declines(guardrail: NeedlepathGuardrail):
+    """Screening that flagged nothing leaves the ordinary declines in force."""
+    post = AsyncMock(
+        return_value=_select_response(
+            gate={"engaged": False, "reason": "standdown:source_support_missing"},
+            format_metrics={"security_mode": "enforce", "security_action": "keep", "security_findings_count": 0},
+        )
+    )
+    inputs = _apply_inputs(AGENT_MESSAGES)
+    with patch.object(guardrail.async_handler, "post", post):
+        result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request")
+    assert result is inputs
+    _assert_byte_identical(result)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_blocked(guardrail: NeedlepathGuardrail):
+    """blocked delivers none of the caller's content: a decline, as in the
+    published clients. The placeholder line is never written into a message."""
+    post = AsyncMock(
+        return_value=_select_response(
+            rendered_context="Relevant State: none selected.",
+            records_selected=0,
+            tokens_after=6,
+            outcome="blocked",
+            format_metrics={"security_mode": "enforce", "security_action": "block", "security_findings_count": 3},
+        )
+    )
+    inputs = _apply_inputs(AGENT_MESSAGES)
+    with patch.object(guardrail.async_handler, "post", post):
+        result = await guardrail.apply_guardrail(inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request")
+    assert result is inputs
+    _assert_byte_identical(result)
+
+
+@pytest.mark.asyncio
+async def test_fail_open_blocked_even_with_selected_records(guardrail: NeedlepathGuardrail):
+    """The outcome decides, not the counters next to it."""
+    post = AsyncMock(return_value=_select_response(outcome="blocked"))
+    await _run_expect_unchanged(guardrail, post)
+
+
+@pytest.mark.asyncio
+async def test_one_message_declining_does_not_block_another():
+    """Selection is per message: a decline on one leaves the other applied."""
+    guardrail = _make_guardrail(select_system=True)
+
+    def _by_kind(**kwargs):
+        if kwargs["json"]["records"][0]["kind"] == "tool_result":
+            return _select_response()
+        return _select_response(gate={"engaged": False, "reason": "standdown:flat_gap"})
+
+    post = AsyncMock(side_effect=_by_kind)
+    result = await _run(guardrail, post)
+
+    out = result["structured_messages"]
+    assert out[TOOL_INDEX]["content"] == SELECTED_BLOCK
+    # The system message stood down, so it is byte-identical.
+    assert out[0] == AGENT_MESSAGES[0]
+
+
+# ── fan-out bounds ────────────────────────────────────────────────────
+
+
+def _many_tool_messages(count: int) -> list:
+    """A user question plus `count` eligible tool outputs of strictly increasing size."""
+    messages = [{"role": "user", "content": USER_QUESTION}]
+    for i in range(count):
+        messages.append(
+            {
+                "role": "tool",
+                "tool_call_id": f"call_{i}",
+                # All above the min-chars threshold; index i is the (i+1)-th smallest.
+                "content": "a" * (DEFAULT_MIN_CHARS_TO_SELECT + 20 + i),
+            }
+        )
+    return messages
+
+
+@pytest.mark.asyncio
+async def test_target_cap_selects_only_the_largest_messages(guardrail: NeedlepathGuardrail):
+    """Past the per-request cap, only the largest messages are selected.
+
+    The smallest overflow messages come back byte-identical and no service
+    call is made for them.
+    """
+    overflow = 4
+    messages = _many_tool_messages(_MAX_TARGETS_PER_REQUEST + overflow)
+
+    post = AsyncMock(return_value=_select_response())
+    result = await _run(guardrail, post, messages=messages)
+
+    assert post.await_count == _MAX_TARGETS_PER_REQUEST
+    out = result["structured_messages"]
+    # Tool messages start at index 1 and grow with index: the first `overflow`
+    # are the smallest, so they are the ones left untouched.
+    for idx in range(1, 1 + overflow):
+        assert out[idx] == messages[idx]
+    for idx in range(1 + overflow, len(messages)):
+        assert out[idx]["content"] == SELECTED_BLOCK
+
+
+@pytest.mark.asyncio
+async def test_concurrent_selections_are_bounded(guardrail: NeedlepathGuardrail):
+    """No more than _MAX_CONCURRENT_SELECTIONS service calls are in flight at once."""
+    in_flight = 0
+    high_water = 0
+
+    async def _tracking_post(*args, **kwargs):
+        nonlocal in_flight, high_water
+        in_flight += 1
+        high_water = max(high_water, in_flight)
+        # Yield twice so every scheduled call gets a chance to start before
+        # this one finishes; without the semaphore all of them would overlap.
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        in_flight -= 1
+        return _select_response()
+
+    messages = _many_tool_messages(_MAX_TARGETS_PER_REQUEST)
+    result = await _run(guardrail, _tracking_post, messages=messages)
+
+    assert high_water <= _MAX_CONCURRENT_SELECTIONS
+    assert all(m["content"] == SELECTED_BLOCK for m in result["structured_messages"][1:])
+
+
+# ── /v1/responses ─────────────────────────────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_responses_api_function_call_output_is_selected_and_written_back(
+    guardrail: NeedlepathGuardrail,
+):
+    """Through LiteLLM's Responses translation, a function_call_output is
+    selected against its call's intent and the block is written back."""
+    from litellm.llms.openai.responses.guardrail_translation.handler import (
+        OpenAIResponsesHandler,
+    )
+
+    data = {
+        "model": "gpt-4o",
+        "input": [
+            {"role": "user", "content": USER_QUESTION},
+            {
+                "type": "function_call",
+                "call_id": "call_1",
+                "name": "edgar_search",
+                "arguments": '{"cik": "0000320193"}',
+            },
+            {"type": "function_call_output", "call_id": "call_1", "output": TOOL_OUTPUT},
+        ],
+    }
+    post = AsyncMock(return_value=_select_response())
+    with patch.object(guardrail.async_handler, "post", post):
+        out = await OpenAIResponsesHandler().process_input_messages(data, guardrail)
+
+    assert post.call_count == 1
+    assert post.call_args.kwargs["json"]["task"]["prompt"] == 'edgar_search: {"cik": "0000320193"}'
+    assert out["input"][2]["output"] == SELECTED_BLOCK
+    assert out["input"][0] == {"role": "user", "content": USER_QUESTION}