litellm/tests/unit/proxy/guardrails/test_content_utils.py

"""Tests for the shared guardrail content extraction helpers."""

import json
from functools import reduce
from typing import Final

import pytest
from pydantic import JsonValue

from litellm.proxy.guardrails._content_utils import (
    apply_redacted_messages_back,
    build_inspection_messages,
    has_non_string_content,
    is_non_conversational_call_type,
    is_non_conversational_request,
    is_string_batch_input,
    iter_message_text,
    iter_tool_json_text,
    map_tool_json_text,
    walk_user_text,
)

# ── iter_message_text ────────────────────────────────────────────────────────────


def test_iter_message_text_string_messages():
    data = {
        "messages": [
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hi"},
        ]
    }
    assert list(iter_message_text(data)) == ["hello", "hi"]


def test_iter_message_text_multimodal_list_content():
    """VERIA-11: list-format content must be inspected, not silently skipped."""
    data = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "AWS_KEY=AKIA..."},
                    {"type": "image_url", "image_url": {"url": "..."}},
                    {"type": "text", "text": "more secrets"},
                ],
            }
        ]
    }
    assert list(iter_message_text(data)) == ["AWS_KEY=AKIA...", "more secrets"]


def test_iter_message_text_responses_api_string_input():
    """fniVO9-F: Responses-API ``input`` must be inspectable when ``messages`` absent."""
    data = {"input": "tell me a secret"}
    assert list(iter_message_text(data)) == ["tell me a secret"]


def test_iter_message_text_responses_api_list_input_messages():
    data = {
        "input": [
            {"role": "user", "content": "first"},
            {"role": "user", "content": "second"},
        ]
    }
    assert list(iter_message_text(data)) == ["first", "second"]


def test_iter_message_text_responses_api_list_input_content_parts():
    data = {
        "input": [
            {"type": "text", "text": "alpha"},
            {"type": "image_url", "image_url": {"url": "..."}},
            {"type": "text", "text": "beta"},
        ]
    }
    assert list(iter_message_text(data)) == ["alpha", "beta"]


def test_iter_message_text_responses_api_list_input_mixed_dicts_and_strings():
    """Greptile P2: mixed-list ``input`` with content-part dicts AND bare
    strings must yield every text fragment — read helpers used to truncate
    the bare strings."""
    data = {
        "input": [
            {"type": "text", "text": "from-dict"},
            "from-bare-string",
            {"type": "image_url", "image_url": {"url": "..."}},
            "another-bare-string",
        ]
    }
    assert list(iter_message_text(data)) == [
        "from-dict",
        "from-bare-string",
        "another-bare-string",
    ]


def test_iter_message_text_walks_messages_and_input_independently():
    """When both are present (rare), every fragment from either field is
    inspected — a stricter guarantee than "first one wins"."""
    data = {
        "messages": [{"role": "user", "content": "msg-content"}],
        "input": "input-content",
    }
    assert list(iter_message_text(data)) == ["msg-content", "input-content"]


def test_iter_message_text_empty_data():
    assert list(iter_message_text({})) == []
    assert list(iter_message_text({"messages": []})) == []
    assert list(iter_message_text({"input": ""})) == []


def test_iter_message_text_responses_api_input_text_and_output_text_parts():
    """LIT-4294: Responses-API content parts use ``input_text`` (request) and
    ``output_text`` (assistant); reading only ``type == "text"`` skipped every
    ``/v1/responses`` body and every text guardrail was a no-op on that path."""
    data = {
        "input": [
            {
                "type": "message",
                "role": "user",
                "content": [{"type": "input_text", "text": "user text"}],
            },
            {
                "type": "message",
                "role": "assistant",
                "content": [{"type": "output_text", "text": "assistant text"}],
            },
        ]
    }
    assert list(iter_message_text(data)) == ["user text", "assistant text"]


def test_iter_message_text_responses_api_tool_call_taxonomy():
    """LIT-4294: a Responses ``input`` list freely mixes message items,
    ``function_call`` (no ``role``), and ``function_call_output`` items. The
    old ``all(item has 'role')`` gate wrapped the whole list as one blob and
    yielded nothing; every text fragment must be visited independently."""
    data = {
        "input": [
            {
                "type": "message",
                "role": "user",
                "content": [{"type": "input_text", "text": "hello"}],
            },
            {
                "type": "function_call",
                "call_id": "c1",
                "name": "get_weather",
                "arguments": "{}",
            },
            {
                "type": "function_call_output",
                "call_id": "c1",
                "output": [{"type": "input_text", "text": "sunny"}],
            },
        ]
    }
    assert list(iter_message_text(data)) == ["hello", "sunny"]


def test_iter_message_text_inspects_reasoning_content_and_summary():
    """VERIA: reasoning items forwarded as ``reasoning_content`` must be
    inspected, including ``summary`` blocks the bridge reads as a fallback."""
    data = {
        "input": [
            {
                "type": "reasoning",
                "id": "rs_1",
                "content": [{"type": "summary_text", "text": "content secret"}],
                "summary": [{"type": "summary_text", "text": "summary secret"}],
            }
        ]
    }
    assert list(iter_message_text(data)) == ["content secret", "summary secret"]



def test_iter_message_text_inspects_compaction_summary():
    data = {
        "input": [
            {
                "type": "compaction",
                "id": "cmp_1",
                "encrypted_content": json.dumps({"type": "compaction", "content": "banned word"}),
            },
            {"role": "user", "content": "next turn"},
        ]
    }
    assert list(iter_message_text(data)) == ["banned word", "next turn"]


# ── walk_user_text ────────────────────────────────────────────────────────────


def test_walk_user_text_redacts_string_messages_in_place():
    data = {
        "messages": [
            {"role": "user", "content": "leak: AKIAEXAMPLE"},
            {"role": "assistant", "content": "ok"},
        ]
    }
    visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
    assert visited == 2
    assert data["messages"][0]["content"] == "leak: [REDACTED]"
    assert data["messages"][1]["content"] == "ok"


def test_walk_user_text_redacts_multimodal_text_parts():
    """VERIA-11: list-content text parts must be mutable for in-place redaction."""
    data = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "AKIAEXAMPLE here"},
                    {"type": "image_url", "image_url": {"url": "..."}},
                    {"type": "text", "text": "no secret"},
                ],
            }
        ]
    }
    visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
    assert visited == 2
    parts = data["messages"][0]["content"]
    assert parts[0] == {"type": "text", "text": "[REDACTED] here"}
    # Non-text part must be left untouched.
    assert parts[1] == {"type": "image_url", "image_url": {"url": "..."}}
    assert parts[2] == {"type": "text", "text": "no secret"}


def test_walk_user_text_redacts_responses_api_string_input():
    data = {"input": "leak AKIAEXAMPLE"}
    visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
    assert visited == 1
    assert data["input"] == "leak [REDACTED]"


def test_walk_user_text_redacts_responses_api_list_input():
    data = {
        "input": [
            {"type": "text", "text": "AKIAEXAMPLE"},
            {"type": "image_url", "image_url": {"url": "..."}},
        ]
    }
    visited = walk_user_text(data, lambda s: f"[redacted]{s}[/]")
    assert visited == 1
    assert data["input"][0] == {"type": "text", "text": "[redacted]AKIAEXAMPLE[/]"}
    assert data["input"][1] == {"type": "image_url", "image_url": {"url": "..."}}


def test_walk_user_text_redacts_responses_input_text_and_output_text_parts():
    """LIT-4294: ``walk_user_text`` must recognise the Responses text-part
    variants so masking guardrails (secret detection, PII) actually redact
    ``/v1/responses`` bodies instead of no-op'ing on them."""
    data = {
        "input": [
            {
                "type": "message",
                "role": "user",
                "content": [{"type": "input_text", "text": "AKIAEXAMPLE"}],
            },
            {
                "type": "message",
                "role": "assistant",
                "content": [{"type": "output_text", "text": "AKIAEXAMPLE too"}],
            },
        ]
    }
    visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
    assert visited == 2
    assert data["input"][0]["content"][0] == {
        "type": "input_text",
        "text": "[REDACTED]",
    }
    assert data["input"][1]["content"][0] == {
        "type": "output_text",
        "text": "[REDACTED] too",
    }


