mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(guardrails)!: redact secrets in tool traffic and system prompts (#47122)
* fix(guardrails): redact secrets in tool traffic and Anthropic system * fix(guardrails): bound tool payload scans and stop nested tool_result recursion Scan each tool value beside its key instead of an indented dump, read tool_result content one level deep, keep non-finite tool arguments as raw text, and give read-only guardrails one fragment per tool payload * test(guardrails): cover read-only guardrail consumers of tool and system text * fix(guardrails): keep tool argument parsing, scanning and tests within bounds Treat tool arguments that decode to a lone surrogate as raw text, decode JSON-escaped detections back to the leaf value, quote unkeyed leaves without repeating a parent key per list item, drop cycle tracking that parsed JSON cannot reach, and use the CI license instead of a fake license server * fix(guardrails): scan list items beside a capped parent key List items keep their parent key in the scanned JSON line, capped to its last 64 characters, so a credential key over a list is still detected while a long key over many items no longer multiplies the scan text * fix(guardrails): cap the scanned key after JSON escaping A key of control characters escaped to six times its length after the cap, so every list item's scan line grew with it * fix(guardrails): collapse whitespace in the scanned key context Space-heavy keys made every inherited list item scan several times slower than a short key * fix(guardrails): build each key's scan context once per payload Repeating the collapse, escape and cut for every list item made a long key over many items cost key length times item count * fix(guardrails): cover Responses tool call items and instructions Responses carries the same tool call arguments and system prompt as chat and messages under function_call.arguments, custom_tool_call.input and instructions, so secrets there reached the provider unscanned * perf(guardrails): skip rescanning text hide-secrets already found clean Agents resend their whole conversation on every turn, so each turn rescanned every earlier tool result. Each guardrail now keeps a bounded set of digests of scan inputs that held no secret and skips those; text with a secret is always rescanned and no secret material is kept * fix(guardrails): size the clean-scan cache for many sessions and replace longest secrets first 4096 digests held about 20 long agent sessions per worker, so interleaved sessions evicted each other and every resend rescanned its history; 65536 digests (about 10 MB) hold a few hundred. Tool payload secrets are now replaced longest first after JSON unescaping, so a shorter overlapping match cannot leave part of a secret behind. The key-context test checks the work it bounds instead of wall-clock time, which CI could not meet * fix(guardrails): keep unmasked tool argument text and redact keyword matches cut inside an escape * fix(guardrails): write a tool visitor's whole payload when it changes more than string values * fix(guardrails): read message text before tool and system text so a vendor error on the new text cannot skip it * fix(guardrails): moderate tool and system text in its own call so a long payload cannot hide a flagged message * fix(guardrails): skip the empty moderation call when only tool or system text is present * fix(guardrails): scan long multi-line tool strings as their own lines so the scan stays linear * fix(guardrails): keep the key beside every line of a long multi-line tool string * fix(guardrails): confine tool-traffic coverage to hide-secrets Restore iter_message_text to its main behavior, so banned keywords, Azure Content Safety, IBM detectors and the OpenAI and Google moderation hooks read exactly what they read on main. Only walk_user_text, which hide-secrets alone calls, still reaches system, instructions and tool payloads. * test(guardrails): stop the tool-traffic integration tests from flaking on a busy host Use the shared per-case timeout for the rig client, and open a new connection per call in the worker-kill burst so a pooled keep-alive socket to the killed worker is never reused. * test(guardrails): drop the wall-clock check from the deep tool arguments unit test The integration test test_deeply_nested_tool_arguments_do_not_stall_the_worker already covers the timing; the unit test keeps the redaction check. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
parent
b9e4298506
commit
15eca24fc2
6 changed files with 1828 additions and 15 deletions
|
|
@ -13,13 +13,19 @@ sys.path.insert(
|
|||
) # Adds the parent directory to the system path
|
||||
import configparser
|
||||
import contextlib
|
||||
import hashlib
|
||||
import itertools
|
||||
import json
|
||||
import re
|
||||
import tempfile
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Generator, Iterator, Sequence
|
||||
from contextvars import ContextVar
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
|
|
@ -27,7 +33,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
log_guardrail_information,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails._content_utils import walk_user_text
|
||||
from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
|
@ -40,7 +46,7 @@ GUARDRAIL_PROVIDER = "hide-secrets"
|
|||
|
||||
# Per-invocation tally of redacted secrets by detect-secrets plugin type; None
|
||||
# means the guardrail did not run, so _process_response records nothing.
|
||||
_masked_entity_count: ContextVar[Optional[dict]] = ContextVar(
|
||||
_masked_entity_count: ContextVar[Optional[dict[str, int]]] = ContextVar(
|
||||
"hide_secrets_masked_entity_count", default=None
|
||||
)
|
||||
|
||||
|
|
@ -449,6 +455,19 @@ _default_detect_secrets_config = {
|
|||
|
||||
_CONFIG_SECTION: Final = "litellm-prompt"
|
||||
|
||||
_JSON_STRING: Final = TypeAdapter(str)
|
||||
|
||||
_KEY_CONTEXT: Final = 64
|
||||
|
||||
# Longest JSON line scanned for a multi-line tool string. Escaping puts the whole string on one line, and
|
||||
# detect-secrets regexes backtrack across it, so a longer multi-line string is cut into runs of whole lines,
|
||||
# each scanned as its own JSON line beside the key. Every run holds at least two lines, so the high-entropy
|
||||
# detectors, which match only a quoted string made wholly of their charset, never judge a lone line.
|
||||
_JSON_LINE_LIMIT: Final = 4096
|
||||
|
||||
# A detector can stop a match inside an escape of a JSON-escaped line, as at the backslash of \".
|
||||
_CUT_ESCAPE: Final = re.compile(r"(?<=.)\\(?:u[0-9A-Fa-f]{0,3})?\Z", re.DOTALL)
|
||||
|
||||
_ASSIGNMENT_LINE: Final = re.compile(r"[^\s\[#;:=][^:=]*[:=]")
|
||||
|
||||
_SHELL_ASSIGNMENT: Final = re.compile(r"(?P<key>[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P<value>\S+)")
|
||||
|
|
@ -459,6 +478,10 @@ _SHELL_TRAILER: Final = re.compile(r"\\|#.*|-*\w[\w.-]*=\S*")
|
|||
|
||||
_SCAN_SUFFIX: Final = ".py"
|
||||
|
||||
# Digests of scan inputs that held no secret, so the history an agent resends every turn is scanned once. 65536
|
||||
# digests take about 10 MB and hold a few hundred interleaved agent sessions per worker.
|
||||
_CLEAN_SCAN_CACHE_SIZE: Final = 65536
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _temp_file(text: str) -> Generator[str, None, None]:
|
||||
|
|
@ -532,6 +555,57 @@ def _quoted_assignments(text: str) -> tuple[str, ...]:
|
|||
)
|
||||
|
||||
|
||||
def _json_unescaped(value: str) -> str:
|
||||
for candidate in (value, _CUT_ESCAPE.sub("", value)):
|
||||
try:
|
||||
return _JSON_STRING.validate_json(f'"{candidate}"')
|
||||
except ValueError:
|
||||
continue
|
||||
return value
|
||||
|
||||
|
||||
def _key_context(key: str) -> str:
|
||||
return json.dumps(" ".join(key.split()), ensure_ascii=False)[1:-1][-_KEY_CONTEXT:]
|
||||
|
||||
|
||||
def _json_line(context: str | None, text: str) -> str:
|
||||
value: Final = json.dumps(text, ensure_ascii=False)
|
||||
return value if context is None else f'"{context}": {value}'
|
||||
|
||||
|
||||
def _json_lines(context: str | None, text: str) -> Iterator[str]:
|
||||
"""Yield ``text`` as one JSON line, or a long multi-line ``text`` as JSON lines of two or more whole lines."""
|
||||
line: Final = _json_line(context, text)
|
||||
lines: Final = text.splitlines(keepends=True)
|
||||
if len(line) <= _JSON_LINE_LIMIT or len(lines) < 2:
|
||||
yield line
|
||||
return
|
||||
budget: Final = _JSON_LINE_LIMIT - len(_json_line(context, ""))
|
||||
sizes: Final = (len(json.dumps(part, ensure_ascii=False)) - 2 for part in lines)
|
||||
states: Final = itertools.accumulate(sizes, partial(_next_run_state, budget), initial=(0, 0, 0))
|
||||
run_ids: Final = (state[0] for state in itertools.islice(states, 1, None))
|
||||
runs: Final = tuple(
|
||||
tuple(part for part, _ in run) for _, run in itertools.groupby(zip(lines, run_ids), key=lambda pair: pair[1])
|
||||
)
|
||||
whole_runs: Final = (*runs[:-2], runs[-2] + runs[-1]) if len(runs[-1]) < 2 else runs
|
||||
for run in whole_runs:
|
||||
yield _json_line(context, "".join(run[:-1]) + run[-1].splitlines()[0])
|
||||
|
||||
|
||||
def _next_run_state(budget: int, state: tuple[int, int, int], size: int) -> tuple[int, int, int]:
|
||||
"""Advance ``(run id, lines in run, escaped size of run)`` by one line of escaped ``size``."""
|
||||
run_id, count, total = state
|
||||
if count >= 2 and total + size > budget:
|
||||
return run_id + 1, 1, size
|
||||
return run_id, count + 1, total + size
|
||||
|
||||
|
||||
def _value_in_texts(value: str, texts: Sequence[str]) -> str | None:
|
||||
"""Return the detected value as it appears in ``texts``, decoding it when it came from a JSON-escaped key line."""
|
||||
candidates: Final = (value, _json_unescaped(value))
|
||||
return next((candidate for candidate in candidates if any(candidate in text for text in texts)), None)
|
||||
|
||||
|
||||
class _ENTERPRISE_SecretDetection(CustomGuardrail):
|
||||
# Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
|
||||
# path skips should_run_check and never sees data["prompt"]).
|
||||
|
|
@ -539,11 +613,16 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
|
|||
|
||||
def __init__(self, detect_secrets_config: Optional[dict] = None, **kwargs):
|
||||
self.user_defined_detect_secrets_config = detect_secrets_config
|
||||
self._clean_scans: OrderedDict[bytes, None] = OrderedDict()
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def scan_message_for_secrets(self, message_content: str):
|
||||
def scan_message_for_secrets(self, message_content: str) -> list[dict[str, str]]:
|
||||
from detect_secrets.settings import transient_settings
|
||||
|
||||
digest: Final = hashlib.sha256(message_content.encode()).digest()
|
||||
if digest in self._clean_scans:
|
||||
self._clean_scans.move_to_end(digest)
|
||||
return []
|
||||
detect_secrets_config = (
|
||||
self.user_defined_detect_secrets_config or _default_detect_secrets_config
|
||||
)
|
||||
|
|
@ -552,6 +631,10 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
|
|||
(*message_content.splitlines(), *_quoted_assignments(message_content))
|
||||
)
|
||||
|
||||
if not found:
|
||||
self._clean_scans[digest] = None
|
||||
if len(self._clean_scans) > _CLEAN_SCAN_CACHE_SIZE:
|
||||
self._clean_scans.popitem(last=False)
|
||||
return [
|
||||
{"type": secret_type, "value": value}
|
||||
for value, secret_type in sorted(
|
||||
|
|
@ -580,6 +663,35 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
|
|||
)
|
||||
return pattern.sub("[REDACTED]", text)
|
||||
|
||||
def redact_tool_payload(self, payload: JsonValue) -> JsonValue:
|
||||
"""Redact secrets in a tool payload's string values, scanning each value alone and as its JSON lines."""
|
||||
fields: Final = tuple(iter_tool_json_fields(payload))
|
||||
texts: Final = tuple(text for _, text in fields)
|
||||
contexts: Final = {key: _key_context(key) for key in {key for key, _ in fields if key is not None}}
|
||||
keyed: Final = tuple(
|
||||
itertools.chain.from_iterable(
|
||||
_json_lines(None if key is None else contexts[key], text) for key, text in fields
|
||||
)
|
||||
)
|
||||
detected: Final = tuple(
|
||||
{**secret, "value": value}
|
||||
for secret in self.scan_message_for_secrets("\n".join((*texts, *keyed)))
|
||||
if (value := _value_in_texts(secret["value"], texts)) is not None
|
||||
)
|
||||
if not detected:
|
||||
return payload
|
||||
counts: Final = _masked_entity_count.get()
|
||||
if counts is not None:
|
||||
for secret in detected:
|
||||
counts[secret["type"]] = counts.get(secret["type"], 0) + 1
|
||||
verbose_proxy_logger.warning(
|
||||
"Detected and redacted secrets in tool payload: %s",
|
||||
sorted(dict.fromkeys(secret["type"] for secret in detected)),
|
||||
)
|
||||
values: Final = sorted({secret["value"] for secret in detected}, key=len, reverse=True)
|
||||
pattern: Final = re.compile("|".join(re.escape(value) for value in values))
|
||||
return map_tool_json_text(payload, lambda text: pattern.sub("[REDACTED]", text))
|
||||
|
||||
async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
if user_api_key_dict.permissions is not None:
|
||||
if GUARDRAIL_NAME in user_api_key_dict.permissions:
|
||||
|
|
@ -641,7 +753,7 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
|
|||
_masked_entity_count.set({})
|
||||
|
||||
# Covers multimodal list content + Responses-API input.
|
||||
inspected = walk_user_text(data, self.redact_text) + self._redact_prompt(data)
|
||||
inspected = walk_user_text(data, self.redact_text, self.redact_tool_payload) + self._redact_prompt(data)
|
||||
|
||||
if inspected == 0:
|
||||
# Image-only, empty-text, and unsupported payloads inspected
|
||||
|
|
|
|||
|
|
@ -8,9 +8,15 @@ skip the other shapes — these helpers normalise that so every hook sees
|
|||
every text fragment.
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import iter_items
|
||||
from litellm.responses.litellm_completion_transformation.compaction_items import (
|
||||
compaction_summary,
|
||||
|
|
@ -60,11 +66,157 @@ TEXT_PART_TYPES: Final[frozenset[str]] = frozenset(
|
|||
{"text", "input_text", "output_text", "summary_text", "reasoning_text"}
|
||||
)
|
||||
|
||||
_OBJECT_LIST: Final[TypeAdapter[Sequence[object]]] = TypeAdapter(Sequence[object])
|
||||
_TEXT_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _object_list(value: object) -> Sequence[object] | None:
|
||||
return _OBJECT_LIST.validate_python(value) if isinstance(value, list) else None
|
||||
|
||||
|
||||
def _text_object(value: object) -> Mapping[str, object] | None:
|
||||
if not isinstance(value, dict):
|
||||
return None
|
||||
try:
|
||||
return _TEXT_OBJECT.validate_python(value)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _tool_json_object(pairs: Sequence[tuple[str, JsonValue]]) -> JsonValue:
|
||||
value: Final = dict(pairs)
|
||||
if len(value) != len(pairs):
|
||||
raise ValueError("Duplicate tool argument keys")
|
||||
return value
|
||||
|
||||
|
||||
def _finite_number(literal: str) -> float:
|
||||
number: Final = float(literal)
|
||||
if not math.isfinite(number):
|
||||
raise ValueError("Non-finite tool argument number")
|
||||
return number
|
||||
|
||||
|
||||
def _tool_arguments(arguments: str) -> JsonValue:
|
||||
parsed: Final = _tool_json(
|
||||
json.loads(
|
||||
arguments,
|
||||
object_pairs_hook=_tool_json_object,
|
||||
parse_float=_finite_number,
|
||||
parse_constant=_finite_number,
|
||||
)
|
||||
)
|
||||
json.dumps(parsed, ensure_ascii=False).encode("utf-8") # raises for lone surrogates, which stay raw text
|
||||
return parsed
|
||||
|
||||
|
||||
_JSON_STRING_TOKEN: Final = re.compile(r'"[^"\\]*(?:\\.[^"\\]*)*"')
|
||||
|
||||
_KEY_SUFFIX: Final = re.compile(r"[ \t\n\r]*:")
|
||||
|
||||
|
||||
def _value_string_spans(arguments: str) -> Iterator[tuple[int, int]]:
|
||||
"""Yield the span of every non-empty string value in the valid JSON ``arguments``, in document order."""
|
||||
for token in _JSON_STRING_TOKEN.finditer(arguments):
|
||||
if token.end() - token.start() > 2 and _KEY_SUFFIX.match(arguments, token.end()) is None:
|
||||
yield token.span()
|
||||
|
||||
|
||||
def _respliced(arguments: str, parsed: JsonValue, rewritten: JsonValue) -> str:
|
||||
"""Serialize ``rewritten`` by replacing only its changed string values in ``arguments``, when nothing else changed."""
|
||||
serialized: Final = json.dumps(rewritten, ensure_ascii=False)
|
||||
strings: Final = iter(tuple(iter_tool_json_text(rewritten)))
|
||||
if json.dumps(map_tool_json_text(parsed, lambda text: next(strings, text)), ensure_ascii=False) != serialized:
|
||||
return serialized
|
||||
values: Final = zip(_value_string_spans(arguments), iter_tool_json_text(parsed), iter_tool_json_text(rewritten))
|
||||
changed: Final = tuple((span, new) for span, old, new in values if new != old)
|
||||
starts: Final = (0, *(end for (_, end), _ in changed))
|
||||
ends: Final = (*(start for (start, _), _ in changed), len(arguments))
|
||||
replacements: Final = (*(json.dumps(new, ensure_ascii=False) for _, new in changed), "")
|
||||
return "".join(arguments[start:end] + text for start, end, text in zip(starts, ends, replacements))
|
||||
|
||||
|
||||
def iter_tool_json_fields(value: object) -> Iterator[tuple[str | None, str]]:
|
||||
"""Yield ``(nearest key, text)`` for every non-empty string leaf; list items take their parent's key."""
|
||||
pending: Final[list[tuple[object, str | None]]] = [(value, None)] # mutable-ok: owned DFS stack avoids recursion
|
||||
while pending:
|
||||
current, key = pending.pop()
|
||||
if isinstance(current, str):
|
||||
if current:
|
||||
yield key, current
|
||||
continue
|
||||
if current is None or isinstance(current, (bool, int, float)):
|
||||
continue
|
||||
fields = _text_object(current)
|
||||
if fields is not None:
|
||||
pending.extend((child, child_key) for child_key, child in reversed(tuple(fields.items())))
|
||||
continue
|
||||
items = _object_list(current)
|
||||
if items is None:
|
||||
raise ValueError("Invalid tool JSON payload")
|
||||
pending.extend((child, key) for child in reversed(items))
|
||||
|
||||
|
||||
def iter_tool_json_text(value: object) -> Iterator[str]:
|
||||
for _, text in iter_tool_json_fields(value):
|
||||
yield text
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ToolFrame:
|
||||
source: object
|
||||
parent: dict[str, JsonValue] | list[JsonValue] # mutable-ok: frame writes only into the mapper's fresh output tree
|
||||
key: str | int
|
||||
|
||||
def write(self, value: JsonValue) -> None:
|
||||
if isinstance(self.parent, dict) and isinstance(self.key, str):
|
||||
self.parent[self.key] = value
|
||||
return
|
||||
if isinstance(self.parent, list) and isinstance(self.key, int):
|
||||
self.parent[self.key] = value
|
||||
return
|
||||
raise ValueError("Invalid tool JSON payload")
|
||||
|
||||
|
||||
def map_tool_json_text(value: object, visit: Callable[[str], str]) -> JsonValue:
|
||||
root: Final[list[JsonValue]] = [None] # mutable-ok: one owned root slot receives the iterative result
|
||||
pending: Final[list[_ToolFrame]] = [_ToolFrame(value, root, 0)] # mutable-ok: owned traversal stack
|
||||
while pending:
|
||||
frame = pending.pop()
|
||||
current = frame.source
|
||||
if isinstance(current, str):
|
||||
frame.write(visit(current) if current else current)
|
||||
continue
|
||||
if current is None or isinstance(current, (bool, int, float)):
|
||||
frame.write(current)
|
||||
continue
|
||||
fields = _text_object(current)
|
||||
if fields is not None:
|
||||
mapped: dict[str, JsonValue] = {key: None for key in fields} # mutable-ok: owned output slots
|
||||
frame.write(mapped)
|
||||
pending.extend(_ToolFrame(child, mapped, key) for key, child in reversed(tuple(fields.items())))
|
||||
continue
|
||||
children = _object_list(current)
|
||||
if children is None:
|
||||
raise ValueError("Invalid tool JSON payload")
|
||||
items: list[JsonValue] = [None] * len(children) # mutable-ok: fresh indexed output slots
|
||||
frame.write(items)
|
||||
pending.extend(_ToolFrame(child, items, index) for index, child in reversed(tuple(enumerate(children))))
|
||||
return root[0]
|
||||
|
||||
|
||||
def _tool_json(value: object) -> JsonValue:
|
||||
return map_tool_json_text(value, lambda text: text)
|
||||
|
||||
|
||||
# Responses-API item types whose ``output`` field carries user/tool text
|
||||
# that guardrails should inspect. ``function_call_output`` is the
|
||||
# built-in shape; ``custom_tool_call_output`` is the custom-tool
|
||||
# counterpart (see ``ChatCompletionCustomToolCallOutput``).
|
||||
_OUTPUT_ITEM_TYPES: Final[frozenset[str]] = frozenset({"function_call_output", "custom_tool_call_output"})
|
||||
# Responses input items carrying the model's tool call: ``function_call`` holds
|
||||
# JSON ``arguments`` and ``custom_tool_call`` holds a freeform ``input`` string.
|
||||
_CALL_ITEM_FIELDS: Final[Mapping[str, str]] = {"function_call": "arguments", "custom_tool_call": "input"}
|
||||
|
||||
|
||||
def _part_text(part: Mapping[str, object]) -> str | None:
|
||||
|
|
@ -153,16 +305,65 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]:
|
|||
yield from _iter_text_parts_in_content(message.get("content"))
|
||||
|
||||
|
||||
def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int:
|
||||
def _call_item(item: object) -> tuple[str, object] | None:
|
||||
"""Return ``(field, value)`` for a Responses tool-call input item, else None."""
|
||||
fields: Final = _text_object(item)
|
||||
field: Final = _CALL_ITEM_FIELDS.get(str(fields.get("type"))) if fields is not None else None
|
||||
return (field, fields.get(field)) if fields is not None and field is not None and field in fields else None
|
||||
|
||||
|
||||
def walk_user_text(
|
||||
data: dict[str, object],
|
||||
visit: Callable[[str], str],
|
||||
visit_tool: Callable[[JsonValue], JsonValue] | None = None,
|
||||
) -> int:
|
||||
"""Rewrite every text fragment in place via ``visit``.
|
||||
|
||||
Mutates ``data["messages"]`` and ``data["input"]``. Returns the number
|
||||
Mutates ``data["messages"]`` and ``data["input"]`` (including tool payloads), ``data["system"]`` and
|
||||
``data["instructions"]``. Returns the number
|
||||
of fragments visited so callers can short-circuit when nothing was
|
||||
inspected.
|
||||
"""
|
||||
visited = 0
|
||||
|
||||
def _rewrite_content(content: object) -> object:
|
||||
def _rewrite_tool(value: object) -> JsonValue:
|
||||
nonlocal visited
|
||||
payload: Final = _tool_json(value)
|
||||
if not tuple(iter_tool_json_text(payload)):
|
||||
return payload
|
||||
visited += 1 # rebind-ok: tally the structured payload inspected by this callback
|
||||
return visit_tool(payload) if visit_tool is not None else map_tool_json_text(payload, visit)
|
||||
|
||||
def _rewrite_arguments(arguments: object) -> object:
|
||||
nonlocal visited
|
||||
if isinstance(arguments, dict):
|
||||
return _rewrite_tool(arguments)
|
||||
if not isinstance(arguments, str) or not arguments:
|
||||
return arguments
|
||||
try:
|
||||
parsed: Final = _tool_arguments(arguments)
|
||||
except (ValueError, RecursionError):
|
||||
visited += 1 # rebind-ok: malformed arguments still undergo the text callback
|
||||
return visit(arguments)
|
||||
if not isinstance(parsed, (dict, list)):
|
||||
visited += 1 # rebind-ok: scalar arguments still undergo the text callback
|
||||
return visit(arguments)
|
||||
rewritten: Final = _rewrite_tool(parsed)
|
||||
return arguments if rewritten == parsed else _respliced(arguments, parsed, rewritten)
|
||||
|
||||
def _rewrite_part(part: Mapping[str, object], native: bool) -> Mapping[str, object]:
|
||||
nonlocal visited
|
||||
text: Final = _part_text(part)
|
||||
rewritten: Final = {**part, **({"text": visit(text)} if text is not None else {})}
|
||||
if text is not None:
|
||||
visited += 1 # rebind-ok: tally each existing text field alongside the native payload
|
||||
if native and part.get("type") == "tool_result" and "content" in part:
|
||||
return {**rewritten, "content": _rewrite_content(part["content"], native=False)}
|
||||
if native and part.get("type") == "tool_use" and "input" in part:
|
||||
return {**rewritten, "input": _rewrite_tool(part["input"])}
|
||||
return rewritten
|
||||
|
||||
def _rewrite_content(content: object, native: bool = True) -> object:
|
||||
nonlocal visited
|
||||
if isinstance(content, str):
|
||||
if content:
|
||||
|
|
@ -175,9 +376,8 @@ def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int:
|
|||
if isinstance(part, str) and part:
|
||||
visited += 1
|
||||
new_parts.append(visit(part))
|
||||
elif isinstance(part, dict) and _part_text(part) is not None:
|
||||
visited += 1
|
||||
new_parts.append({**part, "text": visit(part["text"])})
|
||||
elif isinstance(part, dict):
|
||||
new_parts.append(_rewrite_part(part, native))
|
||||
else:
|
||||
new_parts.append(part)
|
||||
return new_parts
|
||||
|
|
@ -186,8 +386,21 @@ def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int:
|
|||
messages: Final = data.get("messages")
|
||||
if isinstance(messages, list):
|
||||
for message in messages:
|
||||
if isinstance(message, dict) and "content" in message:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
if "content" in message:
|
||||
message["content"] = _rewrite_content(message["content"])
|
||||
calls = message.get("tool_calls")
|
||||
if isinstance(calls, list):
|
||||
for call in calls:
|
||||
function = call.get("function") if isinstance(call, dict) else None
|
||||
if isinstance(function, dict) and "arguments" in function:
|
||||
function["arguments"] = _rewrite_arguments(function["arguments"])
|
||||
|
||||
if "system" in data:
|
||||
data["system"] = _rewrite_content(data["system"]) # rebind-ok: preserve the native hook's in-place contract
|
||||
if "instructions" in data:
|
||||
data["instructions"] = _rewrite_content(data["instructions"]) # rebind-ok: Responses system prompt, in place
|
||||
|
||||
input_value: Final = data.get("input")
|
||||
if isinstance(input_value, str):
|
||||
|
|
@ -217,6 +430,8 @@ def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int:
|
|||
item["content"] = _rewrite_content(item["content"])
|
||||
elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item:
|
||||
item["output"] = _rewrite_content(item["output"])
|
||||
elif (call := _call_item(item)) is not None:
|
||||
item[call[0]] = _rewrite_arguments(call[1])
|
||||
return visited
|
||||
|
||||
return visited
|
||||
|
|
|
|||
|
|
@ -0,0 +1,520 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.budget import case_timeout_seconds
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request
|
||||
from integration.security.test_hide_secrets_tool_traffic import (
|
||||
SECRET,
|
||||
TEXT,
|
||||
Rig,
|
||||
_input,
|
||||
_messages,
|
||||
_provider,
|
||||
tool_rig,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
_MARKER: Final = re.compile(r"case_[0-9a-f]{32}")
|
||||
|
||||
|
||||
def _identity() -> str:
|
||||
return f"case_{uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def _peer(request: Request) -> Reply:
|
||||
if request.method != "POST":
|
||||
return Reply()
|
||||
body: Final = json.loads(request.body)
|
||||
marker: Final = _MARKER.search(request.body.decode())
|
||||
assert marker is not None, request.body
|
||||
identity: Final = marker.group()
|
||||
if request.target.endswith("/messages") and (
|
||||
not isinstance(body.get("messages"), list)
|
||||
or any(not isinstance(message, dict) for message in body.get("messages", []))
|
||||
or ("system" in body and not isinstance(body["system"], (str, list)))
|
||||
):
|
||||
return Reply(
|
||||
status=400,
|
||||
body=b'{"type":"error","error":{"type":"invalid_request_error","message":"invalid shape"}}',
|
||||
)
|
||||
if "provider-error" in request.body.decode():
|
||||
return Reply(
|
||||
status=401,
|
||||
body=b'{"type":"error","error":{"message":"scripted deployment refused","type":"authentication_error"}}',
|
||||
)
|
||||
if request.target.endswith("/responses"):
|
||||
answer: Final = {
|
||||
"id": f"resp_{identity}",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": body["model"],
|
||||
"error": None,
|
||||
"incomplete_details": None,
|
||||
"output": [
|
||||
{
|
||||
"id": f"msg_{identity}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "accepted", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
|
||||
}
|
||||
return Reply(body=json.dumps(answer).encode())
|
||||
if body.get("stream") and not request.target.endswith("/messages"):
|
||||
chunk: Final = {
|
||||
"id": f"chatcmpl_{identity}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": body["model"],
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "accepted"}, "finish_reason": None}],
|
||||
}
|
||||
finish: Final = {
|
||||
**chunk,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
|
||||
}
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
f"data: {json.dumps(chunk)}\n\n".encode(),
|
||||
f"data: {json.dumps(finish)}\n\n".encode(),
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
)
|
||||
reply: Final = _provider(request)
|
||||
if reply.chunks is not None:
|
||||
return replace(
|
||||
reply, chunks=tuple(re.sub(rb"msg_[a-z0-9_]+", f"msg_{identity}".encode(), chunk) for chunk in reply.chunks)
|
||||
)
|
||||
answer_body: Final = json.loads(reply.body)
|
||||
return replace(
|
||||
reply,
|
||||
body=json.dumps(
|
||||
{**answer_body, "id": f"{'msg' if request.target.endswith('/messages') else 'chatcmpl'}_{identity}"}
|
||||
).encode(),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Caller:
|
||||
rig: Rig
|
||||
key: str
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _key(rig: Rig, *, permissions: dict[str, JsonValue] | None = None) -> Iterator[str]:
|
||||
created: Final = rig.proxy.post(
|
||||
"/key/generate", {"models": ["*"], **({"permissions": permissions} if permissions else {})}
|
||||
)
|
||||
key: Final = created["key"]
|
||||
assert isinstance(key, str), created
|
||||
try:
|
||||
yield key
|
||||
finally:
|
||||
rig.proxy.post("/key/delete", {"keys": [key]})
|
||||
assert (
|
||||
read_rows(
|
||||
'SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
database_url=rig.database_url,
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def caller(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Caller]:
|
||||
with tool_rig(tmp_path_factory.mktemp("hide-secrets-compat"), _peer) as rig:
|
||||
with _key(rig) as key:
|
||||
yield Caller(rig, key)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_caller(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Caller]:
|
||||
with tool_rig(tmp_path_factory.mktemp("hide-secrets-worker-recovery"), _peer) as rig:
|
||||
with _key(rig) as key:
|
||||
yield Caller(rig, key)
|
||||
|
||||
|
||||
def _body(provider: str, marker: str, *, stream: bool = False, secret: bool = True) -> dict[str, JsonValue]:
|
||||
if provider == "anthropic":
|
||||
body: Final = _messages("result-nested")
|
||||
body["messages"][0]["content"] = f"read .env {marker}"
|
||||
return {**body, "stream": stream, "cache": {"no-cache": True}}
|
||||
return {
|
||||
"model": "gpt",
|
||||
"stream": stream,
|
||||
"cache": {"no-cache": True},
|
||||
"messages": [
|
||||
{"role": "user", "content": f"read .env {marker}"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"arguments": json.dumps(_input() if secret else {"path": ".env"}),
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "done"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _wire(caller: Caller, marker: str) -> dict[str, JsonValue]:
|
||||
requests: Final = caller.rig.wire.drain()
|
||||
selected: Final = tuple(request for request in requests if marker.encode() in request.body)
|
||||
assert len(selected) == 1, [(request.target, request.body) for request in requests]
|
||||
assert SECRET.encode() not in selected[0].body, selected[0].body
|
||||
assert b"REDACTED" in selected[0].body, selected[0].body
|
||||
return json.loads(selected[0].body)
|
||||
|
||||
|
||||
def _stored(caller: Caller, response_id: str, *, masked: bool = True, cache: bool = False) -> dict[str, JsonValue]:
|
||||
query: Final = (
|
||||
'SELECT request_id, status, metadata, cache_hit FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE request_id=%s OR litellm_call_id=%s"
|
||||
if not cache
|
||||
else 'SELECT request_id, status, metadata, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s'
|
||||
)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
query,
|
||||
(response_id, response_id) if not cache else (f"{response_id}_cache_hit%",),
|
||||
database_url=caller.rig.database_url,
|
||||
),
|
||||
lambda observed: len(observed) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
information: Final = rows[0]["metadata"].get("guardrail_information") or []
|
||||
matching: Final = [item for item in information if item.get("guardrail_name") == "hide-secrets"]
|
||||
if masked:
|
||||
assert len(matching) == 1, rows[0]
|
||||
assert matching[0]["guardrail_status"] == "success", rows[0]
|
||||
assert sum((matching[0].get("masked_entity_count") or {}).values()) > 0, rows[0]
|
||||
else:
|
||||
assert not any(item.get("masked_entity_count") for item in matching), rows[0]
|
||||
if cache:
|
||||
assert str(rows[0]["cache_hit"]).lower() == "true", rows[0]
|
||||
return rows[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ("anthropic", "openai"))
|
||||
@pytest.mark.parametrize("client_mode", ("sync", "async"))
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
|
||||
@pytest.mark.asyncio
|
||||
async def test_sdk_tool_traffic_is_redacted(caller: Caller, provider: str, client_mode: str, stream: bool) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = _body(provider, marker, stream=stream)
|
||||
options: Final = {"extra_headers": {"x-litellm-call-id": marker}, "extra_body": {"cache": body["cache"]}}
|
||||
if provider == "anthropic":
|
||||
arguments: Final = {
|
||||
"model": "claude",
|
||||
"max_tokens": 32,
|
||||
"messages": body["messages"],
|
||||
"stream": stream,
|
||||
**options,
|
||||
}
|
||||
if client_mode == "sync":
|
||||
with anthropic.Anthropic(
|
||||
api_key=caller.key,
|
||||
base_url=str(caller.rig.proxy.client.base_url),
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(trust_env=False),
|
||||
) as sdk:
|
||||
response = sdk.messages.create(**arguments)
|
||||
text = "".join(event.model_dump_json() for event in response) if stream else response.model_dump_json()
|
||||
else:
|
||||
async with anthropic.AsyncAnthropic(
|
||||
api_key=caller.key,
|
||||
base_url=str(caller.rig.proxy.client.base_url),
|
||||
max_retries=0,
|
||||
http_client=httpx.AsyncClient(trust_env=False),
|
||||
) as sdk:
|
||||
response = await sdk.messages.create(**arguments)
|
||||
text = (
|
||||
"".join([event.model_dump_json() async for event in response])
|
||||
if stream
|
||||
else response.model_dump_json()
|
||||
)
|
||||
response_id = f"msg_{marker}"
|
||||
else:
|
||||
arguments = {"model": "gpt", "messages": body["messages"], "stream": stream, **options}
|
||||
if client_mode == "sync":
|
||||
with openai.OpenAI(
|
||||
api_key=caller.key,
|
||||
base_url=str(caller.rig.proxy.client.base_url).rstrip("/") + "/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(trust_env=False),
|
||||
) as sdk:
|
||||
response = sdk.chat.completions.create(**arguments)
|
||||
text = "".join(event.model_dump_json() for event in response) if stream else response.model_dump_json()
|
||||
else:
|
||||
async with openai.AsyncOpenAI(
|
||||
api_key=caller.key,
|
||||
base_url=str(caller.rig.proxy.client.base_url).rstrip("/") + "/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.AsyncClient(trust_env=False),
|
||||
) as sdk:
|
||||
response = await sdk.chat.completions.create(**arguments)
|
||||
text = (
|
||||
"".join([event.model_dump_json() async for event in response])
|
||||
if stream
|
||||
else response.model_dump_json()
|
||||
)
|
||||
response_id = f"chatcmpl_{marker}"
|
||||
assert "accepted" in text and response_id in text, text
|
||||
_wire(caller, marker)
|
||||
_stored(caller, response_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ("/v1/messages", "/anthropic/v1/messages"), ids=("messages", "anthropic-alias"))
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
|
||||
def test_messages_aliases_redact_system_and_multiple_tool_items(caller: Caller, path: str, stream: bool) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = _body("anthropic", marker, stream=stream)
|
||||
if path.startswith("/anthropic/"):
|
||||
body["model"] = "claude-haiku-5-5"
|
||||
body["system"] = [{"type": "text", "text": TEXT, "cache_control": {"type": "ephemeral"}}]
|
||||
body["messages"][-1]["content"][0]["content"].append({"type": "text", "text": TEXT})
|
||||
response: Final = caller.rig.proxy.request(
|
||||
"POST", path, body, key=caller.key, headers={"x-litellm-call-id": marker}
|
||||
)
|
||||
assert response.status_code == 200 and "accepted" in response.text, response.text
|
||||
observed: Final = _wire(caller, marker)
|
||||
assert observed["system"][0]["cache_control"] == {"type": "ephemeral"}, observed
|
||||
assert len(observed["messages"][-1]["content"][0]["content"]) == 3, observed
|
||||
_stored(caller, f"msg_{marker}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ("valid", "unchanged", "malformed", "duplicate"))
|
||||
def test_tool_argument_compatibility(caller: Caller, shape: str) -> None:
|
||||
marker: Final = _identity()
|
||||
original: Final = {
|
||||
"valid": json.dumps(_input()),
|
||||
"unchanged": '{ "path" : "C:\\\\tmp", "count": 3, "ok": false }',
|
||||
"malformed": '{"env":"' + TEXT + '",broken',
|
||||
"duplicate": '{"env":"' + TEXT + '","env":"safe"}',
|
||||
}[shape]
|
||||
body: Final = _body("openai", marker)
|
||||
body["messages"][1]["tool_calls"][0]["function"]["arguments"] = original
|
||||
response: Final = caller.rig.proxy.request("POST", "/v1/chat/completions", body, key=caller.key)
|
||||
assert response.status_code == 200 and "accepted" in response.text, response.text
|
||||
selected: Final = tuple(request for request in caller.rig.wire.drain() if marker.encode() in request.body)
|
||||
assert len(selected) == 1, selected
|
||||
observed: Final = json.loads(selected[0].body)["messages"][1]["tool_calls"][0]["function"]["arguments"]
|
||||
if shape == "unchanged":
|
||||
assert observed == original, observed
|
||||
else:
|
||||
assert SECRET not in observed and "REDACTED" in observed, observed
|
||||
if shape == "valid":
|
||||
parsed: Final = json.loads(observed)
|
||||
assert parsed["path"] == _input()["path"] and parsed["count"] == 3, parsed
|
||||
assert parsed["enabled"] is False and parsed["optional"] is None, parsed
|
||||
_stored(caller, response.json()["id"], masked=shape != "unchanged")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("original", "masked"),
|
||||
(
|
||||
(
|
||||
'{ "env" : "' + TEXT + '", "amount": 0.123456789012345678901, "price": 1.50, "limit": 1e5 }',
|
||||
'{ "env" : "AWS_ACCESS_KEY_ID=[REDACTED]", "amount": 0.123456789012345678901, "price": 1.50, "limit": 1e5 }',
|
||||
),
|
||||
('{"password": "Zx9qL2mN8vB4\\"kR7Tq", "path": ".env"}', '{"password": "[REDACTED]\\"kR7Tq", "path": ".env"}'),
|
||||
),
|
||||
ids=("number-literals", "keyword-match-inside-escape"),
|
||||
)
|
||||
def test_masked_tool_arguments_keep_their_other_text(caller: Caller, original: str, masked: str) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = _body("openai", marker)
|
||||
body["messages"][1]["tool_calls"][0]["function"]["arguments"] = original
|
||||
response: Final = caller.rig.proxy.request("POST", "/v1/chat/completions", body, key=caller.key)
|
||||
assert response.status_code == 200 and "accepted" in response.text, response.text
|
||||
selected: Final = tuple(request for request in caller.rig.wire.drain() if marker.encode() in request.body)
|
||||
assert len(selected) == 1, selected
|
||||
observed: Final = json.loads(selected[0].body)["messages"][1]["tool_calls"][0]["function"]["arguments"]
|
||||
assert observed == masked, observed
|
||||
_stored(caller, response.json()["id"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ("/v1/responses", "/v1/completions"), ids=("responses", "plain-completion"))
|
||||
def test_existing_text_shapes_still_redact(caller: Caller, path: str) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = {
|
||||
"model": "gpt",
|
||||
"cache": {"no-cache": True},
|
||||
**({"input": f"{marker} {TEXT}"} if path.endswith("responses") else {"prompt": f"{marker} {TEXT}"}),
|
||||
}
|
||||
response: Final = caller.rig.proxy.request(
|
||||
"POST", path, body, key=caller.key, headers={"x-litellm-call-id": marker}
|
||||
)
|
||||
assert response.status_code == 200 and "accepted" in response.text, response.text
|
||||
_wire(caller, marker)
|
||||
request_id: Final = marker if path.endswith("responses") else response.json()["id"]
|
||||
_stored(caller, request_id)
|
||||
|
||||
|
||||
_RESPONSES_ITEMS: Final[dict[str, tuple[JsonValue, ...]]] = {
|
||||
"instructions": (),
|
||||
"function-call": (
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "call_1",
|
||||
"name": "write_file",
|
||||
"arguments": json.dumps({"path": ".env", "content": TEXT}),
|
||||
},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "written"},
|
||||
),
|
||||
"custom-tool-call": (
|
||||
{"type": "custom_tool_call", "call_id": "call_1", "name": "shell", "input": f"export {TEXT}"},
|
||||
{"type": "custom_tool_call_output", "call_id": "call_1", "output": "done"},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", tuple(_RESPONSES_ITEMS))
|
||||
def test_responses_tool_calls_and_instructions_are_redacted(caller: Caller, shape: str) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = {
|
||||
"model": "gpt",
|
||||
"cache": {"no-cache": True},
|
||||
"input": [{"role": "user", "content": f"{marker} write the env file"}, *_RESPONSES_ITEMS[shape]],
|
||||
**({"instructions": TEXT} if shape == "instructions" else {}),
|
||||
}
|
||||
response: Final = caller.rig.proxy.request(
|
||||
"POST", "/v1/responses", body, key=caller.key, headers={"x-litellm-call-id": marker}
|
||||
)
|
||||
assert response.status_code == 200 and "accepted" in response.text, response.text
|
||||
sent: Final = _wire(caller, marker)
|
||||
if shape == "function-call":
|
||||
assert json.loads(sent["input"][1]["arguments"]) == {"path": ".env", "content": "AWS_ACCESS_KEY_ID=[REDACTED]"}
|
||||
if shape == "custom-tool-call":
|
||||
assert sent["input"][1]["input"] == "export AWS_ACCESS_KEY_ID=[REDACTED]"
|
||||
_stored(caller, marker)
|
||||
|
||||
|
||||
def test_cache_hit_keeps_masked_guardrail_metadata(caller: Caller) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = {**_body("openai", marker), "cache": {"ttl": 60}}
|
||||
first: Final = caller.rig.proxy.request("POST", "/v1/chat/completions", body, key=caller.key)
|
||||
assert first.status_code == 200, first.text
|
||||
_wire(caller, marker)
|
||||
_stored(caller, first.json()["id"])
|
||||
second: Final = caller.rig.proxy.request("POST", "/v1/chat/completions", body, key=caller.key)
|
||||
assert second.status_code == 200 and second.json()["id"] == first.json()["id"], second.text
|
||||
assert not tuple(request for request in caller.rig.wire.drain() if marker.encode() in request.body)
|
||||
_stored(caller, first.json()["id"], cache=True)
|
||||
|
||||
|
||||
def test_key_permission_can_disable_default_guardrail(caller: Caller) -> None:
|
||||
marker: Final = _identity()
|
||||
with _key(caller.rig, permissions={"hide_secrets": False}) as key:
|
||||
response: Final = caller.rig.proxy.request("POST", "/v1/chat/completions", _body("openai", marker), key=key)
|
||||
assert response.status_code == 200, response.text
|
||||
selected: Final = tuple(request for request in caller.rig.wire.drain() if marker.encode() in request.body)
|
||||
assert len(selected) == 1 and SECRET.encode() in selected[0].body, selected
|
||||
_stored(caller, response.json()["id"], masked=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ("anthropic", "openai"))
|
||||
def test_provider_error_occurs_after_redaction(caller: Caller, provider: str) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = _body(provider, marker)
|
||||
body["messages"][0]["content"] += " provider-error"
|
||||
path: Final = "/v1/messages" if provider == "anthropic" else "/v1/chat/completions"
|
||||
response: Final = caller.rig.proxy.request(
|
||||
"POST", path, body, key=caller.key, headers={"x-litellm-call-id": marker}
|
||||
)
|
||||
assert response.status_code == 401 and "scripted deployment refused" in response.text, response.text
|
||||
_wire(caller, marker)
|
||||
row: Final = _stored(caller, response.headers["x-litellm-call-id"])
|
||||
assert row["status"] == "failure", row
|
||||
|
||||
|
||||
@pytest.mark.parametrize("problem", ("unauthenticated", "messages-int", "messages-list", "system-int"))
|
||||
def test_invalid_requests_preserve_redaction_and_proxy_health(caller: Caller, problem: str) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = _body("anthropic", marker)
|
||||
if problem == "messages-int":
|
||||
body["messages"] = 17
|
||||
elif problem == "messages-list":
|
||||
body["messages"] = [17]
|
||||
elif problem == "system-int":
|
||||
body["system"] = 17
|
||||
body["metadata"] = {"user_id": marker}
|
||||
response: Final = caller.rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
body,
|
||||
key="sk-deliberately-unknown" if problem == "unauthenticated" else caller.key,
|
||||
)
|
||||
assert response.status_code >= 400 and "error" in response.text.lower(), response.text
|
||||
selected: Final = tuple(request for request in caller.rig.wire.drain() if marker.encode() in request.body)
|
||||
if problem == "unauthenticated":
|
||||
assert not selected, selected
|
||||
else:
|
||||
assert all(SECRET.encode() not in request.body for request in selected), selected
|
||||
health: Final = caller.rig.proxy.request("GET", "/health/liveliness")
|
||||
assert health.status_code == 200, health.text
|
||||
|
||||
|
||||
def test_surviving_worker_redacts_mixed_burst_and_replacement_recovers(isolated_caller: Caller) -> None:
|
||||
caller: Final = isolated_caller
|
||||
workers: Final = caller.rig.owned.worker_pids()
|
||||
assert len(workers) == 2, workers
|
||||
markers: Final = tuple(_identity() for _ in range(20))
|
||||
|
||||
def send(item: tuple[int, str]) -> tuple[str, str]:
|
||||
index, marker = item
|
||||
provider: Final = "anthropic" if index % 2 else "openai"
|
||||
path: Final = "/v1/messages" if provider == "anthropic" else "/v1/chat/completions"
|
||||
proxy: Final = caller.rig.proxy
|
||||
# a new connection per call: a pooled keep-alive socket to the killed worker can outlive the kill
|
||||
with httpx.Client(base_url=proxy.client.base_url, timeout=case_timeout_seconds(), trust_env=False) as client:
|
||||
response: Final = Gateway(client, proxy.key, proxy.upstream_url).request(
|
||||
"POST", path, _body(provider, marker, stream=index % 3 == 0), key=caller.key
|
||||
)
|
||||
assert response.status_code == 200 and "accepted" in response.text, response.text
|
||||
return marker, f"{'msg' if provider == 'anthropic' else 'chatcmpl'}_{marker}"
|
||||
|
||||
with caller.rig.owned.replacing_worker(min(workers)) as kill:
|
||||
kill()
|
||||
with ThreadPoolExecutor(max_workers=20) as pool:
|
||||
results: Final = tuple(pool.map(send, enumerate(markers)))
|
||||
captured: Final = caller.rig.wire.drain()
|
||||
for marker, response_id in results:
|
||||
selected: Final = tuple(request for request in captured if marker.encode() in request.body)
|
||||
assert len(selected) == 1 and SECRET.encode() not in selected[0].body, selected
|
||||
_stored(caller, response_id)
|
||||
assert len(caller.rig.owned.worker_pids()) == 2 and min(workers) not in caller.rig.owned.worker_pids()
|
||||
recovery: Final = _identity()
|
||||
response: Final = caller.rig.proxy.request(
|
||||
"POST", "/v1/chat/completions", _body("openai", recovery), key=caller.key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
_wire(caller, recovery)
|
||||
_stored(caller, response.json()["id"])
|
||||
415
tests/integration/security/test_hide_secrets_tool_traffic.py
Normal file
415
tests/integration/security/test_hide_secrets_tool_traffic.py
Normal file
|
|
@ -0,0 +1,415 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import string
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.budget import case_timeout_seconds
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.database import scratch_database
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.timing import assert_not_delayed_by, assert_within_factor_of_control, timed
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
SECRET: Final = "AKIAIOSFODNN7EXAMPLE"
|
||||
TEXT: Final = f"AWS_ACCESS_KEY_ID={SECRET}"
|
||||
MASTER_KEY: Final = "sk-integration-hide-secrets-7426-synthetic-master"
|
||||
|
||||
|
||||
def _provider(request: Request) -> Reply:
|
||||
if request.method != "POST":
|
||||
return Reply()
|
||||
body: Final = json.loads(request.body)
|
||||
identity: Final = uuid.uuid4().hex
|
||||
if request.target.endswith("/messages"):
|
||||
response: Final = {
|
||||
"id": f"msg_{identity}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": body["model"],
|
||||
"content": [{"type": "text", "text": "accepted"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 1},
|
||||
}
|
||||
if body.get("stream"):
|
||||
events: Final = (
|
||||
("message_start", {"type": "message_start", "message": {**response, "content": []}}),
|
||||
(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "accepted"}},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
(
|
||||
"message_delta",
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}},
|
||||
),
|
||||
("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"event: {name}\ndata: {json.dumps(value)}\n\n".encode() for name, value in events),
|
||||
)
|
||||
return Reply(body=json.dumps(response).encode())
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl_{identity}",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": body["model"],
|
||||
"choices": [
|
||||
{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "accepted"}}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
wire: Wire
|
||||
database_url: str
|
||||
owned: OwnedProxy
|
||||
environment: Mapping[str, str]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def tool_rig(
|
||||
directory: Path,
|
||||
respond: Callable[[Request], Reply] = _provider,
|
||||
*,
|
||||
configure: Callable[[dict[str, JsonValue]], None] | None = None,
|
||||
environment: Mapping[str, str] | None = None,
|
||||
workers: int = 2,
|
||||
) -> Iterator[Rig]:
|
||||
assert os.environ.get("LITELLM_LICENSE"), "LITELLM_LICENSE must be set: hide-secrets is an enterprise feature"
|
||||
product_root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
|
||||
with (
|
||||
wire_server(respond) as wire,
|
||||
scratch_database("hide_secrets") as database,
|
||||
):
|
||||
config: Final = directory / "config.yaml"
|
||||
configuration: Final = {
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": "claude",
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-5",
|
||||
"api_key": "synthetic-provider-key",
|
||||
"api_base": wire.url,
|
||||
"num_retries": 0,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5",
|
||||
"api_key": "synthetic-provider-key",
|
||||
"api_base": wire.url,
|
||||
"num_retries": 0,
|
||||
},
|
||||
},
|
||||
],
|
||||
"general_settings": {
|
||||
"master_key": MASTER_KEY,
|
||||
"store_prompts_in_spend_logs": True,
|
||||
"disable_spend_logs": False,
|
||||
"proxy_batch_write_at": 1,
|
||||
"proxy_batch_polling_interval": 1,
|
||||
},
|
||||
"router_settings": {"num_retries": 0},
|
||||
"litellm_settings": {
|
||||
"cache": True,
|
||||
"cache_params": {
|
||||
"type": "redis",
|
||||
"host": "os.environ/REDIS_HOST",
|
||||
"port": "os.environ/REDIS_PORT",
|
||||
},
|
||||
},
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": "hide-secrets",
|
||||
"litellm_params": {"guardrail": "hide-secrets", "mode": "pre_call", "default_on": True},
|
||||
}
|
||||
],
|
||||
}
|
||||
if configure is not None:
|
||||
configure(configuration)
|
||||
config.write_text(yaml.safe_dump(configuration))
|
||||
child_environment: Final = {
|
||||
"DATABASE_URL": database,
|
||||
"LITELLM_DISABLE_NO_REDIS_WARNING": "true",
|
||||
"ANTHROPIC_API_BASE": wire.url,
|
||||
"ANTHROPIC_API_KEY": "synthetic-provider-key",
|
||||
"LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "True",
|
||||
"PYTHONPATH": os.pathsep.join(
|
||||
(str(product_root / "enterprise"), str(product_root), os.environ.get("PYTHONPATH", ""))
|
||||
),
|
||||
**(environment or {}),
|
||||
}
|
||||
with httpx.Client() as client:
|
||||
seed: Final = Gateway(client, MASTER_KEY, wire.url)
|
||||
with owned_proxy_process(
|
||||
seed,
|
||||
directory,
|
||||
child_environment,
|
||||
config=config,
|
||||
workers=workers,
|
||||
client_timeout=case_timeout_seconds(),
|
||||
) as owned:
|
||||
yield Rig(owned.gateway, wire, database, owned, child_environment)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with tool_rig(tmp_path_factory.mktemp("hide-secrets-tools")) as value:
|
||||
yield value
|
||||
|
||||
|
||||
def _input() -> dict[str, JsonValue]:
|
||||
return {
|
||||
"env": TEXT,
|
||||
"path": 'C:\\tmp\\"config"',
|
||||
"count": 3,
|
||||
"enabled": False,
|
||||
"optional": None,
|
||||
"nested": ["safe", {"value": TEXT}],
|
||||
}
|
||||
|
||||
|
||||
def _messages(shape: str) -> dict[str, JsonValue]:
|
||||
user: Final[dict[str, JsonValue]] = {"role": "user", "content": "read .env"}
|
||||
tool: Final[dict[str, JsonValue]] = {
|
||||
"type": "tool_use",
|
||||
"id": "tool_1",
|
||||
"name": "read_file",
|
||||
"input": {"path": ".env"},
|
||||
}
|
||||
prefix: Final[list[JsonValue]] = [user, {"role": "assistant", "content": [tool]}]
|
||||
if shape == "plain":
|
||||
return {"model": "claude", "max_tokens": 32, "messages": [{"role": "user", "content": TEXT}]}
|
||||
if shape.startswith("result"):
|
||||
content: Final[JsonValue] = (
|
||||
TEXT
|
||||
if shape == "result-string"
|
||||
else [
|
||||
{"type": "text", "text": TEXT},
|
||||
{"type": "text", "text": "keep me"},
|
||||
]
|
||||
)
|
||||
return {
|
||||
"model": "claude",
|
||||
"max_tokens": 32,
|
||||
"messages": [
|
||||
*prefix,
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tool_1",
|
||||
"content": content,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
{"type": "text", "text": "summarize"},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
if shape.startswith("input"):
|
||||
arguments: Final = (
|
||||
_input()
|
||||
if shape == "input-dict"
|
||||
else {"command": "export API_TOKEN=0123456789abcdef0123456789abcdef01234567", "enabled": False}
|
||||
)
|
||||
return {
|
||||
"model": "claude",
|
||||
"max_tokens": 32,
|
||||
"messages": [
|
||||
user,
|
||||
{"role": "assistant", "content": [{**tool, "input": arguments}]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tool_1", "content": "done"}]},
|
||||
],
|
||||
}
|
||||
return {
|
||||
"model": "claude",
|
||||
"max_tokens": 32,
|
||||
"messages": [user],
|
||||
"system": TEXT
|
||||
if shape == "system-string"
|
||||
else [{"type": "text", "text": TEXT, "cache_control": {"type": "ephemeral"}}],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
("plain", "result-string", "result-nested", "input-dict", "input-assignment", "system-string", "system-blocks"),
|
||||
)
|
||||
@pytest.mark.parametrize("stream", (False, True))
|
||||
def test_messages_secret_never_reaches_provider(rig: Rig, shape: str, stream: bool) -> None:
|
||||
rig.wire.drain()
|
||||
request: Final = {**_messages(shape), "stream": stream, "cache": {"no-cache": True}}
|
||||
response: Final = rig.proxy.request("POST", "/v1/messages", request)
|
||||
assert response.status_code == 200, response.text
|
||||
assert "accepted" in response.text
|
||||
sent: Final = rig.wire.drain()
|
||||
assert len(sent) == 1
|
||||
assert SECRET.encode() not in sent[0].body
|
||||
assert b"0123456789abcdef0123456789abcdef01234567" not in sent[0].body
|
||||
assert b"REDACTED" in sent[0].body
|
||||
observed: Final = json.loads(sent[0].body)
|
||||
if shape == "input-dict":
|
||||
value: Final = observed["messages"][1]["content"][0]["input"]
|
||||
assert value.keys() == _input().keys()
|
||||
assert value["path"] == _input()["path"]
|
||||
assert value["count"] == 3
|
||||
assert value["enabled"] is False
|
||||
assert value["optional"] is None
|
||||
|
||||
|
||||
def test_chat_tool_arguments_remain_valid_json(rig: Rig) -> None:
|
||||
rig.wire.drain()
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "gpt",
|
||||
"cache": {"no-cache": True},
|
||||
"messages": [
|
||||
{"role": "user", "content": "read .env"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "read_file", "arguments": json.dumps(_input())},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "done"},
|
||||
],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["choices"][0]["message"]["content"] == "accepted"
|
||||
sent: Final = rig.wire.drain()
|
||||
assert len(sent) == 1
|
||||
assert SECRET.encode() not in sent[0].body
|
||||
arguments: Final = json.loads(json.loads(sent[0].body)["messages"][1]["tool_calls"][0]["function"]["arguments"])
|
||||
assert arguments.keys() == _input().keys()
|
||||
assert arguments["count"] == 3
|
||||
assert arguments["path"] == _input()["path"]
|
||||
assert arguments["enabled"] is False
|
||||
assert arguments["optional"] is None
|
||||
|
||||
|
||||
def test_deeply_nested_tool_arguments_do_not_stall_the_worker(rig: Rig) -> None:
|
||||
def call(depth: int) -> float:
|
||||
rig.wire.drain()
|
||||
arguments: Final = '{"k":' * depth + json.dumps(TEXT) + "}" * depth
|
||||
sent: Final = timed(
|
||||
lambda: rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "gpt",
|
||||
"cache": {"no-cache": True},
|
||||
"messages": [
|
||||
{"role": "user", "content": "read .env"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "read", "arguments": arguments},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "done"},
|
||||
],
|
||||
},
|
||||
)
|
||||
)
|
||||
assert sent.value.status_code == 200, sent.value.text
|
||||
upstream: Final = rig.wire.drain()
|
||||
assert len(upstream) == 1
|
||||
assert SECRET.encode() not in upstream[0].body
|
||||
return sent.seconds
|
||||
|
||||
control: Final = tuple(call(1) for _ in range(3))
|
||||
assert_not_delayed_by((call(300),), control=control, injected_delay=5.0)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def single_worker(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with tool_rig(tmp_path_factory.mktemp("hide-secrets-one-worker"), workers=1) as value:
|
||||
yield value
|
||||
|
||||
|
||||
def _tool_history(turns: int, seed: int, shape: str = "anthropic") -> tuple[str, dict[str, JsonValue]]:
|
||||
generator: Final = random.Random(seed)
|
||||
alphabet: Final = string.ascii_letters + string.digits
|
||||
messages: Final[list[JsonValue]] = [{"role": "user", "content": "read the project"}]
|
||||
for turn in range(turns):
|
||||
result = "\n".join("".join(generator.choices(alphabet, k=60)) for _ in range(64))
|
||||
if shape == "anthropic":
|
||||
call: JsonValue = {
|
||||
"type": "tool_use",
|
||||
"id": f"tool_{turn}",
|
||||
"name": "read",
|
||||
"input": {"path": f"{turn}.py"},
|
||||
}
|
||||
messages.append({"role": "assistant", "content": [call]})
|
||||
messages.append(
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": f"tool_{turn}", "content": result}]}
|
||||
)
|
||||
else:
|
||||
function: JsonValue = {"name": "read", "arguments": json.dumps({"path": f"{turn}.py"})}
|
||||
messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{"id": f"call_{turn}", "type": "function", "function": function}],
|
||||
}
|
||||
)
|
||||
messages.append({"role": "tool", "tool_call_id": f"call_{turn}", "content": result})
|
||||
if shape == "anthropic":
|
||||
return "/v1/messages", {"model": "claude", "max_tokens": 32, "messages": messages, "cache": {"no-cache": True}}
|
||||
return "/v1/chat/completions", {"model": "gpt", "messages": messages, "cache": {"no-cache": True}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ("anthropic", "chat"))
|
||||
def test_resent_tool_history_is_scanned_once(single_worker: Rig, shape: str) -> None:
|
||||
path, body = _tool_history(150, seed=2 if shape == "anthropic" else 4, shape=shape)
|
||||
turns: Final = body["messages"]
|
||||
assert isinstance(turns, list)
|
||||
first: Final = timed(lambda: single_worker.proxy.request("POST", path, body))
|
||||
following: Final = {
|
||||
**body,
|
||||
"messages": [*turns, {"role": "assistant", "content": "ok"}, {"role": "user", "content": "next"}],
|
||||
}
|
||||
second: Final = timed(lambda: single_worker.proxy.request("POST", path, following))
|
||||
assert first.value.status_code == 200 and second.value.status_code == 200, (first.value.text, second.value.text)
|
||||
assert_within_factor_of_control(second.seconds, control=first.seconds, factor=0.5)
|
||||
|
|
@ -11,13 +11,15 @@ Covers the three defects from the ticket:
|
|||
"""
|
||||
|
||||
import cProfile
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from litellm_enterprise.enterprise_callbacks import secret_detection
|
||||
from litellm_enterprise.enterprise_callbacks.secret_detection import (
|
||||
_default_detect_secrets_config,
|
||||
_ENTERPRISE_SecretDetection,
|
||||
|
|
@ -48,6 +50,251 @@ def _guardrail() -> _ENTERPRISE_SecretDetection:
|
|||
return _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets", event_hook="pre_call", default_on=True)
|
||||
|
||||
|
||||
def test_tool_payload_scans_raw_values_and_key_context_without_changing_keys():
|
||||
payload = {
|
||||
"password": "hunter2brahms9x",
|
||||
"command": "export API_TOKEN=0123456789abcdef0123456789abcdef01234567",
|
||||
AWS_KEY: "safe",
|
||||
"path": 'C:\\tmp\\"config"',
|
||||
"count": 7,
|
||||
"optional": None,
|
||||
"enabled": False,
|
||||
"nested": [f"AWS_ACCESS_KEY_ID={AWS_KEY}\nkeep this line", "safe"],
|
||||
}
|
||||
result = _guardrail().redact_tool_payload(payload)
|
||||
assert result.keys() == payload.keys()
|
||||
assert result["password"] == "[REDACTED]"
|
||||
assert result["command"] == "export API_TOKEN=[REDACTED]"
|
||||
assert result["nested"] == ["AWS_ACCESS_KEY_ID=[REDACTED]\nkeep this line", "safe"]
|
||||
assert result[AWS_KEY] == "safe"
|
||||
assert result["path"] == payload["path"]
|
||||
assert result["count"] == 7
|
||||
assert result["optional"] is None
|
||||
assert result["enabled"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"arguments",
|
||||
(
|
||||
'{"env":"AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE","env":"safe"}',
|
||||
'{"env":"AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE",broken',
|
||||
json.dumps({"password": "hunter2brahms9x", "ok": True}),
|
||||
),
|
||||
)
|
||||
async def test_tool_arguments_native_hook_cannot_forward_unredacted_credentials(arguments):
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [{"id": "call_1", "function": {"name": "save_config", "arguments": arguments}}],
|
||||
}
|
||||
]
|
||||
}
|
||||
await _guardrail().async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion")
|
||||
rewritten = data["messages"][0]["tool_calls"][0]["function"]["arguments"]
|
||||
assert AWS_KEY not in rewritten
|
||||
assert "hunter2brahms9x" not in rewritten
|
||||
assert "[REDACTED]" in rewritten
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deeply_nested_tool_arguments_are_redacted():
|
||||
arguments = '{"k":' * 200 + json.dumps(f"AWS_ACCESS_KEY_ID={AWS_KEY}") + "}" * 200
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [{"id": "call_1", "function": {"name": "save_config", "arguments": arguments}}],
|
||||
}
|
||||
]
|
||||
}
|
||||
await _guardrail().async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion")
|
||||
rewritten = data["messages"][0]["tool_calls"][0]["function"]["arguments"]
|
||||
assert json.loads(rewritten) == json.loads(arguments.replace(AWS_KEY, "[REDACTED]"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("password", ("hunter2\\brahms9x", "hunter2\tbrahms9x"))
|
||||
def test_tool_payload_redacts_keyword_values_that_need_json_escaping(password):
|
||||
assert _guardrail().redact_tool_payload({"password": password, "path": ".env"}) == {
|
||||
"password": "[REDACTED]",
|
||||
"path": ".env",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("password", "redacted"),
|
||||
(('Zx9qL2mN8vB4"kR7Tq', '[REDACTED]"kR7Tq'), ('Zx9qL2mN8vB4\\"kR7Tq', '[REDACTED]"kR7Tq')),
|
||||
ids=("quote", "backslash-quote"),
|
||||
)
|
||||
def test_tool_payload_redacts_a_keyword_match_that_stops_inside_an_escape(password: str, redacted: str) -> None:
|
||||
assert _guardrail().redact_tool_payload({"password": password}) == {"password": redacted}
|
||||
|
||||
|
||||
def _recorded_scans(monkeypatch: pytest.MonkeyPatch) -> list[str]:
|
||||
scanned: Final[list[str]] = []
|
||||
scan: Final = secret_detection._scan_lines
|
||||
|
||||
def recorded(lines: Sequence[str]) -> frozenset[tuple[str, str]]:
|
||||
scanned.append(lines[0])
|
||||
return scan(lines)
|
||||
|
||||
monkeypatch.setattr(secret_detection, "_scan_lines", recorded)
|
||||
return scanned
|
||||
|
||||
|
||||
def test_clean_scans_are_reused_and_secret_scans_are_not(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
scanned: Final = _recorded_scans(monkeypatch)
|
||||
guardrail: Final = _guardrail()
|
||||
for _ in range(2):
|
||||
assert guardrail.redact_text("nothing sensitive here") == "nothing sensitive here"
|
||||
assert guardrail.redact_text(f"AWS_ACCESS_KEY_ID={AWS_KEY}") == "AWS_ACCESS_KEY_ID=[REDACTED]"
|
||||
assert _guardrail().redact_text("nothing sensitive here") == "nothing sensitive here"
|
||||
assert scanned == [
|
||||
"nothing sensitive here",
|
||||
f"AWS_ACCESS_KEY_ID={AWS_KEY}",
|
||||
f"AWS_ACCESS_KEY_ID={AWS_KEY}",
|
||||
"nothing sensitive here",
|
||||
]
|
||||
|
||||
|
||||
def test_clean_scan_reuse_is_bounded(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
scanned: Final = _recorded_scans(monkeypatch)
|
||||
monkeypatch.setattr(secret_detection, "_CLEAN_SCAN_CACHE_SIZE", 2)
|
||||
guardrail: Final = _guardrail()
|
||||
for text in ("first", "second", "first", "third", "second", "first"):
|
||||
guardrail.scan_message_for_secrets(text)
|
||||
assert scanned == ["first", "second", "third", "second", "first"]
|
||||
|
||||
|
||||
def test_interleaved_agent_sessions_reuse_their_clean_scans(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
scanned: Final[list[str]] = []
|
||||
|
||||
def clean_scan(lines: Sequence[str]) -> frozenset[tuple[str, str]]:
|
||||
scanned.append(lines[0])
|
||||
return frozenset()
|
||||
|
||||
monkeypatch.setattr(secret_detection, "_scan_lines", clean_scan)
|
||||
guardrail: Final = _guardrail()
|
||||
sessions: Final = [[f"session {session} turn {turn}" for turn in range(201)] for session in range(25)]
|
||||
for history in sessions:
|
||||
for text in history:
|
||||
guardrail.redact_text(text)
|
||||
scanned.clear()
|
||||
for history in sessions:
|
||||
for text in history:
|
||||
guardrail.redact_text(text)
|
||||
assert scanned == []
|
||||
|
||||
|
||||
def test_overlapping_secrets_are_replaced_longest_first(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
guardrail: Final = _guardrail()
|
||||
detected: Final = [
|
||||
{"type": "Secret Keyword", "value": "\\u0041\\u0042"},
|
||||
{"type": "Secret Keyword", "value": "ABCD"},
|
||||
]
|
||||
|
||||
def overlapping_scan(text: str) -> list[dict[str, str]]:
|
||||
return detected
|
||||
|
||||
monkeypatch.setattr(guardrail, "scan_message_for_secrets", overlapping_scan)
|
||||
assert guardrail.redact_tool_payload({"token": "ABCDxyz"}) == {"token": "[REDACTED]xyz"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_records_secrets_found_only_in_tool_payloads():
|
||||
data = {
|
||||
"system": "You are a coding agent",
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "t1", "name": "write", "input": {"env": f"AWS_ACCESS_KEY_ID={AWS_KEY}"}}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "written"}]},
|
||||
],
|
||||
"metadata": {},
|
||||
}
|
||||
await _guardrail().async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="anthropic_messages"
|
||||
)
|
||||
assert data["messages"][0]["content"][0]["input"] == {"env": "AWS_ACCESS_KEY_ID=[REDACTED]"}
|
||||
recorded = _recorded(data)
|
||||
assert recorded["guardrail_response"] == "mask"
|
||||
assert recorded["masked_entity_count"] == {"AWS Access Key": 1}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("payload", "expected"),
|
||||
(
|
||||
({"password": ["hunter2brahms9x", "safe"]}, {"password": ["[REDACTED]", "safe"]}),
|
||||
(["a8f5f167f44f4964e6c998dee827110c", "safe"], ["[REDACTED]", "safe"]),
|
||||
({"password" + " " * 100: ["hunter2brahms9x"]}, {"password" + " " * 100: ["[REDACTED]"]}),
|
||||
),
|
||||
ids=("credential-key-over-list", "unkeyed-high-entropy", "whitespace-padded-credential-key"),
|
||||
)
|
||||
def test_list_items_are_scanned_beside_their_parent_key_or_quoted(payload, expected):
|
||||
assert _guardrail().redact_tool_payload(payload) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key", "items"),
|
||||
(("k" * 10_000, 1000), ("\u0000" * 64, 3000), ("tok " * 250_000, 2000)),
|
||||
ids=("long", "control-characters", "long-over-many-items"),
|
||||
)
|
||||
def test_keys_over_a_list_build_one_bounded_context(monkeypatch: pytest.MonkeyPatch, key: str, items: int) -> None:
|
||||
contexts: Final[list[str]] = []
|
||||
scans: Final[list[str]] = []
|
||||
key_context: Final = secret_detection._key_context
|
||||
guardrail: Final = _guardrail()
|
||||
scan: Final = guardrail.scan_message_for_secrets
|
||||
|
||||
def recorded_context(value: str) -> str:
|
||||
contexts.append(value)
|
||||
return key_context(value)
|
||||
|
||||
def recorded_scan(text: str) -> list[dict[str, str]]:
|
||||
scans.append(text)
|
||||
return scan(text)
|
||||
|
||||
monkeypatch.setattr(secret_detection, "_key_context", recorded_context)
|
||||
monkeypatch.setattr(guardrail, "scan_message_for_secrets", recorded_scan)
|
||||
secret_line: Final = f"AWS_ACCESS_KEY_ID={AWS_KEY}"
|
||||
result: Final = guardrail.redact_tool_payload({key: ["v"] * items + [secret_line]})
|
||||
assert result == {key: ["v"] * items + ["AWS_ACCESS_KEY_ID=[REDACTED]"]}
|
||||
assert contexts == [key]
|
||||
longest: Final = secret_detection._KEY_CONTEXT + len('"": ') + len(json.dumps(secret_line))
|
||||
assert max(len(line) for line in scans[0].splitlines()) == longest
|
||||
|
||||
|
||||
def test_long_multi_line_tool_strings_are_scanned_in_bounded_runs_beside_their_key(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
scans: Final[list[str]] = []
|
||||
guardrail: Final = _guardrail()
|
||||
scan: Final = guardrail.scan_message_for_secrets
|
||||
|
||||
def recorded_scan(text: str) -> list[dict[str, str]]:
|
||||
scans.append(text)
|
||||
return scan(text)
|
||||
|
||||
monkeypatch.setattr(guardrail, "scan_message_for_secrets", recorded_scan)
|
||||
code: Final = "\n".join(f"const value{i} = compute({i}); // step {i}" for i in range(2000))
|
||||
padding: Final = "a" * 5000
|
||||
payload: Final = {
|
||||
"password": "hunter2brahms9x\nsecond",
|
||||
"secret": f"{padding}\nzq7Lm2Xv9Rt4Kp\nthird\n{padding}",
|
||||
"content": f"AWS_ACCESS_KEY_ID={AWS_KEY}\n{code}",
|
||||
}
|
||||
assert guardrail.redact_tool_payload(payload) == {
|
||||
"password": "[REDACTED]",
|
||||
"secret": "[REDACTED]\n[REDACTED]",
|
||||
"content": f"AWS_ACCESS_KEY_ID=[REDACTED]\n{code}",
|
||||
}
|
||||
assert max(len(line) for line in scans[0].splitlines()) < len(code)
|
||||
|
||||
|
||||
def _recorded(request_data: dict) -> dict:
|
||||
entries = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(entries) == 1
|
||||
|
|
|
|||
|
|
@ -1,6 +1,10 @@
|
|||
"""Tests for the shared guardrail content extraction helpers."""
|
||||
|
||||
import json
|
||||
from functools import reduce
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.proxy.guardrails._content_utils import (
|
||||
apply_redacted_messages_back,
|
||||
|
|
@ -9,6 +13,8 @@ from litellm.proxy.guardrails._content_utils import (
|
|||
is_non_conversational_call_type,
|
||||
is_string_batch_input,
|
||||
iter_message_text,
|
||||
iter_tool_json_text,
|
||||
map_tool_json_text,
|
||||
walk_user_text,
|
||||
)
|
||||
|
||||
|
|
@ -554,9 +560,7 @@ def test_build_inspection_messages_includes_reasoning_summary():
|
|||
}
|
||||
]
|
||||
}
|
||||
assert build_inspection_messages(data) == [
|
||||
{"role": "assistant", "content": "secret summary"}
|
||||
]
|
||||
assert build_inspection_messages(data) == [{"role": "assistant", "content": "secret summary"}]
|
||||
|
||||
|
||||
# ── has_non_string_content ────────────────────────────────────────────────────
|
||||
|
|
@ -716,6 +720,7 @@ def test_is_string_batch_input_rejects_other_shapes():
|
|||
# LIT-4302: custom_tool_call_output walking
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_iter_message_text_walks_custom_tool_call_output():
|
||||
"""custom_tool_call_output items should yield their output text."""
|
||||
data = {
|
||||
|
|
@ -724,6 +729,7 @@ def test_iter_message_text_walks_custom_tool_call_output():
|
|||
]
|
||||
}
|
||||
from litellm.proxy.guardrails._content_utils import iter_message_text
|
||||
|
||||
texts = list(iter_message_text(data))
|
||||
assert "tool-secret" in texts
|
||||
|
||||
|
|
@ -740,6 +746,34 @@ def test_walk_user_text_redacts_custom_tool_call_output():
|
|||
assert data["input"][0]["output"] == "[MASKED]"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("item", "field", "masked"),
|
||||
(
|
||||
(
|
||||
{"type": "function_call", "name": "secret_tool_name", "arguments": '{"token": "secret", "count": 7}'},
|
||||
"arguments",
|
||||
'{"token": "[MASKED]", "count": 7}',
|
||||
),
|
||||
({"type": "function_call", "arguments": '{"token":"secret",broken'}, "arguments", '{"token":"[MASKED]",broken'),
|
||||
({"type": "custom_tool_call", "name": "secret_tool_name", "input": "echo secret"}, "input", "echo [MASKED]"),
|
||||
),
|
||||
ids=("function-call", "function-call-malformed", "custom-tool-call"),
|
||||
)
|
||||
def test_responses_tool_call_items_are_masked(item, field, masked):
|
||||
data = {"input": [{"role": "user", "content": "go"}, item]}
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 2
|
||||
assert data["input"][1][field] == masked
|
||||
assert data["input"][1].get("name", "secret_tool_name") == "secret_tool_name"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("instructions", ("secret rules", [{"type": "input_text", "text": "secret rules"}]))
|
||||
def test_responses_instructions_are_masked(instructions):
|
||||
data = {"instructions": instructions, "input": "hi"}
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 2
|
||||
masked = data["instructions"] if isinstance(instructions, str) else data["instructions"][0]["text"]
|
||||
assert masked == "[MASKED] rules"
|
||||
|
||||
|
||||
def test_build_inspection_messages_custom_tool_call_output():
|
||||
"""build_inspection_messages should include custom_tool_call_output text."""
|
||||
data = {
|
||||
|
|
@ -782,3 +816,273 @@ def test_is_non_conversational_call_type_defaults_to_inspecting_unknown_call_typ
|
|||
"""A call type this module has never heard of must still be inspected —
|
||||
failing closed is the point of the deny-list."""
|
||||
assert is_non_conversational_call_type("some_future_call_type") is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_input",
|
||||
(
|
||||
{"token": "secret", "path": 'C:\\tmp\\"file"', "nested": [False, 7, None, "secret"]},
|
||||
{"command": "export API_TOKEN=secret\nnext line", "safe": "unchanged"},
|
||||
),
|
||||
)
|
||||
def test_native_tool_payload_rewrites_values_preserving_structure(tool_input):
|
||||
data = {
|
||||
"system": [{"type": "text", "text": "secret", "cache_control": {"type": "ephemeral"}}],
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_use", "name": "secret_tool_name", "id": "secret_id", "input": tool_input}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "secret_id",
|
||||
"content": [
|
||||
{"type": "text", "text": "secret"},
|
||||
{"type": "image", "source": {"data": "secret"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
visited = walk_user_text(data, lambda value: value.replace("secret", "[MASKED]"))
|
||||
assert visited == 3
|
||||
block = data["messages"][0]["content"][0]
|
||||
assert block["name"] == "secret_tool_name"
|
||||
assert block["id"] == "secret_id"
|
||||
assert block["input"] == json.loads(json.dumps(tool_input).replace("secret", "[MASKED]"))
|
||||
result = data["messages"][1]["content"][0]
|
||||
assert result["tool_use_id"] == "secret_id"
|
||||
assert result["content"] == [{"type": "text", "text": "[MASKED]"}, {"type": "image", "source": {"data": "secret"}}]
|
||||
assert data["system"] == [{"type": "text", "text": "[MASKED]", "cache_control": {"type": "ephemeral"}}]
|
||||
|
||||
|
||||
def _tool_call(arguments: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{"id": "call_1", "function": {"name": "secret_tool_name", "arguments": arguments}}],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"arguments",
|
||||
(
|
||||
'{ "token": "secret", "count": 7, "ok": false, "none": null }',
|
||||
{"token": "secret", "count": 7, "ok": False, "none": None},
|
||||
),
|
||||
ids=("string", "dict"),
|
||||
)
|
||||
def test_native_tool_call_arguments_are_masked_as_json(arguments: JsonValue) -> None:
|
||||
data = _tool_call(arguments)
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1
|
||||
function = data["messages"][0]["tool_calls"][0]["function"]
|
||||
masked = json.loads(function["arguments"]) if isinstance(arguments, str) else function["arguments"]
|
||||
assert masked == {"token": "[MASKED]", "count": 7, "ok": False, "none": None}
|
||||
assert function["name"] == "secret_tool_name"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"arguments",
|
||||
('{"token":"secret", "token":"safe"}', '{"token":"secret",broken', '{"pattern":"[\\ud800-\\udbff] secret"}'),
|
||||
ids=("duplicate-key", "malformed", "lone-surrogate"),
|
||||
)
|
||||
def test_native_tool_call_arguments_that_are_not_clean_json_are_masked_as_raw_text(arguments: str) -> None:
|
||||
data = _tool_call(arguments)
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1
|
||||
assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == arguments.replace("secret", "[MASKED]")
|
||||
|
||||
|
||||
def test_native_tool_arguments_keep_their_bytes_when_nothing_is_masked():
|
||||
arguments = '{ "content" : "caf\\u00e9", "count": 3 }'
|
||||
data = {"messages": [{"tool_calls": [{"function": {"arguments": arguments}}]}]}
|
||||
assert walk_user_text(data, lambda text: text) == 1
|
||||
assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == arguments
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data",
|
||||
(
|
||||
{"system": "payload"},
|
||||
{"system": [{"type": "text", "text": "payload"}]},
|
||||
{"messages": [{"content": [{"type": "tool_result", "content": "payload"}]}]},
|
||||
{"messages": [{"content": [{"type": "tool_result", "content": [{"type": "text", "text": "payload"}]}]}]},
|
||||
{"messages": [{"content": [{"type": "tool_use", "input": {"nested": [False, {"field": "payload"}]}}]}]},
|
||||
{"messages": [{"tool_calls": [{"function": {"arguments": '{"field":"payload"}'}}]}]},
|
||||
{"messages": [{"tool_calls": [{"function": {"arguments": {"field": "payload"}}}]}]},
|
||||
),
|
||||
)
|
||||
def test_native_writer_visits_each_tool_and_system_slot(data):
|
||||
assert walk_user_text(data, lambda text: text.replace("payload", "[MASKED]")) == 1
|
||||
assert "payload" not in json.dumps(data)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"part",
|
||||
(
|
||||
{"type": "tool_result", "content": "secret", "text": "safe"},
|
||||
{"type": "tool_use", "input": {"env": "secret"}, "text": "safe"},
|
||||
),
|
||||
)
|
||||
def test_native_tool_envelope_text_cannot_shadow_its_payload(part):
|
||||
data = {"messages": [{"role": "assistant", "content": [part]}]}
|
||||
walk_user_text(data, lambda text: text.replace("secret", "[REDACTED]"))
|
||||
assert "secret" not in json.dumps(data)
|
||||
assert data["messages"][0]["content"][0]["text"] == "safe"
|
||||
|
||||
|
||||
def _deep_tool_value(depth: int, leaf: JsonValue) -> JsonValue:
|
||||
return reduce(lambda value, level: {"nested": value} if level % 2 else [value], range(depth), leaf)
|
||||
|
||||
|
||||
def _deep_tool_leaf(value: JsonValue, depth: int) -> JsonValue:
|
||||
def child(current: JsonValue, _level: int) -> JsonValue:
|
||||
assert isinstance(current, (list, dict))
|
||||
return current[0] if isinstance(current, list) else current["nested"]
|
||||
|
||||
return reduce(child, range(depth), value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("depth", (255, 1000))
|
||||
def test_deep_tool_input_writer_preserves_values(depth: int) -> None:
|
||||
leaf = {"secret-key": "secret", "empty": "", "enabled": False, "count": 7, "ratio": 1.5, "none": None}
|
||||
payload = _deep_tool_value(depth, leaf)
|
||||
data = {"messages": [{"content": [{"type": "tool_use", "id": "tool", "name": "read", "input": payload}]}]}
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1
|
||||
result = _deep_tool_leaf(data["messages"][0]["content"][0]["input"], depth)
|
||||
assert isinstance(result, dict)
|
||||
assert tuple(result) == tuple(leaf)
|
||||
assert result == {**leaf, "secret-key": "[MASKED]"}
|
||||
assert result["enabled"] is False and result["none"] is None
|
||||
assert type(result["count"]) is int and type(result["ratio"]) is float
|
||||
assert leaf["secret-key"] == "secret"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("depth", (255, 1000))
|
||||
def test_deep_tool_json_direct_map_and_iteration(depth: int) -> None:
|
||||
value = _deep_tool_value(depth, {"secret": "secret", "tail": [False, None, "last", ""]})
|
||||
rewritten = map_tool_json_text(value, lambda text: text.replace("secret", "[MASKED]"))
|
||||
assert tuple(iter_tool_json_text(value)) == ("secret", "last")
|
||||
assert tuple(iter_tool_json_text(rewritten)) == ("[MASKED]", "last")
|
||||
assert _deep_tool_leaf(rewritten, depth) == {"secret": "[MASKED]", "tail": [False, None, "last", ""]}
|
||||
|
||||
|
||||
def test_deep_tool_arguments_decode_values_and_preserve_no_hit_bytes() -> None:
|
||||
arguments = '{ "nested" : ' * 255 + '{"caf\\u00e9":"caf\\u00e9 secret","empty":"","count":7}' + " }" * 255
|
||||
data = {"messages": [{"tool_calls": [{"function": {"name": "secret_tool", "arguments": arguments}}]}]}
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1
|
||||
function = data["messages"][0]["tool_calls"][0]["function"]
|
||||
decoded = json.loads(function["arguments"])
|
||||
for _ in range(255):
|
||||
decoded = decoded["nested"]
|
||||
assert tuple(decoded) == ("café", "empty", "count")
|
||||
assert decoded == {"café": "café [MASKED]", "empty": "", "count": 7}
|
||||
assert function["name"] == "secret_tool"
|
||||
clean = arguments.replace("secret", "safe")
|
||||
untouched = {"messages": [{"tool_calls": [{"function": {"arguments": clean}}]}]}
|
||||
assert walk_user_text(untouched, lambda text: text) == 1
|
||||
assert untouched["messages"][0]["tool_calls"][0]["function"]["arguments"] == clean
|
||||
|
||||
|
||||
def test_deep_tool_arguments_parser_limit_keeps_raw_masking_fallback() -> None:
|
||||
arguments = ""
|
||||
for depth in (1000, 2000, 4000, 8000, 16000):
|
||||
candidate = '{"nested":' * depth + '{"value":"secret"}' + "}" * depth
|
||||
try:
|
||||
json.loads(candidate, object_pairs_hook=dict)
|
||||
except RecursionError:
|
||||
arguments = candidate
|
||||
break
|
||||
assert arguments, "JSON decoder did not reach its recursion guard"
|
||||
data = {"messages": [{"tool_calls": [{"function": {"arguments": arguments}}]}]}
|
||||
try:
|
||||
visited = walk_user_text(data, lambda text: text.replace("secret", "[MASKED]"))
|
||||
except RecursionError:
|
||||
visited = 0
|
||||
assert visited == 1
|
||||
rewritten = data["messages"][0]["tool_calls"][0]["function"]["arguments"]
|
||||
assert rewritten == arguments.replace("secret", "[MASKED]")
|
||||
|
||||
|
||||
def test_tool_json_graph_repeated_shared_containers_keep_order_and_empty_leaves() -> None:
|
||||
shared = {"secret-key": "secret", "empty": "", "scalars": [False, 7, 1.5, None], "tail": "last"}
|
||||
value = {"first": shared, "second": shared, "empty": []}
|
||||
visited = []
|
||||
|
||||
def visit(text: str) -> str:
|
||||
visited.append(text)
|
||||
return text.replace("secret", "[MASKED]")
|
||||
|
||||
rewritten = map_tool_json_text(value, visit)
|
||||
assert visited == ["secret", "last", "secret", "last"]
|
||||
assert tuple(iter_tool_json_text(value)) == ("secret", "last", "secret", "last")
|
||||
assert isinstance(rewritten, dict) and tuple(rewritten) == ("first", "second", "empty")
|
||||
assert rewritten["first"] == rewritten["second"] == {**shared, "secret-key": "[MASKED]"}
|
||||
assert rewritten["empty"] == [] and shared["secret-key"] == "secret"
|
||||
|
||||
|
||||
def test_nested_tool_result_chain_leaves_later_messages_masked() -> None:
|
||||
chain = reduce(
|
||||
lambda content, _: [{"type": "tool_result", "tool_use_id": "tool", "content": content}],
|
||||
range(1500),
|
||||
[{"type": "text", "text": "inner"}],
|
||||
)
|
||||
data = {"messages": [{"role": "assistant", "content": chain}, {"role": "user", "content": "secret"}]}
|
||||
walk_user_text(data, lambda text: text.replace("secret", "[MASKED]"))
|
||||
assert data["messages"][1]["content"] == "[MASKED]"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"arguments",
|
||||
('{"token":"secret","limit":1e400}', '{"token":"secret","limit":-Infinity}', '{"token":"secret","ratio":NaN}'),
|
||||
)
|
||||
def test_tool_arguments_with_non_finite_numbers_keep_their_original_bytes(arguments: str) -> None:
|
||||
data = {"messages": [{"tool_calls": [{"function": {"arguments": arguments}}]}]}
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1
|
||||
assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == arguments.replace("secret", "[MASKED]")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("arguments", "masked"),
|
||||
(
|
||||
(
|
||||
'{ "token" : "secret", "amount": 0.123456789012345678901, "price": 1.50, "limit": 1e5 }',
|
||||
'{ "token" : "[MASKED]", "amount": 0.123456789012345678901, "price": 1.50, "limit": 1e5 }',
|
||||
),
|
||||
(
|
||||
'{"caf\\u00e9":"caf\\u00e9","list":["secret", "", "a\\"secret\\"b"],\n "secret": "keep"}',
|
||||
'{"caf\\u00e9":"caf\\u00e9","list":["[MASKED]", "", "a\\"[MASKED]\\"b"],\n "secret": "keep"}',
|
||||
),
|
||||
),
|
||||
ids=("number-literals", "escapes-and-keys"),
|
||||
)
|
||||
def test_masked_tool_arguments_change_only_the_masked_strings(arguments: str, masked: str) -> None:
|
||||
data = {"messages": [{"tool_calls": [{"function": {"arguments": arguments}}]}]}
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1
|
||||
assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == masked
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"rewritten",
|
||||
(
|
||||
{"token": "[MASKED]"},
|
||||
{"token": "[MASKED]", "renamed": "keep", "count": 7},
|
||||
{"token": "[MASKED]", "note": "keep", "count": 8},
|
||||
{"token": "[MASKED]", "note": "keep", "count": True},
|
||||
),
|
||||
ids=("key-dropped", "key-renamed", "number-changed", "number-retyped"),
|
||||
)
|
||||
def test_tool_visitor_that_changes_more_than_strings_writes_its_payload_as_arguments(rewritten: JsonValue) -> None:
|
||||
data = {
|
||||
"messages": [{"tool_calls": [{"function": {"arguments": '{ "token": "secret", "note": "keep", "count": 1 }'}}]}]
|
||||
}
|
||||
assert walk_user_text(data, lambda text: text, lambda payload: rewritten) == 1
|
||||
assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == json.dumps(rewritten)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue