diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 8b17cf13cc4..65e4d50587f 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -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[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P\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 diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 8581dca603b..b4830b7c771 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -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 diff --git a/tests/integration/security/test_hide_secrets_tool_compatibility.py b/tests/integration/security/test_hide_secrets_tool_compatibility.py new file mode 100644 index 00000000000..d8b900de6d0 --- /dev/null +++ b/tests/integration/security/test_hide_secrets_tool_compatibility.py @@ -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"]) diff --git a/tests/integration/security/test_hide_secrets_tool_traffic.py b/tests/integration/security/test_hide_secrets_tool_traffic.py new file mode 100644 index 00000000000..42fdc0f5a44 --- /dev/null +++ b/tests/integration/security/test_hide_secrets_tool_traffic.py @@ -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) diff --git a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py index 12c54210579..9816eb91102 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py @@ -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 diff --git a/tests/unit/proxy/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py index c14375135dc..a668d5f609d 100644 --- a/tests/unit/proxy/guardrails/test_content_utils.py +++ b/tests/unit/proxy/guardrails/test_content_utils.py @@ -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) +