def test_walk_user_text_redacts_function_call_output_text():
    """LIT-4294: tool-call round-trips carry secrets in
    ``function_call_output.output``; the redact walker must descend into it
    while leaving ``function_call`` items (call_id, arguments) untouched."""
    data = {
        "input": [
            {
                "type": "message",
                "role": "user",
                "content": [{"type": "input_text", "text": "AKIAEXAMPLE user"}],
            },
            {
                "type": "function_call",
                "call_id": "c1",
                "name": "get_weather",
                "arguments": '{"AKIAEXAMPLE": 1}',
            },
            {
                "type": "function_call_output",
                "call_id": "c1",
                "output": [{"type": "input_text", "text": "AKIAEXAMPLE tool"}],
            },
        ]
    }
    visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
    assert visited == 2
    assert data["input"][0]["content"][0]["text"] == "[REDACTED] user"
    assert data["input"][1] == {
        "type": "function_call",
        "call_id": "c1",
        "name": "get_weather",
        "arguments": '{"AKIAEXAMPLE": 1}',
    }
    assert data["input"][2]["output"][0]["text"] == "[REDACTED] tool"


def test_walk_user_text_redacts_function_call_output_string_output():
    """LIT-4294: ``function_call_output.output`` is also a plain string in
    OpenAI's Responses spec; the redact walker must handle both forms."""
    data = {
        "input": [
            {
                "type": "function_call_output",
                "call_id": "c1",
                "output": "AKIAEXAMPLE tool",
            },
        ]
    }
    visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
    assert visited == 1
    assert data["input"][0]["output"] == "[REDACTED] tool"


def test_walk_user_text_redacts_mixed_list_input():
    """Read and write helpers must agree on coverage — bare strings inside
    a mixed ``input`` list are inspected by both."""
    data = {
        "input": [
            {"type": "text", "text": "secret-one"},
            "secret-two",
            {"type": "image_url", "image_url": {"url": "..."}},
        ]
    }
    visited = walk_user_text(data, lambda s: f"<{s}>")
    assert visited == 2
    assert data["input"][0] == {"type": "text", "text": "<secret-one>"}
    assert data["input"][1] == "<secret-two>"
    assert data["input"][2] == {"type": "image_url", "image_url": {"url": "..."}}


def test_walk_user_text_redacts_reasoning_content_and_summary():
    """VERIA: in-place redaction must cover both plaintext shapes the bridge
    forwards from a reasoning item."""
    data = {
        "input": [
            {
                "type": "reasoning",
                "id": "rs_1",
                "content": [{"type": "summary_text", "text": "AKIAEXAMPLE content"}],
                "summary": [{"type": "summary_text", "text": "AKIAEXAMPLE summary"}],
            }
        ]
    }
    visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
    assert visited == 2
    item = data["input"][0]
    assert item["content"][0]["text"] == "[REDACTED] content"
    assert item["summary"][0]["text"] == "[REDACTED] summary"
    assert item["id"] == "rs_1"



def test_walk_user_text_redacts_compaction_summary_and_keeps_the_block():
    data = {
        "input": [
            {
                "type": "compaction",
                "id": "cmp_1",
                "encrypted_content": json.dumps(
                    {"type": "compaction", "content": "AKIAEXAMPLE summary", "signature": "sig"}
                ),
            }
        ]
    }
    visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
    item = data["input"][0]
    assert visited == 1
    assert (item["type"], item["id"]) == ("compaction", "cmp_1")
    assert json.loads(item["encrypted_content"]) == {
        "type": "compaction",
        "content": "[REDACTED] summary",
        "signature": "sig",
    }


# ── build_inspection_messages ─────────────────────────────────────────────────


def test_build_inspection_messages_chat_completion_passthrough():
    data = {
        "messages": [
            {"role": "system", "content": "be helpful"},
            {"role": "user", "content": "hi"},
        ]
    }
    assert build_inspection_messages(data) == [
        {"role": "system", "content": "be helpful"},
        {"role": "user", "content": "hi"},
    ]


def test_build_inspection_messages_joins_multimodal_text_parts():
    data = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "first part"},
                    {"type": "image_url", "image_url": {"url": "..."}},
                    {"type": "text", "text": "second part"},
                ],
            }
        ]
    }
    assert build_inspection_messages(data) == [{"role": "user", "content": "first part\nsecond part"}]


def test_build_inspection_messages_lifts_responses_api_input():
    """fniVO9-F: ``input`` must be visible to hooks that POST messages to a remote API."""
    data = {"input": "responses-api content"}
    assert build_inspection_messages(data) == [{"role": "user", "content": "responses-api content"}]


def test_build_inspection_messages_drops_messages_with_no_text():
    data = {
        "messages": [
            {"role": "user", "content": ""},
            {
                "role": "user",
                "content": [{"type": "image_url", "image_url": {"url": "..."}}],
            },
            {"role": "user", "content": "kept"},
        ]
    }
    assert build_inspection_messages(data) == [{"role": "user", "content": "kept"}]


