From b7af384dce43225d61864df97d02af38f7bafc98 Mon Sep 17 00:00:00 2001 From: "joseph.tu" Date: Thu, 13 Aug 2026 18:36:48 +0800 Subject: [PATCH] feat(guardrails/xecguard): pass the calling virtual key through to the SIEM XecGuard exports a SIEM event per scan. That event carried nothing about which virtual key caused the scan, so attributing one meant a lookup back into the proxy database. `send_meta`, off by default, forwards the calling key as the scan payload's `meta` object: its identity, the attributes the proxy already injects about it (tenancy, spend, request route), and the key's own metadata. Values are coerced to the flat scalars the backend accepts, and oversize or nested ones are skipped rather than failing the scan. Proxy-injected attributes are emitted first and win a name collision, so a key's own metadata cannot shadow them. `meta_identity_format` selects the wire shape of the identity: a bare alias, or `{alias, key_id}` for cases where the alias is absent, renamed or reused. The alias form is the default and stays supported, since a backend that only accepts strings answers the object form with 400. Callback credential slots are never forwarded and cannot be opted back in. The proxy's own per-key control settings (rate limits, budget knobs) are skipped by default and can be opted back in by name; that set is derived from the proxy's own metadata-field lists rather than copied. Also included: `apply_to_aliases` / `except_aliases`, so one guardrail can cover a chosen set of virtual keys instead of every request. All changes are in XecGuard-owned files. The UI form is generated from the config model, so the shared frontend file is untouched. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/xecguard/__init__.py | 5 + .../guardrail_hooks/xecguard/xecguard.py | 419 +++++- .../guardrails/guardrail_hooks/xecguard.py | 78 +- .../guardrail_hooks/test_xecguard.py | 1226 ++++++++++++++++- 4 files changed, 1725 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/__init__.py index a3825cca7bc..917f36f9888 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/__init__.py @@ -19,6 +19,11 @@ def initialize_guardrail( api_key=litellm_params.api_key, xecguard_model=litellm_params.xecguard_model, policy_names=litellm_params.policy_names, + apply_to_aliases=litellm_params.apply_to_aliases, + except_aliases=litellm_params.except_aliases, + send_meta=litellm_params.send_meta, + meta_data_fields=litellm_params.meta_data_fields, + meta_identity_format=litellm_params.meta_identity_format, block_on_error=litellm_params.block_on_error, grounding_strictness=litellm_params.grounding_strictness, guardrail_name=guardrail.get( diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index ddb40dc3ca0..440389a79ae 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -20,10 +20,18 @@ Design notes (intentional divergences from the framework defaults): directly for ``logging_only`` mode - it does NOT bridge to ``apply_guardrail``. Our override runs the scan non-blockingly and swallows every exception. + * When ``send_meta`` is enabled the scan payload carries a ``meta`` + object identifying the calling virtual key. It is correlation data + for XecGuard's SIEM export only and never affects the verdict; the + backend's flat-scalar contract for it is enforced client-side so a + malformed key metadata entry cannot fail an otherwise valid scan. """ import asyncio +import json import os +import re +from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional @@ -40,6 +48,10 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy._types import ( + LiteLLM_ManagementEndpoint_MetadataFields, + LiteLLM_ManagementEndpoint_MetadataFields_Premium, +) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( GenericGuardrailAPIInputs, @@ -76,6 +88,70 @@ _DEFAULT_POLICIES: Final = [ "Default_Policy_GeneralPromptAttackProtection", ] +# ``meta`` contract of POST /xecguard/v1/scan: an optional object carrying caller +# context that takes no part in detection. XecGuard flattens it into the SIEM +# event (``virtualkey`` -> ``ctx_virtualkey``, ``data.X`` -> ``ctx_X``), and SIEM +# index fields only accept flat scalars - anything else is rejected with 400. So +# every value is coerced or dropped here rather than risking a scan failure. +_METADATA_KEY_METADATA_FIELD: Final = "user_api_key_metadata" +_META_NAME_PATTERN: Final = re.compile(r"^[A-Za-z_][A-Za-z0-9_.\-]{0,63}$") +_META_CONTROL_CHARS: Final = re.compile(r"[\x00-\x1f\x7f]") +_META_MAX_DATA_FIELDS: Final = 32 +_META_MAX_VALUE_CHARS: Final = 512 +_META_MAX_SERIALIZED_BYTES: Final = 4096 +# Never forwarded, and ``meta_data_fields`` cannot opt them back in: these slots +# hold credentials, so there is no configuration under which shipping them to an +# external SIEM is right. +_META_EXCLUDED_DATA_FIELDS: Final = frozenset({"logging", "callback_settings", "secret_manager_settings"}) + +# The proxy stores its own per-key control settings inside key metadata - rate +# limits, budget knobs, enforced params, ``disable_global_guardrails``. They sit +# in the same dict as the admin's own fields but they are proxy configuration, +# not caller identity: noise in a SIEM, they eat the 32-field / 4096-byte budget, +# and a couple of them describe the key's security posture. Skipped by default, +# but an admin who explicitly names one in ``meta_data_fields`` gets it - unlike +# the credential slots above, forwarding these is a judgement call, not a bug. +# +# Taken from the proxy's own lists rather than copied, so a field litellm adds +# later is covered without an edit here. +_META_CONTROL_DATA_FIELDS: Final = ( + frozenset(LiteLLM_ManagementEndpoint_MetadataFields) | frozenset(LiteLLM_ManagementEndpoint_MetadataFields_Premium) +) - _META_EXCLUDED_DATA_FIELDS + +# Two shapes for ``meta.virtualkey``. "string" is the identity as a bare string, +# which is all the currently deployed backend accepts. "object" carries the alias +# and the key id side by side, so a SIEM event is attributable even when the alias +# is absent, renamed, or reused - it needs a backend that validates the object +# form, hence the switch rather than a straight cutover. +_META_IDENTITY_FORMATS: Final = ("string", "object") +_DEFAULT_META_IDENTITY_FORMAT: Final = "string" + +# Virtual-key attributes the proxy injects alongside every request, forwarded as +# ``meta.data`` so a SIEM event can be attributed without a lookup back into the +# proxy database. Ordered: identity first, then tenancy, then commercials, so the +# fields that survive the 32-field / 4096-byte caps are the ones worth keeping. +# +# This set deliberately includes PII (``user_email``) and commercial figures +# (``spend``, ``max_budget``). Both leave the proxy only when ``send_meta`` is +# explicitly enabled, and ``meta_data_fields`` narrows the set for deployments +# that must not egress them. +_META_AUTO_DATA_FIELDS: Final[tuple[tuple[str, str], ...]] = ( + ("key_id", "user_api_key_hash"), + ("key_alias", "user_api_key_alias"), + ("team_id", "user_api_key_team_id"), + ("team_alias", "user_api_key_team_alias"), + ("user_id", "user_api_key_user_id"), + ("user_email", "user_api_key_user_email"), + ("org_id", "user_api_key_org_id"), + ("org_alias", "user_api_key_org_alias"), + ("project_id", "user_api_key_project_id"), + ("project_alias", "user_api_key_project_alias"), + ("end_user_id", "user_api_key_end_user_id"), + ("spend", "user_api_key_spend"), + ("max_budget", "user_api_key_max_budget"), + ("request_route", "user_api_key_request_route"), +) + class XecGuardMissingCredentials(Exception): pass @@ -88,6 +164,11 @@ class XecGuardGuardrail(CustomGuardrail): api_base: str | None = None, xecguard_model: str | None = None, policy_names: list[str] | None = None, + apply_to_aliases: Sequence[str] | None = None, + except_aliases: Sequence[str] | None = None, + send_meta: bool | None = None, + meta_data_fields: Sequence[str] | None = None, + meta_identity_format: str | None = None, block_on_error: bool | None = None, grounding_strictness: str | None = None, **kwargs: Any, @@ -105,6 +186,40 @@ class XecGuardGuardrail(CustomGuardrail): self.xecguard_model = xecguard_model or _DEFAULT_MODEL self.policy_names = policy_names + # Guardrail-side key targeting (free, OSS). Normalized to lists. + self.apply_to_aliases = apply_to_aliases or () + self.except_aliases = except_aliases or () + + # Caller context forwarded as the scan payload's ``meta``. Opt-in: turning + # it on sends the calling key's alias and its admin-set metadata to + # XecGuard, which is a data-egress change no upgrade should make silently. + if send_meta is None: + self.send_meta = os.environ.get("XECGUARD_SEND_META", "false").lower() in ( + "true", + "1", + "yes", + ) + else: + self.send_meta = send_meta + self.meta_data_fields = tuple(meta_data_fields) if meta_data_fields else () + + # Wire shape of ``meta.virtualkey``. Defaults to the string form: a backend + # that only accepts strings answers the object form with 400, and with + # ``block_on_error`` on (the default) that turns every request into a block. + # An unknown value falls back rather than raising - a typo in the UI should + # not take the gateway down. + requested_format = ( + (meta_identity_format or os.environ.get("XECGUARD_META_IDENTITY_FORMAT") or "").strip().lower() + ) + if requested_format and requested_format not in _META_IDENTITY_FORMATS: + verbose_proxy_logger.warning( + "XecGuard: unknown meta_identity_format %r - falling back to %r (valid: %s)", + requested_format, + _DEFAULT_META_IDENTITY_FORMAT, + ", ".join(_META_IDENTITY_FORMATS), + ) + requested_format = "" + self.meta_identity_format = requested_format or _DEFAULT_META_IDENTITY_FORMAT if block_on_error is None: env: Final = os.environ.get("XECGUARD_BLOCK_ON_ERROR", "true") @@ -143,6 +258,92 @@ class XecGuardGuardrail(CustomGuardrail): GuardrailEventHooks.logging_only, ] + @staticmethod + def _calling_key_identity( + request_data: Mapping[str, Any] | None, + ) -> tuple[str | None, str | None]: + """Return (key_alias, key_hash) of the calling virtual key from the + proxy-injected request metadata. Both may be None (e.g. master key).""" + alias: str | None = None + key_hash: str | None = None + if isinstance(request_data, dict): + for meta_key in ("metadata", "litellm_metadata"): + md = request_data.get(meta_key) + if isinstance(md, dict): + alias = alias or md.get("user_api_key_alias") + key_hash = key_hash or md.get("user_api_key_hash") + return alias, key_hash + + def _key_is_targeted(self, request_data: Mapping[str, Any] | None) -> bool: + """Guardrail-side key targeting. With no allow/block list configured, + every key is scanned. Otherwise the calling key is matched by alias + (preferred) or hashed token: + * blocklist (except_aliases): listed keys are NOT scanned; + * allowlist (apply_to_aliases): only listed keys are scanned. + When both are set, a key is scanned iff it is in the allowlist AND not + in the blocklist. + """ + allowlist: Final = self.apply_to_aliases or () + blocklist: Final = self.except_aliases or () + if not allowlist and not blocklist: + return True + + alias, key_hash = self._calling_key_identity(request_data) + identifiers: Final = tuple(ident for ident in (alias, key_hash) if ident) + + # Deny wins, and it is checked first so that precedence stays visible + # rather than folded into the allowlist expression below. + if blocklist and any(ident in blocklist for ident in identifiers): + return False + if not allowlist: + return True + return any(ident in allowlist for ident in identifiers) + + # Metadata fields the proxy injects to identify the calling virtual key. + _KEY_IDENTITY_FIELDS = ("user_api_key_alias", "user_api_key_hash") + + @classmethod + def _key_context(cls, data: Mapping[str, Any] | None) -> Mapping[str, Any] | None: + """Return a mapping ``_calling_key_identity`` can read the key fields from. + + That reader looks for top-level ``metadata`` / ``litellm_metadata``. On the + pre/during/post_call paths the proxy already puts the injected key fields + there, so ``data`` is handed back untouched -- reshaping to a single key would + drop the other location it also reads. Only the logging path needs help: there + ``data`` is ``model_call_details``, which carries the same fields one level + down under ``litellm_params``. + """ + if not isinstance(data, dict): + return data + for meta_key in ("metadata", "litellm_metadata"): + md = data.get(meta_key) + if isinstance(md, dict) and any(field in md for field in cls._KEY_IDENTITY_FIELDS): + return data + nested = data.get("litellm_params") + if isinstance(nested, dict): + for meta_key in ("metadata", "litellm_metadata"): + md = nested.get(meta_key) + if isinstance(md, dict): + return {meta_key: md} # mutable-ok: lifts nested metadata to the readers' shape + return data + + def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool: + """Gate on the calling virtual key in addition to the native checks. + + Deciding here rather than inside ``apply_guardrail`` is what makes LiteLLM + record the guardrail as not having run for a key this guardrail does not + cover, instead of logging a "success"/"allow" entry for a request it never + evaluated. ``super()`` is consulted first so the native decisions -- global + opt-outs, event-hook matching, tag-based modes -- keep precedence. + + The gates are still enforced inside ``apply_guardrail`` and + ``async_logging_hook`` as well: ``POST /guardrails/apply_guardrail`` invokes + ``apply_guardrail`` directly and never reaches this method. + """ + if not super().should_run_guardrail(data, event_type): + return False + return self._key_is_targeted(self._key_context(data)) + @log_guardrail_information async def apply_guardrail( self, @@ -151,6 +352,13 @@ class XecGuardGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: + # Guardrail-side key targeting (allowlist / blocklist by key alias): + # skip scanning entirely for keys this guardrail does not cover. + # should_run_guardrail already gates the proxy's own dispatch paths; this + # also covers POST /guardrails/apply_guardrail, which calls straight in. + if not self._key_is_targeted(self._key_context(request_data)): + return inputs + messages: Final = self._build_full_history( request_data=request_data, inputs=inputs, @@ -160,7 +368,11 @@ class XecGuardGuardrail(CustomGuardrail): return inputs scan_type: Final = "input" if input_type == "request" else "response" - scan_result: Final = await self._call_scan(messages=messages, scan_type=scan_type) + scan_result: Final = await self._call_scan( + messages=messages, + scan_type=scan_type, + request_data=request_data, + ) if scan_result is None: return inputs @@ -214,6 +426,12 @@ class XecGuardGuardrail(CustomGuardrail): ): return kwargs, result + # Same key targeting as apply_guardrail. logging_only reaches the guardrail + # through this hook rather than apply_guardrail, so the gate is repeated here; + # without it an excluded key's content would still be sent to XecGuard. + if not self._key_is_targeted(self._key_context(kwargs)): + return kwargs, result + start_time: Final = datetime.now() try: assistant_text: Final = self._extract_assistant_text_from_response(result) @@ -240,6 +458,7 @@ class XecGuardGuardrail(CustomGuardrail): scan_result: Final = await self._call_scan( messages=messages, scan_type=scan_type, + request_data=request_data, suppress_errors=True, ) if scan_result is None: @@ -300,6 +519,200 @@ class XecGuardGuardrail(CustomGuardrail): ) return kwargs, result + # ------------------------------------------------------------------ + # Caller context (scan payload ``meta``) - SIEM correlation only + # ------------------------------------------------------------------ + + def _build_scan_meta(self, context: Mapping[str, Any] | None) -> Mapping[str, Any] | None: + """Assemble the scan payload's ``meta`` object, or None to omit it. + + ``virtualkey`` is the identity this guardrail filtered on and ``data`` is + the calling key's proxy-injected attributes plus its own metadata as set on + the Virtual Keys page. Neither participates in detection - XecGuard forwards + them to the SIEM so a scan can be traced back to the virtual key that caused + it. + + ``meta`` is optional in the contract, so anything that cannot be made to + satisfy it is left out instead of turning a scan into a 400. + """ + if not self.send_meta: + return None + + virtualkey: str | Mapping[str, str] | None + if self.meta_identity_format == "object": + virtualkey = self._scan_meta_virtualkey_object(context) + else: + virtualkey = self._scan_meta_virtualkey(context) + if not virtualkey: + verbose_proxy_logger.debug( + "XecGuard: omitting scan meta - the calling key has no alias or hash matching the " + "backend's virtualkey pattern (give the key a key_alias to enable SIEM correlation)" + ) + return None + + meta: dict[str, Any] = {"virtualkey": virtualkey} # mutable-ok: the JSON object being assembled + data = self._build_scan_meta_data(context, virtualkey=virtualkey) + if data: + meta["data"] = data + return meta + + def _scan_meta_virtualkey(self, context: Mapping[str, Any] | None) -> str | None: + """The key identity to report, or None when there is no usable one. + + Alias first: that is what an operator types into ``apply_to_aliases`` / + ``except_aliases``, so the value in the SIEM matches the value in the + guardrail config. The hashed token is the fallback for keys created + without an alias; a master-key call has neither. Note that the token + hash only satisfies the backend's pattern when it happens to start with + a hex letter - aliasless keys are not reliably correlatable. + """ + for candidate in self._calling_key_identity(context): + if isinstance(candidate, str) and _META_NAME_PATTERN.match(candidate): + return candidate + return None + + def _scan_meta_virtualkey_object(self, context: Mapping[str, Any] | None) -> Mapping[str, str] | None: + """The object form of ``virtualkey``: ``{"alias": ..., "key_id": ...}``. + + Either member may be absent - a key created without an alias has only an + id, and a master-key call has neither (in which case ``meta`` is omitted). + Unlike the string form this does not require the alias to satisfy the + backend's identifier pattern: the pattern exists because a bare string + becomes a SIEM field *value* directly, whereas here each member is + sanitized the same way ``meta.data`` values are. That makes keys whose + alias contains spaces or CJK correlatable, which the string form cannot do. + """ + alias, key_hash = self._calling_key_identity(context) + obj: dict[str, str] = {} # mutable-ok: the JSON object being assembled + for name, raw in (("alias", alias), ("key_id", key_hash)): + value = self._coerce_meta_value(raw) + if value is not None: + obj[name] = value + return obj or None + + @staticmethod + def _calling_key_metadata(context: Mapping[str, Any] | None) -> Mapping[object, Any]: + """The calling virtual key's own metadata, as injected by the proxy. + + This is the JSON an admin typed into the key's Metadata box on the + Virtual Keys page (minus the callback-credential slots, which the proxy + strips before injecting). Team metadata is deliberately not merged in: + ``meta.data`` is meant to describe the key that made the call. + + The key type is ``object``, not ``str``: nothing between the database and + here validates it, and the caller drops a non-str key rather than letting + it reach ``re.match`` and raise. Narrowing this to ``str`` would make that + guard look redundant to a type checker and invite its removal. + """ + if not isinstance(context, dict): + return {} # mutable-ok: "this key has no metadata"; the caller only reads it + for meta_key in ("metadata", "litellm_metadata"): + md = context.get(meta_key) + if isinstance(md, dict): + key_metadata = md.get(_METADATA_KEY_METADATA_FIELD) + if isinstance(key_metadata, dict): + return key_metadata + return {} # mutable-ok: same empty result, no metadata field was injected + + @classmethod + def _auto_meta_data_items(cls, context: Mapping[str, Any] | None) -> tuple[tuple[str, Any], ...]: + """The proxy-injected virtual-key attributes, in ``_META_AUTO_DATA_FIELDS`` + order regardless of how the proxy ordered its metadata dict. + + Absent and null fields are skipped, so a key with no team contributes no + ``team_id`` rather than an empty one - a SIEM query for "scans with no + team" then means it, instead of matching every key. + """ + injected: dict[str, Any] = {} # mutable-ok: accumulator keyed by meta.data name + if isinstance(context, dict): + for meta_key in ("metadata", "litellm_metadata"): + md = context.get(meta_key) + if not isinstance(md, dict): + continue + for name, source_field in _META_AUTO_DATA_FIELDS: + if name not in injected and md.get(source_field) is not None: + injected[name] = md[source_field] + return tuple((name, injected[name]) for name, _ in _META_AUTO_DATA_FIELDS if name in injected) + + def _build_scan_meta_data( + self, context: Mapping[str, Any] | None, virtualkey: str | Mapping[str, str] + ) -> Mapping[str, str]: + """Coerce the calling key's attributes and metadata into ``meta.data``. + + Two sources, in this order: the attributes the proxy injects about the + calling key (identity, tenancy, budget), then the free-form metadata an + admin typed into the key's Metadata box. Proxy-injected attributes go + first and win a name collision, so an admin cannot shadow ``key_id`` with + a field of their own and mislead an investigation. + + Fields are kept while they satisfy the contract: a name matching the + backend's pattern, a flat scalar value, at most 32 fields, and a + serialized ``meta`` within the 4096-byte cap. Oversize fields are skipped + rather than ending the scan, so a later small field still gets through. + Dropped names are logged without their values - both sources can hold + sensitive strings. + """ + source = self._calling_key_metadata(context) + data: dict[str, str] = {} # mutable-ok: accumulator, re-measured as it grows + # Re-measured against the real payload shape each time, so the cap holds + # regardless of how long the virtualkey and the field names are. + probe: dict[str, Any] = {"virtualkey": virtualkey, "data": data} # mutable-ok: views `data` + dropped: list[str] = [] # mutable-ok: skipped field names, for one debug line + + for name, raw_value in (*self._auto_meta_data_items(context), *source.items()): + if self.meta_data_fields: + if name not in self.meta_data_fields: + continue + elif name in _META_CONTROL_DATA_FIELDS: + # proxy config rather than caller identity - opt in by name + continue + if name in _META_EXCLUDED_DATA_FIELDS: + continue + if name in data: # a proxy-injected attribute already claimed this name + dropped.append(str(name)) + continue + if not isinstance(name, str) or not _META_NAME_PATTERN.match(name): + dropped.append(str(name)) + continue + if len(data) >= _META_MAX_DATA_FIELDS: + dropped.append(name) + continue + value = self._coerce_meta_value(raw_value) + if value is None: + dropped.append(name) + continue + data[name] = value + if len(json.dumps(probe, ensure_ascii=False).encode("utf-8")) > _META_MAX_SERIALIZED_BYTES: + del data[name] + dropped.append(name) + + if dropped: + verbose_proxy_logger.debug( + "XecGuard: scan meta.data dropped %d field(s) (names only): %s", + len(dropped), + dropped, + ) + return data + + @staticmethod + def _coerce_meta_value(value: object) -> str | None: + """Coerce one key-metadata value to the contract, or None to drop it. + + Scalars are stringified so an admin writing ``{"tier": 3}`` still gets a + usable ``ctx_tier``. Nested objects and lists have no flat representation + a SIEM index field can hold, so they are dropped. + """ + if isinstance(value, bool): + text = "true" if value else "false" + elif isinstance(value, str): + text = value + elif isinstance(value, (int, float)): + text = str(value) + else: + return None + text = _META_CONTROL_CHARS.sub("", text)[:_META_MAX_VALUE_CHARS] + return text or None + # ------------------------------------------------------------------ # HTTP helpers # ------------------------------------------------------------------ @@ -308,6 +721,7 @@ class XecGuardGuardrail(CustomGuardrail): self, messages: list[dict], scan_type: str, + request_data: Mapping[str, Any] | None = None, suppress_errors: bool = False, ) -> dict | None: payload: Final[dict[str, Any]] = { @@ -316,6 +730,9 @@ class XecGuardGuardrail(CustomGuardrail): "messages": messages, "policy_names": (self.policy_names if self.policy_names else _DEFAULT_POLICIES), } + meta = self._build_scan_meta(self._key_context(request_data)) + if meta is not None: + payload["meta"] = meta return await self._post( path=_SCAN_ENDPOINT, payload=payload, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/xecguard.py b/litellm/types/proxy/guardrails/guardrail_hooks/xecguard.py index 74a95898cb3..ce1aa7d9868 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/xecguard.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/xecguard.py @@ -1,6 +1,6 @@ from typing import Final, Literal -from pydantic import Field +from pydantic import Field, field_validator from .base import GuardrailConfigModel @@ -48,6 +48,69 @@ class XecGuardConfigModel(GuardrailConfigModel): "options": list(XECGUARD_DEFAULT_POLICY_OPTIONS), }, ) + apply_to_aliases: str | list[str] | None = Field( # mutable-ok: list sets UI type; Sequence ambiguous + default=None, + description=( + "Allowlist of virtual-key aliases: only requests from keys whose " + "alias is listed here are scanned by this guardrail. Leave empty to " + "apply to all keys (subject to the exclude list below). Accepts a " + "list or a comma-separated string." + ), + ) + except_aliases: str | list[str] | None = Field( # mutable-ok: list sets UI type; Sequence ambiguous + default=None, + description=( + "Exclude list of virtual-key aliases: requests from keys whose alias " + "is listed here are NOT scanned by this guardrail (exempted), even " + "when the allowlist is empty. Accepts a list or a comma-separated " + "string." + ), + ) + send_meta: bool | None = Field( + default=None, + description=( + "Forward caller context to XecGuard as the scan payload's `meta` " + "object: `meta.virtualkey` is the calling virtual key's alias (its " + "token hash when it has no alias) and `meta.data` is that key's own " + "metadata from the Virtual Keys page. It takes no part in the " + "verdict - XecGuard flattens it into the SIEM event (ctx_virtualkey, " + "ctx_) so scans can be traced back to the key that caused " + "them. Defaults to false; falls back to the XECGUARD_SEND_META env " + "var." + ), + ) + meta_data_fields: str | list[str] | None = Field( # mutable-ok: list sets UI type; Sequence ambiguous + default=None, + description=( + "Restrict which of the virtual key's metadata fields are forwarded " + "in `meta.data`. Leave empty to send every field that fits the " + "backend's contract (flat scalar values, at most 32 fields, 512 " + "characters each), minus the proxy's own per-key control settings " + "(rate limits, budget knobs, enforced params) which are skipped as " + "SIEM noise - naming one here opts it back in. Callback credential " + "slots are never forwarded either way. Accepts a list or a " + "comma-separated string. Only used when `send_meta` is enabled." + ), + ) + # Named `meta_identity_format`, not `meta_virtualkey_format`: the proxy masks + # any litellm_param whose name contains "key" before serving it back, so a + # `virtualkey` in the name means the UI form prefills "ob****ct" and saving + # the form writes that back - the plugin then falls through to the default and + # the admin's choice is lost with no error. See the masking regression test. + meta_identity_format: Literal["string", "object"] | None = Field( + default=None, + description=( + "Wire shape of `meta.virtualkey`. 'string' (default) sends the alias " + "as a bare string and is what current XecGuard backends accept. " + "'object' sends `{alias, key_id}` so a scan stays attributable when " + "the key has no alias or the alias was renamed or reused, and lifts " + "the identifier-pattern restriction on aliases - it requires a " + "backend that validates the object form, otherwise every scan is " + "rejected with 400. Falls back to the " + "XECGUARD_META_IDENTITY_FORMAT env var. Only used when `send_meta` " + "is enabled." + ), + ) block_on_error: bool | None = Field( default=None, description=( @@ -67,6 +130,19 @@ class XecGuardConfigModel(GuardrailConfigModel): ), ) + @field_validator("apply_to_aliases", "except_aliases", "meta_data_fields", mode="before") + @classmethod + def _normalize_alias_list(cls, v: object) -> object: + """Accept either a list or a comma-separated string (the UI submits a + plain text box as a string; YAML users may write a list) and normalize + to a de-whitespaced, empties-dropped list of aliases / field names.""" + if v is None: + return None + items: Final = v.split(",") if isinstance(v, str) else v + if isinstance(items, (list, tuple)): + return [s.strip() for s in items if isinstance(s, str) and s.strip()] # mutable-ok: tests assert this list + return v + @staticmethod def ui_friendly_name() -> str: return "XecGuard" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py index 6e601df897b..e3e4d236826 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py @@ -6,17 +6,24 @@ branch coverage. Network calls are always mocked; the companion live suite lives in ``test_xecguard_live.py``. """ +import inspect +import json import os from unittest.mock import MagicMock, patch import httpx import pytest +from pydantic import ValidationError from fastapi.exceptions import HTTPException +from litellm.proxy.guardrails.guardrail_hooks.xecguard import ( + xecguard as xecguard_module, +) from litellm.proxy.guardrails.guardrail_hooks.xecguard.xecguard import ( XecGuardGuardrail, XecGuardMissingCredentials, ) +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import ( XecGuardConfigModel, ) @@ -1230,7 +1237,10 @@ class TestXecGuardMessageAssembly: def test_synthesize_user_joins_strings(self, xecguard_guardrail): assert xecguard_guardrail._synthesize_user_from_inputs( {"texts": ["a", "b"]} - ) == {"role": "user", "content": "a\nb"} + ) == { + "role": "user", + "content": "a\nb", + } def test_extract_last_text_by_role_not_found(self, xecguard_guardrail): assert ( @@ -1975,3 +1985,1217 @@ class TestXecGuardInitializer: assert isinstance(cb, XecGuardGuardrail) assert cb.api_key == "xgs_init" assert cb.guardrail_name == "xg-test" + + def test_initializer_forwards_every_configurable_field(self): + """The initializer names each param explicitly, so a field added to the + config model but not forwarded here is silently inert: the UI shows the + control, the operator sets it, and nothing happens. Assert on the config + model's own field list so adding a field without wiring it fails here.""" + from litellm.proxy.guardrails.guardrail_hooks.xecguard import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + params = LitellmParams( + guardrail="xecguard", + mode="pre_call", + api_key="xgs_init", + api_base="https://api.test.xecguard.local", + xecguard_model="xecguard_v2", + policy_names=["Default_Policy_SkillsProtection"], + apply_to_aliases="key-allow", + except_aliases="key-deny", + send_meta=True, + meta_data_fields="cost_center, owner", + meta_identity_format="object", + block_on_error=False, + grounding_strictness="STRICT", + default_on=True, + ) + cb = initialize_guardrail( + litellm_params=params, guardrail={"guardrail_name": "xg"} + ) + + assert cb.xecguard_model == "xecguard_v2" + assert cb.policy_names == ["Default_Policy_SkillsProtection"] + # the alias lists keep the normalized list the config model produced + assert cb.apply_to_aliases == ["key-allow"] + assert cb.except_aliases == ["key-deny"] + assert cb.send_meta is True + assert cb.meta_data_fields == ("cost_center", "owner") + assert cb.meta_identity_format == "object" + assert cb.block_on_error is False + assert cb.grounding_strictness == "STRICT" + + # Every field the UI renders for this provider must have landed on the + # instance -- api_key/api_base are asserted by the test above. + source = inspect.getsource(initialize_guardrail) + unwired = [ + name + for name in XecGuardConfigModel.model_fields + if name not in ("optional_params",) + and f"litellm_params.{name}" not in source + ] + assert unwired == [], ( + f"XecGuardConfigModel fields not forwarded by initialize_guardrail: {unwired}" + ) + + +# =========================================================================== +# Extension layered on top of the original integration: per-virtual-key +# filtering (key allow/block lists + per-key policy subset via native metadata). +# +# The fixture below is intentionally NOT autouse so the exhaustive suite +# above keeps its original environment handling; it is opted into per class +# via @pytest.mark.usefixtures. +# =========================================================================== + + +@pytest.fixture +def _clean_env(monkeypatch): + """Keep XecGuard env vars from leaking into these tests.""" + for var in ( + "XECGUARD_API_KEY", + "XECGUARD_API_BASE", + "XECGUARD_BLOCK_ON_ERROR", + "XECGUARD_SEND_META", + "XECGUARD_META_IDENTITY_FORMAT", + ): + monkeypatch.delenv(var, raising=False) + yield + + +def _extension_guardrail(**overrides): + params = dict( + api_base="https://api.test.xecguard.local", + api_key="xgs_test_scan_secret", + guardrail_name="test-xecguard", + event_hook="pre_call", + default_on=True, + ) + params.update(overrides) + return XecGuardGuardrail(**params) + + +# --------------------------------------------------------------------------- +# Per-virtual-key filtering +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("_clean_env") +class TestXecGuardCallingKeyIdentity: + def test_from_metadata(self): + data = {"metadata": {"user_api_key_alias": "svc", "user_api_key_hash": "h1"}} + assert XecGuardGuardrail._calling_key_identity(data) == ("svc", "h1") + + def test_from_litellm_metadata(self): + data = {"litellm_metadata": {"user_api_key_alias": "svc2"}} + assert XecGuardGuardrail._calling_key_identity(data) == ("svc2", None) + + def test_master_key_returns_none_none(self): + assert XecGuardGuardrail._calling_key_identity({}) == (None, None) + assert XecGuardGuardrail._calling_key_identity(None) == (None, None) + + +@pytest.mark.usefixtures("_clean_env") +class TestXecGuardKeyIsTargeted: + def test_no_lists_scans_everything(self): + gr = _extension_guardrail() + assert gr._key_is_targeted({"metadata": {"user_api_key_alias": "any"}}) is True + + def test_allowlist_match(self): + gr = _extension_guardrail(apply_to_aliases=["prod"]) + assert gr._key_is_targeted({"metadata": {"user_api_key_alias": "prod"}}) is True + + def test_allowlist_miss(self): + gr = _extension_guardrail(apply_to_aliases=["prod"]) + assert gr._key_is_targeted({"metadata": {"user_api_key_alias": "dev"}}) is False + + def test_blocklist_excludes(self): + gr = _extension_guardrail(except_aliases=["internal"]) + assert ( + gr._key_is_targeted({"metadata": {"user_api_key_alias": "internal"}}) + is False + ) + + def test_blocklist_wins_over_allowlist(self): + gr = _extension_guardrail(apply_to_aliases=["prod"], except_aliases=["prod"]) + assert ( + gr._key_is_targeted({"metadata": {"user_api_key_alias": "prod"}}) is False + ) + + def test_match_by_hash(self): + gr = _extension_guardrail(apply_to_aliases=["hash-abc"]) + assert ( + gr._key_is_targeted({"metadata": {"user_api_key_hash": "hash-abc"}}) is True + ) + + +# --------------------------------------------------------------------------- +# Integration — apply_guardrail respects the key allow/deny lists +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("_clean_env") +class TestXecGuardApplyGuardrailKeyTargeting: + @pytest.mark.asyncio + async def test_key_not_in_allowlist_skips_scan(self): + gr = _extension_guardrail(apply_to_aliases=["prod"]) + data = { + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"user_api_key_alias": "dev"}, + } + with patch.object(gr.async_handler, "post") as post: + result = await gr.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=data, + input_type="request", + ) + assert result == {"texts": ["hi"]} + post.assert_not_called() + + @pytest.mark.asyncio + async def test_blocklisted_key_skips_scan(self): + gr = _extension_guardrail(except_aliases=["internal"]) + data = { + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"user_api_key_alias": "internal"}, + } + with patch.object(gr.async_handler, "post") as post: + result = await gr.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=data, + input_type="request", + ) + assert result == {"texts": ["hi"]} + post.assert_not_called() + + @pytest.mark.asyncio + async def test_targeted_key_uses_config_policies(self): + gr = _extension_guardrail( + apply_to_aliases=["prod"], policy_names=["Config_Level_Policy"] + ) + data = { + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"user_api_key_alias": "prod"}, + } + resp = _make_response({"decision": "SAFE", "trace_id": "tr"}) + with patch.object(gr.async_handler, "post", return_value=resp) as post: + await gr.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=data, + input_type="request", + ) + post.assert_called_once() + assert post.call_args.kwargs["json"]["policy_names"] == ["Config_Level_Policy"] + + +# --------------------------------------------------------------------------- +# Config-model guards — the alias fields are plain manual text input (no live +# dropdown), so this feature stays entirely off shared UI-renderer / options +# code. These lock in the two facts that guarantee that. +# --------------------------------------------------------------------------- + + +class TestXecGuardAliasConfigModel: + def test_validator_splits_comma_separated_string(self): + # The UI submits a plain text box as a single string. + cfg = XecGuardConfigModel(apply_to_aliases="prod, staging , ,dev") + assert cfg.apply_to_aliases == ["prod", "staging", "dev"] + + def test_validator_passes_list_through_and_cleans(self): + # YAML users may write a list; non-str / empty entries are dropped and + # surviving strings are stripped. + cfg = XecGuardConfigModel(except_aliases=[" a ", "b", "", 5, None]) + assert cfg.except_aliases == ["a", "b"] + + def test_validator_none_stays_none(self): + cfg = XecGuardConfigModel() + assert cfg.apply_to_aliases is None + assert cfg.except_aliases is None + + def test_ui_type_is_plain_string_not_array(self): + # Optional[Union[str, List[str]]] must resolve to the "string" UI type + # (str is first in the Union), so the generic renderer draws a plain + # with no options list — never the shared multiselect path. + from litellm.proxy.guardrails.guardrail_endpoints import ( + _get_field_type_from_annotation, + ) + + for field_name in ("apply_to_aliases", "except_aliases"): + annotation = XecGuardConfigModel.model_fields[field_name].annotation + assert _get_field_type_from_annotation(annotation) == "string" + + +# --------------------------------------------------------------------------- +# Per-virtual-key filtering in logging_only mode +# --------------------------------------------------------------------------- + + +class TestXecGuardLoggingHookKeyTargeting: + """logging_only dispatches to async_logging_hook, not apply_guardrail, so the + allow/deny lists and the per-key opt-out have to be enforced there as well — + otherwise an excluded key's content is still sent to the XecGuard backend.""" + + @staticmethod + def _kwargs(metadata, nested=True): + """Build logging-hook kwargs. The proxy puts the injected request + metadata under ``litellm_params.metadata``; ``nested=False`` exercises + the top-level ``metadata`` fallback.""" + base = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": {}, + } + if nested: + base["litellm_params"] = {"metadata": metadata} + else: + base["metadata"] = metadata + return base + + @pytest.mark.asyncio + async def test_blocklisted_key_is_not_scanned(self): + gr = _extension_guardrail(except_aliases=["internal"]) + kwargs = self._kwargs({"user_api_key_alias": "internal"}) + with patch.object(gr.async_handler, "post") as post: + out_kwargs, out_result = await gr.async_logging_hook( + kwargs=kwargs, + result=_build_model_response("some answer"), + call_type="acompletion", + ) + post.assert_not_called() + assert out_kwargs is kwargs + assert "guardrail_information" not in kwargs["standard_logging_object"] + + @pytest.mark.asyncio + async def test_key_not_in_allowlist_is_not_scanned(self): + gr = _extension_guardrail(apply_to_aliases=["prod"]) + kwargs = self._kwargs({"user_api_key_alias": "dev"}) + with patch.object(gr.async_handler, "post") as post: + await gr.async_logging_hook( + kwargs=kwargs, + result=_build_model_response("some answer"), + call_type="acompletion", + ) + post.assert_not_called() + + @pytest.mark.asyncio + async def test_top_level_metadata_is_also_honoured(self): + gr = _extension_guardrail(except_aliases=["internal"]) + kwargs = self._kwargs({"user_api_key_alias": "internal"}, nested=False) + with patch.object(gr.async_handler, "post") as post: + await gr.async_logging_hook( + kwargs=kwargs, + result=_build_model_response("some answer"), + call_type="acompletion", + ) + post.assert_not_called() + + @pytest.mark.asyncio + async def test_targeted_key_is_still_scanned(self): + gr = _extension_guardrail(apply_to_aliases=["prod"]) + kwargs = self._kwargs({"user_api_key_alias": "prod"}) + resp = _make_response({"decision": "SAFE", "trace_id": "lg-key-1"}) + with patch.object(gr.async_handler, "post", return_value=resp) as post: + await gr.async_logging_hook( + kwargs=kwargs, + result=_build_model_response("some answer"), + call_type="acompletion", + ) + post.assert_called_once() + info_list = kwargs["standard_logging_object"]["guardrail_information"] + assert info_list[0]["guardrail_response"]["trace_id"] == "lg-key-1" + + def test_key_context_lifts_nested_metadata(self): + # logging path: model_call_details carries the key fields one level down + gr = _extension_guardrail() + ctx = gr._key_context( + {"litellm_params": {"metadata": {"user_api_key_alias": "nested"}}} + ) + assert ctx["metadata"]["user_api_key_alias"] == "nested" + + def test_key_context_returns_proxy_shape_untouched(self): + # pre/during/post_call: reshaping to a single key would drop the other + # location _calling_key_identity also reads. + gr = _extension_guardrail() + data = { + "metadata": {"user_api_key_alias": "prod"}, + "litellm_metadata": {"user_api_key_hash": "hash-abc"}, + } + assert gr._key_context(data) is data + assert gr._calling_key_identity(gr._key_context(data)) == ("prod", "hash-abc") + + def test_key_context_top_level_wins_over_nested(self): + gr = _extension_guardrail() + ctx = gr._key_context( + { + "metadata": {"user_api_key_alias": "top"}, + "litellm_params": {"metadata": {"user_api_key_alias": "nested"}}, + } + ) + assert ctx["metadata"]["user_api_key_alias"] == "top" + + def test_key_context_passes_through_when_nothing_to_lift(self): + gr = _extension_guardrail() + assert gr._key_context(None) is None + empty: dict = {} + assert gr._key_context(empty) is empty + odd = {"litellm_params": {"metadata": "not-a-dict"}} + assert gr._key_context(odd) is odd + + +# --------------------------------------------------------------------------- +# should_run_guardrail: gate before LiteLLM records the guardrail as having run +# --------------------------------------------------------------------------- + + +class TestXecGuardShouldRunGuardrail: + @staticmethod + def _data(metadata): + return {"messages": [{"role": "user", "content": "hi"}], "metadata": metadata} + + def test_targeted_key_runs(self): + gr = _extension_guardrail(apply_to_aliases=["prod"]) + assert ( + gr.should_run_guardrail( + self._data({"user_api_key_alias": "prod"}), GuardrailEventHooks.pre_call + ) + is True + ) + + def test_key_not_in_allowlist_does_not_run(self): + gr = _extension_guardrail(apply_to_aliases=["prod"]) + assert ( + gr.should_run_guardrail( + self._data({"user_api_key_alias": "dev"}), GuardrailEventHooks.pre_call + ) + is False + ) + + def test_blocklisted_key_does_not_run(self): + gr = _extension_guardrail(except_aliases=["internal"]) + assert ( + gr.should_run_guardrail( + self._data({"user_api_key_alias": "internal"}), + GuardrailEventHooks.pre_call, + ) + is False + ) + + def test_no_lists_configured_runs(self): + gr = _extension_guardrail() + assert ( + gr.should_run_guardrail( + self._data({"user_api_key_alias": "anything"}), + GuardrailEventHooks.pre_call, + ) + is True + ) + + def test_native_decision_still_wins(self): + # super() says no (wrong event type for this hook) -> we must not override it + gr = _extension_guardrail(apply_to_aliases=["prod"]) # event_hook="pre_call" + assert ( + gr.should_run_guardrail( + self._data({"user_api_key_alias": "prod"}), + GuardrailEventHooks.post_call, + ) + is False + ) + + def test_native_opt_out_still_wins(self): + # admin-set disable_global_guardrails is honoured by super() even for a + # key that our own allow list would otherwise target + gr = _extension_guardrail(apply_to_aliases=["prod"]) + data = self._data( + { + "user_api_key_alias": "prod", + "user_api_key_metadata": {"disable_global_guardrails": True}, + } + ) + assert gr.should_run_guardrail(data, GuardrailEventHooks.pre_call) is False + + def test_logging_path_shape_is_understood(self): + # on the logging path `data` is model_call_details: key fields sit under + # litellm_params.metadata, and the gate must still find them + gr = _extension_guardrail( + except_aliases=["internal"], event_hook="logging_only" + ) + data = {"litellm_params": {"metadata": {"user_api_key_alias": "internal"}}} + assert gr.should_run_guardrail(data, GuardrailEventHooks.logging_only) is False + data_ok = {"litellm_params": {"metadata": {"user_api_key_alias": "other"}}} + assert ( + gr.should_run_guardrail(data_ok, GuardrailEventHooks.logging_only) is True + ) + + def test_mode_mismatch_is_left_to_super(self): + # a pre_call guardrail must not run for logging_only, whatever the lists say + gr = _extension_guardrail(event_hook="pre_call") + data = {"litellm_params": {"metadata": {"user_api_key_alias": "anything"}}} + assert gr.should_run_guardrail(data, GuardrailEventHooks.logging_only) is False + + +# --------------------------------------------------------------------------- +# Scan payload `meta` — caller context for XecGuard's SIEM export. +# +# Contract (POST /xecguard/v1/scan): meta is optional; when present virtualkey +# is required and must match ^[A-Za-z_][A-Za-z0-9_.-]{0,63}$; data is a flat +# string->string map of at most 32 fields, keys on the same pattern, values +# <=512 chars, whole meta <=4096 bytes serialized. A violation is a 400, which +# block_on_error would turn into a user-visible block — so everything is coerced +# or dropped here instead. +# --------------------------------------------------------------------------- + + +def _meta_request_data(key_metadata=None, alias="team-alpha", **metadata): + """A pre/during/post_call request_data with the proxy-injected key fields.""" + meta: dict = {"messages": [{"role": "user", "content": "hi"}]} + injected = {"user_api_key_alias": alias, **metadata} + if key_metadata is not None: + injected["user_api_key_metadata"] = key_metadata + meta["metadata"] = injected + return meta + + +def _admin_data(meta): + """``meta.data`` with the proxy-injected attributes removed. + + ``send_meta`` forwards two merged sources: the attributes the proxy injects + about the calling key, and the free-form metadata an admin typed on the + Virtual Keys page. The tests below are about the second source's coercion + rules, so they filter the first out rather than restating it 20 times -- + ``TestXecGuardScanMetaAutoFields`` covers the injected half on its own. + """ + auto = {name for name, _ in xecguard_module._META_AUTO_DATA_FIELDS} + return {k: v for k, v in (meta.get("data") or {}).items() if k not in auto} + + +async def _sent_payload(gr, request_data): + """Run one safe scan and return the JSON body that reached the backend.""" + resp = _make_response( + {"decision": "SAFE", "trace_id": "meta-tr", "xecguard_result": []} + ) + with patch.object(gr.async_handler, "post", return_value=resp) as mock_post: + await gr.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + return mock_post.call_args.kwargs["json"] + + +@pytest.mark.usefixtures("_clean_env") +class TestXecGuardScanMeta: + @pytest.mark.asyncio + async def test_meta_absent_by_default(self): + # send_meta is opt-in: enabling it forwards the key alias and the key's + # admin-set metadata to XecGuard, so an upgrade must not start doing it. + gr = _extension_guardrail() + assert gr.send_meta is False + payload = await _sent_payload(gr, _meta_request_data({"cost_center": "CC-42"})) + assert "meta" not in payload + + @pytest.mark.asyncio + async def test_virtualkey_is_the_alias_the_guardrail_filtered_on(self): + gr = _extension_guardrail(send_meta=True, apply_to_aliases=["team-alpha"]) + payload = await _sent_payload(gr, _meta_request_data(alias="team-alpha")) + assert payload["meta"]["virtualkey"] == "team-alpha" + assert _admin_data(payload["meta"]) == {} + + @pytest.mark.asyncio + async def test_data_is_the_virtual_keys_page_metadata(self): + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _meta_request_data({"cost_center": "CC-42", "owner": "alice@corp"}), + ) + assert payload["meta"]["virtualkey"] == "team-alpha" + assert _admin_data(payload["meta"]) == { + "cost_center": "CC-42", + "owner": "alice@corp", + } + + @pytest.mark.asyncio + async def test_meta_does_not_disturb_the_existing_payload(self): + gr = _extension_guardrail(send_meta=True, policy_names=["jailbreak"]) + payload = await _sent_payload(gr, _meta_request_data({"a": "b"})) + assert payload["scan_type"] == "input" + assert payload["policy_names"] == ["jailbreak"] + assert payload["messages"] == [{"role": "user", "content": "hi"}] + + @pytest.mark.asyncio + async def test_hash_is_the_fallback_when_the_key_has_no_alias(self): + gr = _extension_guardrail(send_meta=True) + data = _meta_request_data(alias=None, user_api_key_hash="abc123def") + payload = await _sent_payload(gr, data) + assert payload["meta"]["virtualkey"] == "abc123def" + + @pytest.mark.asyncio + async def test_meta_omitted_when_no_identity_matches_the_pattern(self): + # master-key calls carry neither field; a hash starting with a digit + # cannot satisfy the backend pattern. Omit meta rather than 400 the scan. + gr = _extension_guardrail(send_meta=True) + assert "meta" not in await _sent_payload(gr, _meta_request_data(alias=None)) + digit_hash = _meta_request_data(alias=None, user_api_key_hash="9abcdef") + assert "meta" not in await _sent_payload(gr, digit_hash) + + @pytest.mark.asyncio + async def test_alias_failing_the_pattern_falls_through_to_the_hash(self): + gr = _extension_guardrail(send_meta=True) + data = _meta_request_data(alias="has spaces", user_api_key_hash="deadbeef") + payload = await _sent_payload(gr, data) + assert payload["meta"]["virtualkey"] == "deadbeef" + + @pytest.mark.asyncio + async def test_scalars_are_stringified_and_nested_values_dropped(self): + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _meta_request_data( + { + "tier": 3, + "ratio": 1.5, + "beta": True, + "ga": False, + "absent": None, + "nested": {"x": 1}, + "listed": ["a"], + "blank": "", + } + ), + ) + assert _admin_data(payload["meta"]) == { + "tier": "3", + "ratio": "1.5", + "beta": "true", + "ga": "false", + } + + @pytest.mark.asyncio + async def test_a_non_string_field_name_is_dropped_not_raised(self): + """The metadata dict is decoded JSON that nothing validates on the way in. + + `meta.data` promises to coerce or drop, never to fail the scan, and it is + built outside any try on the pre/during/post_call paths -- so a key that + reaches `re.match` unchecked would turn one malformed field into a 500 for + every request from that key. + """ + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, _meta_request_data({1: "oops", "tier": "gold"}) + ) + assert _admin_data(payload["meta"]) == {"tier": "gold"} + + @pytest.mark.asyncio + async def test_field_names_off_the_pattern_are_dropped(self): + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _meta_request_data( + {"ok.name-1": "v", "bad name": "v", "9lead": "v", "": "v"} + ), + ) + assert _admin_data(payload["meta"]) == {"ok.name-1": "v"} + + @pytest.mark.asyncio + async def test_control_characters_are_stripped_and_values_truncated(self): + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _meta_request_data({"note": "a\x00b\x1fc\x7fd", "long": "x" * 700}), + ) + assert payload["meta"]["data"]["note"] == "abcd" + assert len(payload["meta"]["data"]["long"]) == 512 + + @pytest.mark.asyncio + async def test_at_most_32_fields(self): + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, _meta_request_data({f"f{i}": str(i) for i in range(40)}) + ) + assert len(payload["meta"]["data"]) == 32 + assert "f0" in payload["meta"]["data"] and "f39" not in payload["meta"]["data"] + + @pytest.mark.asyncio + async def test_serialized_cap_sheds_fields_but_keeps_virtualkey(self): + gr = _extension_guardrail(send_meta=True) + # 20 x ~500 chars is far past 4096 bytes; a later small field still fits. + key_metadata = {f"big{i}": "y" * 500 for i in range(20)} + key_metadata["small"] = "s" + payload = await _sent_payload(gr, _meta_request_data(key_metadata)) + meta = payload["meta"] + assert meta["virtualkey"] == "team-alpha" + assert len(json.dumps(meta, ensure_ascii=False).encode("utf-8")) <= 4096 + assert meta["data"]["small"] == "s" + + @pytest.mark.asyncio + async def test_utf8_values_are_measured_in_bytes(self): + gr = _extension_guardrail(send_meta=True) + # 3 bytes per CJK char: 500 chars pass the char cap but eat 1500 bytes. + key_metadata = {f"cjk{i}": "資" * 500 for i in range(6)} + payload = await _sent_payload(gr, _meta_request_data(key_metadata)) + blob = json.dumps(payload["meta"], ensure_ascii=False).encode("utf-8") + assert len(blob) <= 4096 + + @pytest.mark.asyncio + async def test_meta_data_fields_narrows_the_forwarded_set(self): + gr = _extension_guardrail(send_meta=True, meta_data_fields=["cost_center"]) + payload = await _sent_payload( + gr, _meta_request_data({"cost_center": "CC-42", "owner": "alice@corp"}) + ) + assert payload["meta"]["data"] == {"cost_center": "CC-42"} + + @pytest.mark.asyncio + async def test_callback_credential_slots_are_never_forwarded(self): + # the proxy strips these before injecting; belt-and-braces here because a + # leak would ship per-key integration credentials to an external service + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _meta_request_data( + { + "logging": "langfuse-secret", + "callback_settings": "s", + "secret_manager_settings": "sm", + "keep": "v", + } + ), + ) + assert _admin_data(payload["meta"]) == {"keep": "v"} + + @pytest.mark.asyncio + async def test_credential_slots_cannot_be_opted_back_in(self): + # meta_data_fields is an admin convenience, not an override for the + # credential blocklist -- naming one must not spring the leak + gr = _extension_guardrail(send_meta=True, meta_data_fields=["logging", "keep"]) + payload = await _sent_payload( + gr, _meta_request_data({"logging": "langfuse-secret", "keep": "v"}) + ) + assert _admin_data(payload["meta"]) == {"keep": "v"} + + @pytest.mark.asyncio + async def test_proxy_control_settings_are_skipped_by_default(self): + # The Virtual Keys page writes the proxy's own per-key knobs into the same + # metadata dict as the admin's fields. They are configuration, not caller + # identity: noise in a SIEM, they eat the 32-field budget, and + # disable_global_guardrails describes the key's security posture. + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _meta_request_data( + { + "tag_rpm_limit": "{}", + "throttle_on_budget_exceeded": False, + "disable_global_guardrails": True, + "enforced_params": "x", + "cost_center": "CC-42", + } + ), + ) + assert _admin_data(payload["meta"]) == {"cost_center": "CC-42"} + + @pytest.mark.asyncio + async def test_a_named_control_setting_is_opted_back_in(self): + # unlike the credential slots, forwarding these is a judgement call -- + # a deployment that wants its rate-limit config in the SIEM can say so + gr = _extension_guardrail( + send_meta=True, + meta_data_fields=["throttle_on_budget_exceeded", "cost_center"], + ) + payload = await _sent_payload( + gr, + _meta_request_data( + {"throttle_on_budget_exceeded": False, "cost_center": "CC-42"} + ), + ) + assert _admin_data(payload["meta"]) == { + "throttle_on_budget_exceeded": "false", + "cost_center": "CC-42", + } + + def test_the_control_field_list_tracks_the_proxys_own(self): + """The skip list is derived from the proxy's lists, not copied. + + A hardcoded copy silently rots: litellm adds a metadata-backed knob, it + starts appearing as ctx_ in the customer's SIEM, and nobody + notices. Assert the derivation instead of the contents. + """ + from litellm.proxy._types import ( + LiteLLM_ManagementEndpoint_MetadataFields, + LiteLLM_ManagementEndpoint_MetadataFields_Premium, + ) + + proxy_fields = set(LiteLLM_ManagementEndpoint_MetadataFields) | set( + LiteLLM_ManagementEndpoint_MetadataFields_Premium + ) + control = xecguard_module._META_CONTROL_DATA_FIELDS + excluded = xecguard_module._META_EXCLUDED_DATA_FIELDS + assert control == proxy_fields - excluded + # the two tiers must not overlap, or the opt-in path would reach a + # credential slot + assert not (control & excluded) + # sanity: the fields that prompted this are actually covered + assert {"tag_rpm_limit", "throttle_on_budget_exceeded"} <= control + assert "logging" in excluded and "logging" not in control + + @pytest.mark.asyncio + async def test_data_omitted_when_nothing_survives_the_filter(self): + # empty `data` is not the same as absent: send meta without the key. + # Reaching that now takes a meta_data_fields matching neither source, + # since a key with an alias always contributes at least `key_alias`. + gr = _extension_guardrail(send_meta=True, meta_data_fields=["no_such_field"]) + for request_data in (_meta_request_data(), _meta_request_data({})): + payload = await _sent_payload(gr, request_data) + assert payload["meta"] == {"virtualkey": "team-alpha"} + + @pytest.mark.asyncio + async def test_data_is_only_the_injected_attributes_when_the_key_has_no_metadata( + self, + ): + gr = _extension_guardrail(send_meta=True) + for request_data in (_meta_request_data(), _meta_request_data({})): + payload = await _sent_payload(gr, request_data) + assert payload["meta"]["data"] == {"key_alias": "team-alpha"} + + @pytest.mark.asyncio + async def test_response_scan_carries_meta_too(self): + gr = _extension_guardrail(send_meta=True, event_hook="post_call") + data = _meta_request_data({"cost_center": "CC-42"}) + data["response"] = _build_model_response("the answer") + resp = _make_response({"decision": "SAFE", "trace_id": "meta-post"}) + with patch.object(gr.async_handler, "post", return_value=resp) as mock_post: + await gr.apply_guardrail( + inputs={"texts": ["the answer"]}, + request_data=data, + input_type="response", + ) + payload = mock_post.call_args.kwargs["json"] + assert payload["scan_type"] == "response" + assert _admin_data(payload["meta"]) == {"cost_center": "CC-42"} + + @pytest.mark.asyncio + async def test_logging_only_path_carries_meta_from_nested_metadata(self): + # logging_only goes through async_logging_hook, where the injected key + # fields sit under litellm_params.metadata rather than at the top level + gr = _extension_guardrail(send_meta=True, event_hook="logging_only") + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": {}, + "litellm_params": { + "metadata": { + "user_api_key_alias": "team-alpha", + "user_api_key_metadata": {"cost_center": "CC-42"}, + } + }, + } + resp = _make_response({"decision": "SAFE", "trace_id": "meta-log"}) + with patch.object(gr.async_handler, "post", return_value=resp) as mock_post: + await gr.async_logging_hook( + kwargs=kwargs, + result=_build_model_response("some answer"), + call_type="acompletion", + ) + payload = mock_post.call_args.kwargs["json"] + assert payload["meta"]["virtualkey"] == "team-alpha" + assert _admin_data(payload["meta"]) == {"cost_center": "CC-42"} + + @pytest.mark.asyncio + async def test_grounding_call_does_not_get_meta(self): + # the contract defines meta for /scan only + gr = _extension_guardrail(send_meta=True, event_hook="post_call") + data = _meta_request_data({"cost_center": "CC-42"}) + data["response"] = _build_model_response("Peggy Seeger was American.") + data["metadata"]["xecguard_grounding_documents"] = [ + {"document_id": "d1", "context": "Peggy Seeger is American."} + ] + scan_ok = _make_response({"decision": "SAFE", "trace_id": "g1"}) + grounding_ok = _make_response({"decision": "SAFE", "trace_id": "g2"}) + with patch.object( + gr.async_handler, "post", side_effect=[scan_ok, grounding_ok] + ) as mock_post: + await gr.apply_guardrail( + inputs={"texts": ["text"]}, + request_data=data, + input_type="response", + ) + assert "meta" in mock_post.call_args_list[0].kwargs["json"] + assert "meta" not in mock_post.call_args_list[1].kwargs["json"] + + def test_env_var_enables_meta(self): + with patch.dict(os.environ, {"XECGUARD_SEND_META": "true"}, clear=False): + assert _extension_guardrail().send_meta is True + with patch.dict(os.environ, {"XECGUARD_SEND_META": "false"}, clear=False): + assert _extension_guardrail().send_meta is False + + def test_explicit_config_beats_the_env_var(self): + with patch.dict(os.environ, {"XECGUARD_SEND_META": "true"}, clear=False): + assert _extension_guardrail(send_meta=False).send_meta is False + + def test_config_model_normalizes_meta_data_fields(self): + cfg = XecGuardConfigModel(meta_data_fields="cost_center, owner ") + assert cfg.meta_data_fields == ["cost_center", "owner"] + assert XecGuardConfigModel().meta_data_fields is None + + def test_an_explicit_none_stays_none(self): + # distinct from omitting the field: a before-validator does not run on an + # unset default, so this is the only way that branch is reached + assert XecGuardConfigModel(apply_to_aliases=None).apply_to_aliases is None + assert XecGuardConfigModel(meta_data_fields=None).meta_data_fields is None + + def test_a_wrong_type_is_handed_back_for_pydantic_to_reject(self): + """The normalizer returns an unrecognised value unchanged on purpose. + + Coercing it -- to None, or to [] -- would turn a typo in the config into a + silently empty allowlist, which for `apply_to_aliases` means scanning every + key instead of the chosen ones. Letting pydantic reject it keeps the + mistake loud. + """ + for bad in (5, {"a": 1}, True): + with pytest.raises(ValidationError): + XecGuardConfigModel(apply_to_aliases=bad) + + +# --------------------------------------------------------------------------- +# meta.data also carries what the *proxy* knows about the calling key, not just +# what an admin typed. Without it a SIEM event is a bare alias, and attributing +# it needs a lookup back into the proxy database that whoever reads the SIEM +# usually cannot make. These fields include PII (user_email) and commercials +# (spend, max_budget) and leave the proxy only when send_meta is on. +# --------------------------------------------------------------------------- + + +def _injected_request_data(**injected): + """request_data whose metadata holds only proxy-injected key attributes.""" + return {"messages": [{"role": "user", "content": "hi"}], "metadata": injected} + + +@pytest.mark.usefixtures("_clean_env") +class TestXecGuardScanMetaAutoFields: + @pytest.mark.asyncio + async def test_the_full_injected_set_is_forwarded(self): + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _injected_request_data( + user_api_key_hash="abc123", + user_api_key_alias="team-alpha", + user_api_key_team_id="t-1", + user_api_key_team_alias="platform", + user_api_key_user_id="u-1", + user_api_key_user_email="alice@corp", + user_api_key_org_id="o-1", + user_api_key_org_alias="acme", + user_api_key_project_id="p-1", + user_api_key_project_alias="proj", + user_api_key_end_user_id="eu-1", + user_api_key_spend=1.5, + user_api_key_max_budget=100, + user_api_key_request_route="/chat/completions", + ), + ) + assert payload["meta"]["data"] == { + "key_id": "abc123", + "key_alias": "team-alpha", + "team_id": "t-1", + "team_alias": "platform", + "user_id": "u-1", + "user_email": "alice@corp", + "org_id": "o-1", + "org_alias": "acme", + "project_id": "p-1", + "project_alias": "proj", + "end_user_id": "eu-1", + "spend": "1.5", + "max_budget": "100", + "request_route": "/chat/completions", + } + + @pytest.mark.asyncio + async def test_absent_and_null_attributes_are_skipped(self): + # a key with no team must contribute no team_id rather than an empty one: + # a SIEM query for "scans with no team" should mean it, not match every key + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _injected_request_data( + user_api_key_alias="team-alpha", + user_api_key_team_id=None, + user_api_key_user_email=None, + ), + ) + assert payload["meta"]["data"] == {"key_alias": "team-alpha"} + + @pytest.mark.asyncio + async def test_injected_order_is_the_constants_order(self): + # identity, then tenancy, then commercials - so the fields that survive + # the 32-field / 4096-byte caps are the ones worth keeping, whatever + # order the proxy happened to build its metadata dict in + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _injected_request_data( + user_api_key_spend=1.5, + user_api_key_team_id="t-1", + user_api_key_alias="team-alpha", + user_api_key_hash="abc123", + ), + ) + assert list(payload["meta"]["data"]) == [ + "key_id", + "key_alias", + "team_id", + "spend", + ] + + @pytest.mark.asyncio + async def test_an_admin_cannot_shadow_an_injected_attribute(self): + # otherwise a key whose owner controls its metadata could write + # key_id/team_id of someone else and mislead an investigation + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _meta_request_data( + {"key_id": "not-mine", "team_id": "not-mine", "mine": "ok"}, + alias="team-alpha", + user_api_key_hash="abc123", + user_api_key_team_id="t-1", + ), + ) + data = payload["meta"]["data"] + assert data["key_id"] == "abc123" + assert data["team_id"] == "t-1" + assert data["mine"] == "ok" + + @pytest.mark.asyncio + async def test_meta_data_fields_narrows_the_injected_set_too(self): + # the allowlist is what a deployment forbidden to egress PII or spend + # figures uses, so it has to bind the injected half as well + gr = _extension_guardrail( + send_meta=True, meta_data_fields=["key_id", "cost_center"] + ) + payload = await _sent_payload( + gr, + _meta_request_data( + {"cost_center": "CC-42", "owner": "alice@corp"}, + alias="team-alpha", + user_api_key_hash="abc123", + user_api_key_user_email="alice@corp", + user_api_key_spend=1.5, + ), + ) + assert payload["meta"]["data"] == {"key_id": "abc123", "cost_center": "CC-42"} + + @pytest.mark.asyncio + async def test_nested_metadata_is_read_on_the_logging_only_path(self): + # async_logging_hook sees the injected fields under litellm_params.metadata + gr = _extension_guardrail(send_meta=True, event_hook="logging_only") + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": {}, + "litellm_params": { + "metadata": { + "user_api_key_alias": "team-alpha", + "user_api_key_hash": "abc123", + "user_api_key_team_alias": "platform", + } + }, + } + resp = _make_response({"decision": "SAFE", "trace_id": "meta-auto-log"}) + with patch.object(gr.async_handler, "post", return_value=resp) as mock_post: + await gr.async_logging_hook( + kwargs=kwargs, + result=_build_model_response("some answer"), + call_type="acompletion", + ) + assert mock_post.call_args.kwargs["json"]["meta"]["data"] == { + "key_id": "abc123", + "key_alias": "team-alpha", + "team_alias": "platform", + } + + @pytest.mark.asyncio + async def test_injected_attributes_obey_the_value_coercion_rules(self): + gr = _extension_guardrail(send_meta=True) + payload = await _sent_payload( + gr, + _injected_request_data( + user_api_key_alias="team-alpha", + user_api_key_user_email="a\x00b@corp", + user_api_key_team_alias="y" * 700, + user_api_key_model_max_budget={"gpt-4o": 1}, + ), + ) + data = payload["meta"]["data"] + assert data["user_email"] == "ab@corp" + assert len(data["team_alias"]) == 512 + # model_max_budget is not in the forwarded set, and would be dropped as a + # nested value even if it were + assert "model_max_budget" not in data + + @pytest.mark.asyncio + async def test_nothing_is_forwarded_while_send_meta_is_off(self): + gr = _extension_guardrail() + payload = await _sent_payload( + gr, + _injected_request_data( + user_api_key_alias="team-alpha", user_api_key_hash="abc123" + ), + ) + assert "meta" not in payload + + +# --------------------------------------------------------------------------- +# meta.virtualkey has two wire shapes. The string form is the alias, which is +# all the deployed backend accepts; the object form carries {alias, key_id} so a +# scan stays attributable when the alias is absent, renamed, or reused. Sending +# the object form to a string-only backend is a 400, which block_on_error turns +# into a block for every request - hence a switch, defaulting to string. +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("_clean_env") +class TestXecGuardScanMetaVirtualkeyObject: + def test_string_is_the_default(self): + assert _extension_guardrail(send_meta=True).meta_identity_format == "string" + + def test_env_var_selects_the_format(self): + with patch.dict( + os.environ, {"XECGUARD_META_IDENTITY_FORMAT": "OBJECT"}, clear=False + ): + assert _extension_guardrail(send_meta=True).meta_identity_format == "object" + + def test_explicit_config_beats_the_env_var(self): + with patch.dict( + os.environ, {"XECGUARD_META_IDENTITY_FORMAT": "object"}, clear=False + ): + gr = _extension_guardrail(send_meta=True, meta_identity_format="string") + assert gr.meta_identity_format == "string" + + def test_an_unknown_value_falls_back_instead_of_raising(self): + # a typo in the UI must not take the gateway down at startup + gr = _extension_guardrail(send_meta=True, meta_identity_format="objekt") + assert gr.meta_identity_format == "string" + + def test_config_model_offers_both_shapes_to_the_ui(self): + assert XecGuardConfigModel().meta_identity_format is None + assert ( + XecGuardConfigModel(meta_identity_format="object").meta_identity_format + == "object" + ) + with pytest.raises(ValidationError): + XecGuardConfigModel(meta_identity_format="objekt") + + def test_no_non_secret_field_is_masked_on_the_way_back_to_the_ui(self): + """Only `api_key` may be masked when the proxy serves this provider's + params back to the UI. + + ``_get_masked_values`` matches on *substrings* of the field name -- + "key", "token", "secret", "credentials", "password" -- and the guardrail + endpoints run every ``litellm_params`` dict through it. A non-secret + field caught by that heuristic is served to the UI as "ob****ct"; the + edit form prefills it, saving writes the masked string back, and the + value is silently wrong from then on. That is why this field is + ``meta_identity_format`` and not ``meta_virtualkey_format``: the latter + matches on "virtual*key*". + + Enum and free-text fields are the dangerous ones. Booleans and lists + survive by type (``_mask_value`` returns non-str unchanged), so a name + collision there is latent rather than live -- still worth failing on, + since changing such a field to a string would spring the trap. + """ + from litellm.litellm_core_utils.litellm_logging import _get_masked_values + + probe = {name: "sentinel" for name in XecGuardConfigModel.model_fields} + served = _get_masked_values(probe, unmasked_length=4, number_of_asterisks=4) + masked = {name for name, value in served.items() if value != "sentinel"} + assert masked == {"api_key"}, ( + "these XecGuard params would reach the UI masked and be corrupted by a " + f"form save: {sorted(masked - {'api_key'})}. Rename them off the " + "substrings _get_masked_values matches on." + ) + + @pytest.mark.asyncio + async def test_object_form_carries_alias_and_key_id(self): + gr = _extension_guardrail(send_meta=True, meta_identity_format="object") + payload = await _sent_payload( + gr, _meta_request_data(alias="team-alpha", user_api_key_hash="abc123") + ) + assert payload["meta"]["virtualkey"] == { + "alias": "team-alpha", + "key_id": "abc123", + } + + @pytest.mark.asyncio + async def test_either_member_may_be_absent(self): + gr = _extension_guardrail(send_meta=True, meta_identity_format="object") + no_alias = await _sent_payload( + gr, _injected_request_data(user_api_key_hash="abc123") + ) + assert no_alias["meta"]["virtualkey"] == {"key_id": "abc123"} + no_hash = await _sent_payload(gr, _meta_request_data(alias="team-alpha")) + assert no_hash["meta"]["virtualkey"] == {"alias": "team-alpha"} + + @pytest.mark.asyncio + async def test_meta_is_omitted_when_the_key_has_neither(self): + # a master-key call: nothing to correlate on, so omit meta rather than + # send an empty object and collect a 400 + gr = _extension_guardrail(send_meta=True, meta_identity_format="object") + assert "meta" not in await _sent_payload(gr, _injected_request_data()) + + @pytest.mark.asyncio + async def test_the_object_form_lifts_the_identifier_pattern_on_aliases(self): + # the pattern exists because a bare string becomes a SIEM field *value* + # directly; an object member is sanitized like a meta.data value instead, + # so aliases with spaces or CJK become correlatable + gr = _extension_guardrail(send_meta=True, meta_identity_format="object") + payload = await _sent_payload( + gr, _meta_request_data(alias="研發 team", user_api_key_hash="abc123") + ) + assert payload["meta"]["virtualkey"] == { + "alias": "研發 team", + "key_id": "abc123", + } + # the string form cannot: it falls through the failing alias to the hash + string_gr = _extension_guardrail(send_meta=True) + string_payload = await _sent_payload( + string_gr, _meta_request_data(alias="研發 team", user_api_key_hash="abc123") + ) + assert string_payload["meta"]["virtualkey"] == "abc123" + + @pytest.mark.asyncio + async def test_object_members_are_sanitized_and_truncated(self): + gr = _extension_guardrail(send_meta=True, meta_identity_format="object") + payload = await _sent_payload( + gr, _meta_request_data(alias="a\x00b\x1fc", user_api_key_hash="h" * 700) + ) + virtualkey = payload["meta"]["virtualkey"] + assert virtualkey["alias"] == "abc" + assert len(virtualkey["key_id"]) == 512 + + @pytest.mark.asyncio + async def test_the_serialized_cap_still_holds_with_the_object_form(self): + gr = _extension_guardrail(send_meta=True, meta_identity_format="object") + key_metadata = {f"big{i}": "y" * 500 for i in range(20)} + payload = await _sent_payload( + gr, + _meta_request_data( + key_metadata, alias="team-alpha", user_api_key_hash="abc123" + ), + ) + meta = payload["meta"] + assert meta["virtualkey"] == {"alias": "team-alpha", "key_id": "abc123"} + assert len(json.dumps(meta, ensure_ascii=False).encode("utf-8")) <= 4096 + + @pytest.mark.asyncio + async def test_data_is_unaffected_by_the_format(self): + gr = _extension_guardrail(send_meta=True, meta_identity_format="object") + payload = await _sent_payload( + gr, + _meta_request_data( + {"cost_center": "CC-42"}, alias="team-alpha", user_api_key_hash="abc123" + ), + ) + assert _admin_data(payload["meta"]) == {"cost_center": "CC-42"} + assert payload["meta"]["data"]["key_id"] == "abc123"