def test_build_inspection_messages_responses_api_tool_call_taxonomy():
    """LIT-4294: mixed Responses ``input`` (message + function_call +
    function_call_output) must produce a non-empty inspection list. The
    customer's writeup reproduced a 422 from AIM's ``/fw/v1/analyze``
    (``No messages in the request``) when this synthesised list came back
    empty; every other guardrail silently scanned nothing on the same
    input."""
    data = {
        "input": [
            {
                "type": "message",
                "role": "user",
                "content": [{"type": "input_text", "text": "hello"}],
            },
            {
                "type": "function_call",
                "call_id": "c1",
                "name": "get_weather",
                "arguments": "{}",
            },
            {
                "type": "function_call_output",
                "call_id": "c1",
                "output": [{"type": "input_text", "text": "sunny"}],
            },
        ]
    }
    assert build_inspection_messages(data) == [
        {"role": "user", "content": "hello"},
        {"role": "tool", "content": "sunny"},
    ]


def test_build_inspection_messages_function_call_output_defaults_to_tool():
    """LIT-4294: a Responses ``function_call_output`` item is the semantic
    equivalent of a chat-completions ``role: "tool"`` message, so the shared
    helper synthesises ``role: "tool"`` when the item has no explicit role.
    AIM's schema-safe coercion happens at the AIM call site, not here."""
    data = {
        "input": [
            {
                "type": "function_call_output",
                "call_id": "c1",
                "output": [{"type": "input_text", "text": "tool text"}],
            },
        ]
    }
    assert build_inspection_messages(data) == [{"role": "tool", "content": "tool text"}]


def test_build_inspection_messages_function_call_output_preserves_explicit_role():
    """When ``function_call_output`` carries a caller-supplied ``role`` the
    shared helper preserves it rather than synthesising ``tool``."""
    data = {
        "input": [
            {
                "type": "function_call_output",
                "role": "assistant",
                "call_id": "c1",
                "output": [{"type": "input_text", "text": "tool text"}],
            },
        ]
    }
    assert build_inspection_messages(data) == [{"role": "assistant", "content": "tool text"}]


def test_build_inspection_messages_bare_content_part_preserves_explicit_role():
    """A bare content-part dict with an explicit ``role`` keeps it. Only
    absent roles get defaulted to ``user``."""
    data = {
        "input": [
            {"type": "input_text", "text": "no role"},
            {"type": "output_text", "role": "assistant", "text": "with role"},
        ]
    }
    assert build_inspection_messages(data) == [
        {"role": "user", "content": "no role"},
        {"role": "assistant", "content": "with role"},
    ]


def test_build_inspection_messages_message_item_preserves_role():
    """Responses message items carry a role explicitly; the shared helper
    passes it through untouched."""
    data = {
        "input": [
            {"type": "message", "role": "system", "content": [{"type": "input_text", "text": "sys"}]},
            {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "asst"}]},
        ]
    }
    assert build_inspection_messages(data) == [
        {"role": "system", "content": "sys"},
        {"role": "assistant", "content": "asst"},
    ]


def test_build_inspection_messages_empty_data():
    assert build_inspection_messages({}) == []
    assert build_inspection_messages({"messages": []}) == []
    assert build_inspection_messages({"input": ""}) == []


def test_build_inspection_messages_includes_reasoning_summary():
    """VERIA: remote guardrail APIs must see reasoning summaries even when
    the reasoning item has no ``content`` field."""
    data = {
        "input": [
            {
                "type": "reasoning",
                "id": "rs_1",
                "summary": [{"type": "summary_text", "text": "secret summary"}],
            }
        ]
    }
    assert build_inspection_messages(data) == [{"role": "assistant", "content": "secret summary"}]


# ── has_non_string_content ────────────────────────────────────────────────────


def test_has_non_string_content_string_messages():
    data = {"messages": [{"role": "user", "content": "hello"}]}
    assert has_non_string_content(data) is False


def test_has_non_string_content_multimodal_messages():
    data = {"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]}
    assert has_non_string_content(data) is True


def test_has_non_string_content_responses_api_string_input():
    assert has_non_string_content({"input": "plain string"}) is False


def test_has_non_string_content_responses_api_list_input():
    assert has_non_string_content({"input": ["a", "b"]}) is True


def test_has_non_string_content_empty_data():
    assert has_non_string_content({}) is False
    assert has_non_string_content({"messages": []}) is False
    assert has_non_string_content({"input": ""}) is False


# ── apply_redacted_messages_back ──────────────────────────────────────────────


def test_apply_redacted_messages_back_chat_completion():
    data = {"messages": [{"role": "user", "content": "secret"}]}
    apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}])
    assert data["messages"] == [{"role": "user", "content": "[REDACTED]"}]
    assert "input" not in data


def test_apply_redacted_messages_back_responses_api_string_input():
    """A Responses-API request reads ``data["input"]``; writing only to
    ``messages`` would let unredacted text reach the LLM."""
    data = {"input": "secret payload"}
    apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}])
    assert data["input"] == "[REDACTED]"


def test_apply_redacted_messages_back_both_fields():
    """Defensive: when both fields are present, both are updated."""
    data = {
        "messages": [{"role": "user", "content": "old"}],
        "input": "old",
    }
    apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}])
    assert data["messages"] == [{"role": "user", "content": "[REDACTED]"}]
    assert data["input"] == "[REDACTED]"


def test_apply_redacted_messages_back_skips_input_when_not_string():
    """List ``input`` (multimodal Responses-API) is left alone — the
    multimodal-degrades-to-block guard runs upstream."""
    data = {"input": [{"type": "text", "text": "leak"}]}
    apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}])
    assert data["input"] == [{"type": "text", "text": "leak"}]


def test_apply_redacted_messages_back_rewrites_string_batches():
    """An /embeddings batch is a list of plain strings; each is rewritten in place
    from the matching redacted message so no element reaches the LLM unredacted."""
    data = {"input": ["first SSN", "second SSN"]}
    apply_redacted_messages_back(
        data,
        [
            {"role": "user", "content": "first [REDACTED]"},
            {"role": "user", "content": "second [REDACTED]"},
        ],
    )
    assert data["input"] == ["first [REDACTED]", "second [REDACTED]"]


def test_apply_redacted_messages_back_keeps_batch_elements_aligned():
    """A guardrail that redacts a whole element away returns it as empty text.
    Each element still has to take its own redaction, never the next one's."""
    data = {"input": ["all secret", "second doc", "third doc"]}
    apply_redacted_messages_back(
        data,
        [
            {"role": "user", "content": ""},
            {"role": "user", "content": "second doc"},
            {"role": "user", "content": "third doc"},
        ],
    )
    assert data["input"] == ["", "second doc", "third doc"]


def test_apply_redacted_messages_back_skips_empty_batch_elements():
    """Empty elements are never sent to the guardrail, so the redactions line up
    with the elements that were."""
    data = {"input": ["", "secret doc"]}
    assert apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED] doc"}]) is True
    assert data["input"] == ["", "[REDACTED] doc"]


def test_apply_redacted_messages_back_rejects_short_batch_response():
    """A guardrail that returns fewer messages than were inspected cannot be
    applied element-wise: writing the prefix would forward the rest of the batch
    unredacted, so nothing is written and the caller has to block."""
    data = {"input": ["first SSN", "second SSN", "third SSN"]}
    assert apply_redacted_messages_back(data, [{"role": "user", "content": "first [REDACTED]"}]) is False
    assert data["input"] == ["first SSN", "second SSN", "third SSN"]


def test_apply_redacted_messages_back_rejects_long_batch_response():
    """More redactions than inspected elements means the alignment is unknown."""
    data = {"input": ["only SSN"]}
    assert (
        apply_redacted_messages_back(
            data,
            [
                {"role": "user", "content": "only [REDACTED]"},
                {"role": "user", "content": "spurious"},
            ],
        )
        is False
    )
    assert data["input"] == ["only SSN"]


def test_apply_redacted_messages_back_rejects_batch_content_missing():
    """A message without content cannot safely replace the original batch element."""
    data = {"input": ["secret doc"]}
    assert apply_redacted_messages_back(data, [{"role": "user"}]) is False
    assert data["input"] == ["secret doc"]


def test_apply_redacted_messages_back_returns_true_for_non_batch_shapes():
    data = {"messages": [{"role": "user", "content": "secret"}]}
    assert apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) is True


# ── is_string_batch_input ─────────────────────────────────────────────────────


def test_is_string_batch_input_embeddings_batch():
    assert is_string_batch_input({"input": ["a", "b"]}) is True


def test_is_string_batch_input_rejects_other_shapes():
    assert is_string_batch_input({"input": "a"}) is False
    assert is_string_batch_input({"input": []}) is False
    assert is_string_batch_input({"input": [1, 2]}) is False
    assert is_string_batch_input({"input": ["a", {"type": "text", "text": "b"}]}) is False
    assert is_string_batch_input({"messages": [], "input": ["a"]}) is False


# -------------------------------------------------------------------
# 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 = {
        "input": [
            {"type": "custom_tool_call_output", "output": "tool-secret"},
        ]
    }
    from litellm.proxy.guardrails._content_utils import iter_message_text

    texts = list(iter_message_text(data))
    assert "tool-secret" in texts


def test_walk_user_text_redacts_custom_tool_call_output():
    """walk_user_text should rewrite text inside custom_tool_call_output."""
    data = {
        "input": [
            {"type": "custom_tool_call_output", "output": "PII-data"},
        ]
    }
    count = walk_user_text(data, lambda t: t.replace("PII-data", "[MASKED]"))
    assert count >= 1
    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 = {
        "input": [
            {"type": "custom_tool_call_output", "output": "custom-tool-leak"},
        ]
    }
    msgs = build_inspection_messages(data)
    assert any("custom-tool-leak" in m["content"] for m in msgs)


# ── is_non_conversational_call_type ──────────────────────────────────────────────


def test_is_non_conversational_call_type_flags_embeddings():
    """An /embeddings body carries documents being indexed, not a prompt."""
    assert is_non_conversational_call_type("embedding") is True
    assert is_non_conversational_call_type("aembedding") is True


def test_is_non_conversational_call_type_passes_every_conversational_call_type():
    """Deliberately a deny-list: ``anthropic_messages``, ``responses`` and
    ``call_mcp_tool`` carry conversations but are absent from
    ``TEXT_CONTENT_CALL_TYPES``, so a guardrail gating on that allow-list would
    stop inspecting them."""
    for call_type in (
        "completion",
        "acompletion",
        "text_completion",
        "responses",
        "aresponses",
        "anthropic_messages",
        "aanthropic_messages",
        "call_mcp_tool",
    ):
        assert is_non_conversational_call_type(call_type) is False


def test_is_non_conversational_call_type_defaults_to_inspecting_unknown_call_types():
    """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(
    ("call_type", "endpoint", "expected"),
    [
        ("aembedding", None, True),
        ("allm_passthrough_route", "embeddings", True),
        ("allm_passthrough_route", "/v1/embeddings/", True),
        ("allm_passthrough_route", "chat/completions", False),
        ("allm_passthrough_route", "model/amazon.titan-embed-text-v2:0/invoke", False),
        ("allm_passthrough_route", "fine_tuned_embeddings_jobs", False),
        ("allm_passthrough_route", "chat/completions#embeddings", False),
        ("allm_passthrough_route", "chat/completions?x=embeddings", False),
        ("allm_passthrough_route", "chat/completions;embeddings", False),
        ("allm_passthrough_route", None, False),
        ("acompletion", "embeddings", False),
    ],
)
def test_is_non_conversational_request_skips_only_embeddings_requests(
    call_type: str, endpoint: str | None, expected: bool
):
    data: Final[dict[str, object]] = {"endpoint": endpoint} if endpoint is not None else {}

    assert is_non_conversational_request(call_type, data) is expected


@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)


def test_skip_tool_message_leaves_tool_results_and_masks_the_rest() -> None:
    data: Final[dict[str, object]] = {
        "messages": [
            {"role": "user", "content": "secret"},
            {
                "role": "assistant",
                "content": [{"type": "tool_use", "id": "tool", "name": "read", "input": {"env": "secret"}}],
                "tool_calls": [{"id": "call", "function": {"name": "read", "arguments": '{"env": "secret"}'}}],
            },
            {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tool", "content": "secret"}]},
            {"role": "tool", "tool_call_id": "call", "content": "secret"},
        ]
    }
    assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]"), skip_tool_message=True) == 3
    assert data == {
        "messages": [
            {"role": "user", "content": "[MASKED]"},
            {
                "role": "assistant",
                "content": [{"type": "tool_use", "id": "tool", "name": "read", "input": {"env": "[MASKED]"}}],
                "tool_calls": [{"id": "call", "function": {"name": "read", "arguments": '{"env": "[MASKED]"}'}}],
            },
            {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tool", "content": "secret"}]},
            {"role": "tool", "tool_call_id": "call", "content": "secret"},
        ]
    }


@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)