"""
Unit tests for OpenAI Chat Completions Guardrail Translation Handler

Tests the handler's ability to process input/output for Chat Completions API
with guardrail transformations, including tool calls.
"""

import json
from typing import Any, Literal, Optional

import pytest


from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.openai.chat.guardrail_translation.handler import (
    OpenAIChatCompletionsHandler,
)
from litellm.types.utils import (
    ChatCompletionMessageToolCall,
    Choices,
    Function,
    GenericGuardrailAPIInputs,
    Message,
    ModelResponse,
)


class MockGuardrail(CustomGuardrail):
    """Mock guardrail for testing that transforms text and tool calls"""

    def __init__(self, guardrail_name: str = "test"):
        super().__init__(guardrail_name=guardrail_name)
        self.last_inputs = None
        self.last_request_data = None
        self.tool_calls_modified = False

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        """Mock apply_guardrail that uppercases text and modifies tool calls"""
        self.last_inputs = inputs
        self.last_request_data = request_data

        # Return modified inputs (uppercase texts for testing)
        texts = inputs.get("texts", [])
        modified_texts = [text.upper() for text in texts]

        # Modify tool calls in place if present
        tool_calls = inputs.get("tool_calls", [])
        if tool_calls:
            self.tool_calls_modified = True
            for tool_call in tool_calls:
                if isinstance(tool_call, dict) and "function" in tool_call:
                    function = tool_call["function"]
                    if "arguments" in function:
                        # Modify arguments to uppercase JSON string
                        try:
                            args_dict = json.loads(function["arguments"])
                            # Uppercase all string values
                            for key, value in args_dict.items():
                                if isinstance(value, str):
                                    args_dict[key] = value.upper()
                            function["arguments"] = json.dumps(args_dict)
                        except json.JSONDecodeError:
                            # If not JSON, just uppercase the string
                            function["arguments"] = function["arguments"].upper()

        # Return modified inputs as GenericGuardrailAPIInputs
        result: GenericGuardrailAPIInputs = {"texts": modified_texts}
        if tool_calls:
            result["tool_calls"] = tool_calls  # type: ignore
        if "images" in inputs:
            result["images"] = inputs["images"]  # type: ignore
        return result


class MockCopiedToolCallGuardrail(CustomGuardrail):
    """Mock guardrail that returns copied tool calls instead of mutating inputs."""

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        tool_calls = inputs.get("tool_calls", [])
        copied_tool_calls = []
        for tool_call in tool_calls:
            copied = dict(tool_call)
            function = dict(copied["function"])
            function["arguments"] = json.dumps({"email": "[EMAIL]"})
            copied["function"] = function
            copied_tool_calls.append(copied)

        return GenericGuardrailAPIInputs(
            texts=inputs.get("texts", []),
            tool_calls=copied_tool_calls,
        )


class MockNonListToolCallGuardrail(CustomGuardrail):
    """Mock guardrail that returns tool_calls as a non-list envelope on the response
    path, as some released guardrails do when they assign a detection API JSON dict."""

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        result = GenericGuardrailAPIInputs(texts=inputs.get("texts", []))
        result["tool_calls"] = {"verdict": "allow", "detections": []}  # type: ignore
        return result


class MockMisalignedToolCallGuardrail(CustomGuardrail):
    """Mock guardrail that returns a tool_calls list whose length differs from the
    input, so it cannot be applied positionally onto the response."""

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        tool_calls = inputs.get("tool_calls", [])
        shortened = []
        if tool_calls:
            first = dict(tool_calls[0])
            first["function"] = {"name": "x", "arguments": json.dumps({"x": 1})}
            shortened.append(first)
        return GenericGuardrailAPIInputs(
            texts=inputs.get("texts", []),
            tool_calls=shortened,
        )


class TestOpenAIChatCompletionsHandlerToolsInput:
    """Test input processing with tools (function definitions)"""

    @pytest.mark.asyncio
    async def test_tools_passed_to_guardrail(self):
        """Test that tools (function definitions) are passed to the guardrail"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        # Create input data with tools (function definitions)
        data = {
            "messages": [
                {"role": "user", "content": "What's the weather in Boston?"},
            ],
            "tools": [
                {
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "description": "Get the current weather in a location",
                        "parameters": {
                            "type": "object",
                            "properties": {
                                "location": {
                                    "type": "string",
                                    "description": "City name",
                                },
                                "unit": {
                                    "type": "string",
                                    "enum": ["celsius", "fahrenheit"],
                                },
                            },
                            "required": ["location"],
                        },
                    },
                }
            ],
        }

        # Process the input
        await handler.process_input_messages(data, guardrail)

        # Verify tools were passed to guardrail
        assert guardrail.last_inputs is not None
        assert "tools" in guardrail.last_inputs
        assert len(guardrail.last_inputs["tools"]) == 1

        tool = guardrail.last_inputs["tools"][0]
        assert tool["type"] == "function"
        assert tool["function"]["name"] == "get_weather"
        assert (
            tool["function"]["description"] == "Get the current weather in a location"
        )
        assert "parameters" in tool["function"]

    @pytest.mark.asyncio
    async def test_multiple_tools_passed_to_guardrail(self):
        """Test that multiple tools are passed to the guardrail"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        data = {
            "messages": [
                {"role": "user", "content": "What's the weather and time?"},
            ],
            "tools": [
                {
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "parameters": {"type": "object", "properties": {}},
                    },
                },
                {
                    "type": "function",
                    "function": {
                        "name": "get_time",
                        "parameters": {"type": "object", "properties": {}},
                    },
                },
            ],
        }

        await handler.process_input_messages(data, guardrail)

        assert guardrail.last_inputs is not None
        assert "tools" in guardrail.last_inputs
        assert len(guardrail.last_inputs["tools"]) == 2
        assert guardrail.last_inputs["tools"][0]["function"]["name"] == "get_weather"
        assert guardrail.last_inputs["tools"][1]["function"]["name"] == "get_time"

    @pytest.mark.asyncio
    async def test_no_tools_in_request(self):
        """Test that requests without tools work correctly"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        data = {
            "messages": [
                {"role": "user", "content": "Hello"},
            ],
        }

        await handler.process_input_messages(data, guardrail)

        assert guardrail.last_inputs is not None
        # tools should not be in inputs if not provided
        assert (
            "tools" not in guardrail.last_inputs
            or guardrail.last_inputs.get("tools") is None
        )

    @pytest.mark.asyncio
    async def test_tools_and_tool_calls_both_passed(self):
        """Test that both tools (definitions) and tool_calls (invocations) are passed"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        data = {
            "messages": [
                {"role": "user", "content": "What's the weather?"},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "call_123",
                            "type": "function",
                            "function": {
                                "name": "get_weather",
                                "arguments": '{"location": "Boston"}',
                            },
                        }
                    ],
                },
            ],
            "tools": [
                {
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "parameters": {
                            "type": "object",
                            "properties": {"location": {"type": "string"}},
                        },
                    },
                }
            ],
        }

        await handler.process_input_messages(data, guardrail)

        assert guardrail.last_inputs is not None
        # Both should be present
        assert "tools" in guardrail.last_inputs
        assert "tool_calls" in guardrail.last_inputs
        assert len(guardrail.last_inputs["tools"]) == 1
        assert len(guardrail.last_inputs["tool_calls"]) == 1


class TestOpenAIChatCompletionsHandlerToolCallsInput:
    """Test input processing with tool calls"""

    @pytest.mark.asyncio
    async def test_extract_tool_calls_from_input_messages(self):
        """Test that tool calls are extracted from input messages"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        # Create input data with tool calls (assistant message)
        data = {
            "messages": [
                {"role": "user", "content": "What's the weather in SF?"},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "call_123",
                            "type": "function",
                            "function": {
                                "name": "get_weather",
                                "arguments": json.dumps(
                                    {"location": "San Francisco", "unit": "celsius"}
                                ),
                            },
                        }
                    ],
                },
            ]
        }

        # Process the input
        await handler.process_input_messages(data, guardrail)

        # Verify tool calls were extracted and passed to guardrail
        assert guardrail.last_inputs is not None
        assert "tool_calls" in guardrail.last_inputs
        assert len(guardrail.last_inputs["tool_calls"]) == 1

        tool_call = guardrail.last_inputs["tool_calls"][0]
        assert tool_call["id"] == "call_123"
        assert tool_call["function"]["name"] == "get_weather"
        # Note: tool call arguments may already be modified by guardrail
        # Check that it contains location parameter
        assert "location" in tool_call["function"]["arguments"]

        # Verify tool call was modified by guardrail
        assert guardrail.tool_calls_modified is True

        # Verify the message was updated with modified tool call
        modified_tool_call = data["messages"][1]["tool_calls"][0]
        args = json.loads(modified_tool_call["function"]["arguments"])
        assert args["location"] == "SAN FRANCISCO"  # Should be uppercased
        assert args["unit"] == "CELSIUS"  # Should be uppercased

    @pytest.mark.asyncio
    async def test_extract_tool_calls_and_text_from_input_messages(self):
        """Test that both tool calls and text content are extracted"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        # Create input data with both text and tool calls
        data = {
            "messages": [
                {"role": "user", "content": "What's the weather?"},
                {
                    "role": "assistant",
                    "content": "Let me check that for you.",
                    "tool_calls": [
                        {
                            "id": "call_456",
                            "type": "function",
                            "function": {
                                "name": "get_current_weather",
                                "arguments": json.dumps({"location": "Boston"}),
                            },
                        }
                    ],
                },
            ]
        }

        # Process the input
        await handler.process_input_messages(data, guardrail)

        # Verify both texts and tool calls were extracted
        assert guardrail.last_inputs is not None
        assert "texts" in guardrail.last_inputs
        assert "tool_calls" in guardrail.last_inputs

        # Should have 2 texts (user message + assistant message)
        assert len(guardrail.last_inputs["texts"]) == 2
        assert "What's the weather?" in guardrail.last_inputs["texts"]
        assert "Let me check that for you." in guardrail.last_inputs["texts"]

        # Should have 1 tool call
        assert len(guardrail.last_inputs["tool_calls"]) == 1
        assert (
            guardrail.last_inputs["tool_calls"][0]["function"]["name"]
            == "get_current_weather"
        )

        # Verify text content was modified
        assert data["messages"][0]["content"] == "WHAT'S THE WEATHER?"
        assert data["messages"][1]["content"] == "LET ME CHECK THAT FOR YOU."

        # Verify tool call was modified
        modified_tool_call = data["messages"][1]["tool_calls"][0]
        args = json.loads(modified_tool_call["function"]["arguments"])
        assert args["location"] == "BOSTON"

    @pytest.mark.asyncio
    async def test_extract_multiple_tool_calls_from_input(self):
        """Test extraction of multiple tool calls"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        # Create input data with multiple tool calls
        data = {
            "messages": [
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "call_1",
                            "type": "function",
                            "function": {
                                "name": "get_weather",
                                "arguments": json.dumps({"location": "NYC"}),
                            },
                        },
                        {
                            "id": "call_2",
                            "type": "function",
                            "function": {
                                "name": "get_time",
                                "arguments": json.dumps({"timezone": "EST"}),
                            },
                        },
                    ],
                }
            ]
        }

        # Process the input
        await handler.process_input_messages(data, guardrail)

        # Verify multiple tool calls were extracted
        assert guardrail.last_inputs is not None
        assert "tool_calls" in guardrail.last_inputs
        assert len(guardrail.last_inputs["tool_calls"]) == 2

        # Verify both tool calls
        tool_calls = guardrail.last_inputs["tool_calls"]
        assert tool_calls[0]["function"]["name"] == "get_weather"
        assert tool_calls[1]["function"]["name"] == "get_time"

        # Verify both were modified
        modified_tool_calls = data["messages"][0]["tool_calls"]
        args1 = json.loads(modified_tool_calls[0]["function"]["arguments"])
        args2 = json.loads(modified_tool_calls[1]["function"]["arguments"])
        assert args1["location"] == "NYC"
        assert args2["timezone"] == "EST"

    @pytest.mark.asyncio
    async def test_tool_calls_separate_from_texts(self):
        """Test that tool calls are passed as a separate parameter, not mixed with texts"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        data = {
            "messages": [
                {"role": "user", "content": "Get weather for LA"},
                {
                    "role": "assistant",
                    "content": "Sure!",
                    "tool_calls": [
                        {
                            "id": "call_xyz",
                            "type": "function",
                            "function": {
                                "name": "get_weather",
                                "arguments": json.dumps({"city": "Los Angeles"}),
                            },
                        }
                    ],
                },
            ]
        }

        # Process the input
        await handler.process_input_messages(data, guardrail)

        # Verify tool calls and texts are separate
        assert guardrail.last_inputs is not None
        texts = guardrail.last_inputs.get("texts", [])
        tool_calls = guardrail.last_inputs.get("tool_calls", [])

        # Texts should only contain the content strings
        assert len(texts) == 2
        assert "Get weather for LA" in texts
        assert "Sure!" in texts

        # Tool call arguments should NOT be in texts
        assert not any("Los Angeles" in text for text in texts)

        # Tool calls should be separate
        assert len(tool_calls) == 1
        assert tool_calls[0]["function"]["name"] == "get_weather"
        # Check that it contains city parameter (may be modified by guardrail)
        assert "city" in tool_calls[0]["function"]["arguments"]

    @pytest.mark.asyncio
    async def test_no_tool_calls_in_input(self):
        """Test that messages without tool calls work correctly"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        # Create input data without tool calls
        data = {
            "messages": [
                {"role": "user", "content": "Hello"},
                {"role": "assistant", "content": "Hi there!"},
            ]
        }

        # Process the input
        await handler.process_input_messages(data, guardrail)

        # Verify no tool calls were passed to guardrail
        assert guardrail.last_inputs is not None
        tool_calls = guardrail.last_inputs.get("tool_calls", [])
        assert len(tool_calls) == 0

        # Verify text was still processed
        assert len(guardrail.last_inputs["texts"]) == 2
        assert data["messages"][0]["content"] == "HELLO"
        assert data["messages"][1]["content"] == "HI THERE!"

    @pytest.mark.asyncio
    async def test_empty_tool_calls_list(self):
        """Test that empty tool_calls list is handled correctly"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        data = {
            "messages": [
                {"role": "assistant", "content": "Hello", "tool_calls": []},
            ]
        }

        # Process the input
        await handler.process_input_messages(data, guardrail)

        # Verify empty tool_calls doesn't cause issues
        assert guardrail.last_inputs is not None
        tool_calls = guardrail.last_inputs.get("tool_calls", [])
        assert len(tool_calls) == 0


class TestOpenAIChatCompletionsHandlerToolCallsOutput:
    """Test output processing with tool calls"""

    @pytest.mark.asyncio
    async def test_extract_tool_calls_from_output_response(self):
        """Test that tool calls are extracted from output responses"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        # Create a mock response with tool calls
        response = ModelResponse(
            id="chatcmpl-123",
            created=1234567890,
            model="gpt-4",
            object="chat.completion",
            choices=[
                Choices(
                    finish_reason="tool_calls",
                    index=0,
                    message=Message(
                        content=None,
                        role="assistant",
                        tool_calls=[
                            ChatCompletionMessageToolCall(
                                id="call_789",
                                type="function",
                                function=Function(
                                    name="search_database",
                                    arguments=json.dumps({"query": "python tutorials"}),
                                ),
                            )
                        ],
                    ),
                )
            ],
        )

        # Process the output
        await handler.process_output_response(response, guardrail)

        # Verify tool calls were extracted and passed to guardrail
        assert guardrail.last_inputs is not None
        assert "tool_calls" in guardrail.last_inputs
        assert len(guardrail.last_inputs["tool_calls"]) == 1

        tool_call = guardrail.last_inputs["tool_calls"][0]
        assert tool_call["function"]["name"] == "search_database"
        # Check that it contains query parameter (may be modified by guardrail)
        assert "query" in tool_call["function"]["arguments"]

        # Verify tool call was modified in response
        response_tool_call = response.choices[0].message.tool_calls[0]
        args = json.loads(response_tool_call.function.arguments)
        assert args["query"] == "PYTHON TUTORIALS"  # Should be uppercased

    @pytest.mark.asyncio
    async def test_extract_tool_calls_and_content_from_output(self):
        """Test extraction of both content and tool calls from output"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        response = ModelResponse(
            id="chatcmpl-456",
            created=1234567890,
            model="gpt-4",
            object="chat.completion",
            choices=[
                Choices(
                    finish_reason="tool_calls",
                    index=0,
                    message=Message(
                        content="I'll search for that information.",
                        role="assistant",
                        tool_calls=[
                            ChatCompletionMessageToolCall(
                                id="call_999",
                                type="function",
                                function=Function(
                                    name="web_search",
                                    arguments=json.dumps(
                                        {"keywords": "litellm documentation"}
                                    ),
                                ),
                            )
                        ],
                    ),
                )
            ],
        )

        # Process the output
        await handler.process_output_response(response, guardrail)

        # Verify both texts and tool calls were extracted
        assert guardrail.last_inputs is not None
        assert "texts" in guardrail.last_inputs
        assert "tool_calls" in guardrail.last_inputs

        assert len(guardrail.last_inputs["texts"]) == 1
        assert "I'll search for that information." in guardrail.last_inputs["texts"]

        assert len(guardrail.last_inputs["tool_calls"]) == 1

        # Verify both were modified
        assert (
            response.choices[0].message.content == "I'LL SEARCH FOR THAT INFORMATION."
        )
        response_tool_call = response.choices[0].message.tool_calls[0]
        args = json.loads(response_tool_call.function.arguments)
        assert args["keywords"] == "LITELLM DOCUMENTATION"

    @pytest.mark.asyncio
    async def test_extract_multiple_tool_calls_from_output(self):
        """Test extraction of multiple tool calls from output"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        response = ModelResponse(
            id="chatcmpl-789",
            created=1234567890,
            model="gpt-4",
            object="chat.completion",
            choices=[
                Choices(
                    finish_reason="tool_calls",
                    index=0,
                    message=Message(
                        content=None,
                        role="assistant",
                        tool_calls=[
                            ChatCompletionMessageToolCall(
                                id="call_1",
                                type="function",
                                function=Function(
                                    name="get_weather",
                                    arguments=json.dumps({"location": "Tokyo"}),
                                ),
                            ),
                            ChatCompletionMessageToolCall(
                                id="call_2",
                                type="function",
                                function=Function(
                                    name="get_news",
                                    arguments=json.dumps({"topic": "technology"}),
                                ),
                            ),
                        ],
                    ),
                )
            ],
        )

        # Process the output
        await handler.process_output_response(response, guardrail)

        # Verify multiple tool calls were extracted
        assert guardrail.last_inputs is not None
        assert "tool_calls" in guardrail.last_inputs
        assert len(guardrail.last_inputs["tool_calls"]) == 2

        # Verify both tool calls
        tool_calls = guardrail.last_inputs["tool_calls"]
        assert tool_calls[0]["function"]["name"] == "get_weather"
        assert tool_calls[1]["function"]["name"] == "get_news"

        # Verify both were modified
        response_tool_calls = response.choices[0].message.tool_calls
        args1 = json.loads(response_tool_calls[0].function.arguments)
        args2 = json.loads(response_tool_calls[1].function.arguments)
        assert args1["location"] == "TOKYO"
        assert args2["topic"] == "TECHNOLOGY"

    @pytest.mark.asyncio
    async def test_extract_tool_calls_from_real_openai_response(self):
        """Test extraction of tool calls from a real OpenAI API response structure"""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockGuardrail()

        # Create a response matching the exact structure from OpenAI API
        response = ModelResponse(
            id="chatcmpl-abc123",
            created=1699896916,
            model="gpt-4o-mini",
            object="chat.completion",
            choices=[
                Choices(
                    finish_reason="tool_calls",
                    index=0,
                    message=Message(
                        content=None,
                        role="assistant",
                        tool_calls=[
                            ChatCompletionMessageToolCall(
                                id="call_abc123",
                                type="function",
                                function=Function(
                                    name="get_current_weather",
                                    arguments='{\n"location": "Boston, MA"\n}',
                                ),
                            )
                        ],
                    ),
                )
            ],
        )

        # Process the output
        await handler.process_output_response(response, guardrail)

        # Verify tool calls were extracted and passed to guardrail
        assert guardrail.last_inputs is not None
        assert "tool_calls" in guardrail.last_inputs
        assert len(guardrail.last_inputs["tool_calls"]) == 1

        # Verify the tool call details
        tool_call = guardrail.last_inputs["tool_calls"][0]
        assert tool_call["id"] == "call_abc123"
        assert tool_call["type"] == "function"
        assert tool_call["function"]["name"] == "get_current_weather"

        # Verify arguments can be parsed
        args = json.loads(tool_call["function"]["arguments"])
        assert "location" in args

        # Verify tool call was modified by guardrail (location should be uppercased)
        response_tool_call = response.choices[0].message.tool_calls[0]
        modified_args = json.loads(response_tool_call.function.arguments)
        assert modified_args["location"] == "BOSTON, MA"  # Should be uppercased

        # Verify response metadata
        assert response.id == "chatcmpl-abc123"
        assert response.model == "gpt-4o-mini"
        assert response.choices[0].finish_reason == "tool_calls"

    @pytest.mark.asyncio
    async def test_output_response_uses_returned_guardrailed_tool_calls(self):
        """Test returned tool_calls are remapped even when guardrail does not mutate inputs."""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockCopiedToolCallGuardrail(guardrail_name="test")

        response = ModelResponse(
            id="chatcmpl-tool-copy",
            created=1234567890,
            model="gpt-4",
            object="chat.completion",
            choices=[
                Choices(
                    finish_reason="tool_calls",
                    index=0,
                    message=Message(
                        content=None,
                        role="assistant",
                        tool_calls=[
                            ChatCompletionMessageToolCall(
                                id="call_email",
                                type="function",
                                function=Function(
                                    name="send_email",
                                    arguments=json.dumps({"email": "john@example.com"}),
                                ),
                            )
                        ],
                    ),
                )
            ],
        )

        await handler.process_output_response(response, guardrail)

        response_tool_call = response.choices[0].message.tool_calls[0]
        assert response_tool_call.function.name == "send_email"
        assert json.loads(response_tool_call.function.arguments) == {"email": "[EMAIL]"}

    @pytest.mark.asyncio
    async def test_output_response_ignores_non_list_returned_tool_calls(self):
        """A guardrail returning tool_calls as a non-list (e.g. a detection-API envelope
        dict) must not crash the remap; the original arguments are preserved."""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockNonListToolCallGuardrail(guardrail_name="test")
        original = json.dumps({"email": "john@example.com"})
        response = ModelResponse(
            id="chatcmpl-nonlist",
            created=1234567890,
            model="gpt-4",
            object="chat.completion",
            choices=[
                Choices(
                    finish_reason="tool_calls",
                    index=0,
                    message=Message(
                        content=None,
                        role="assistant",
                        tool_calls=[
                            ChatCompletionMessageToolCall(
                                id="call_email",
                                type="function",
                                function=Function(
                                    name="send_email", arguments=original
                                ),
                            )
                        ],
                    ),
                )
            ],
        )

        await handler.process_output_response(response, guardrail)

        response_tool_call = response.choices[0].message.tool_calls[0]
        assert response_tool_call.function.arguments == original

    @pytest.mark.asyncio
    async def test_output_response_ignores_misaligned_returned_tool_calls(self):
        """A guardrail returning a tool_calls list of a different length than the input
        cannot be applied positionally; the handler falls back and preserves the
        original arguments instead of writing onto the wrong tool call."""
        handler = OpenAIChatCompletionsHandler()
        guardrail = MockMisalignedToolCallGuardrail(guardrail_name="test")
        first_args = json.dumps({"email": "a@example.com"})
        second_args = json.dumps({"email": "b@example.com"})
        response = ModelResponse(
            id="chatcmpl-misaligned",
            created=1234567890,
            model="gpt-4",
            object="chat.completion",
            choices=[
                Choices(
                    finish_reason="tool_calls",
                    index=0,
                    message=Message(
                        content=None,
                        role="assistant",
                        tool_calls=[
                            ChatCompletionMessageToolCall(
                                id="call_1",
                                type="function",
                                function=Function(
                                    name="send_email", arguments=first_args
                                ),
                            ),
                            ChatCompletionMessageToolCall(
                                id="call_2",
                                type="function",
                                function=Function(
                                    name="send_email", arguments=second_args
                                ),
                            ),
                        ],
                    ),
                )
            ],
        )

        await handler.process_output_response(response, guardrail)

        tool_calls = response.choices[0].message.tool_calls
        assert tool_calls[0].function.arguments == first_args
        assert tool_calls[1].function.arguments == second_args


class MockPassThroughGuardrail(CustomGuardrail):
    """Mock guardrail that passes through without blocking - for testing streaming fallback behavior"""

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        """Simply return inputs unchanged"""
        return inputs


class TestOpenAIChatCompletionsHandlerStreamingOutput:
    """Test streaming output processing functionality"""

    @pytest.mark.asyncio
    async def test_process_output_streaming_response_empty_choices(self):
        """Test that streaming response with empty choices doesn't raise IndexError

        This test verifies the fix for the bug where accessing chunk.choices[0]
        would raise IndexError when a streaming chunk has an empty choices list.
        """
        from litellm.types.utils import ModelResponseStream

        handler = OpenAIChatCompletionsHandler()
        guardrail = MockPassThroughGuardrail(guardrail_name="test")

        # Create a streaming chunk with empty choices
        chunk_with_empty_choices = ModelResponseStream(
            id="chatcmpl-123",
            created=1234567890,
            model="gpt-4",
            object="chat.completion.chunk",
            choices=[],  # Empty choices - this was causing the IndexError
        )

        responses_so_far = [chunk_with_empty_choices]

        # This should not raise IndexError
        result = await handler.process_output_streaming_response(
            responses_so_far=responses_so_far,
            guardrail_to_apply=guardrail,
            litellm_logging_obj=None,
        )

        # Should return the responses unchanged
        assert result == responses_so_far

    @pytest.mark.asyncio
    async def test_process_output_streaming_response_with_valid_choices(self):
        """Test that streaming response with valid choices still works correctly"""
        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices

        handler = OpenAIChatCompletionsHandler()
        guardrail = MockPassThroughGuardrail(guardrail_name="test")

        # Create streaming chunks with valid choices
        chunk1 = ModelResponseStream(
            id="chatcmpl-123",
            created=1234567890,
            model="gpt-4",
            object="chat.completion.chunk",
            choices=[
                StreamingChoices(
                    index=0,
                    delta=Delta(content="Hello"),
                    finish_reason=None,
                )
            ],
        )

        chunk2 = ModelResponseStream(
            id="chatcmpl-123",
            created=1234567890,
            model="gpt-4",
            object="chat.completion.chunk",
            choices=[
                StreamingChoices(
                    index=0,
                    delta=Delta(content=" world"),
                    finish_reason="stop",
                )
            ],
        )

        responses_so_far = [chunk1, chunk2]

        # This should process successfully
        result = await handler.process_output_streaming_response(
            responses_so_far=responses_so_far,
            guardrail_to_apply=guardrail,
            litellm_logging_obj=None,
        )

        # Should return the responses
        assert result == responses_so_far

    @pytest.mark.asyncio
    async def test_process_output_streaming_response_mixed_empty_and_valid_choices_no_finish(
        self,
    ):
        """Test streaming response with mix of empty and valid choices chunks (stream not finished)

        This tests the has_stream_ended check when iterating through chunks with mixed choices.
        The stream hasn't finished yet (no finish_reason), so it won't trigger stream_chunk_builder.
        """
        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices

        handler = OpenAIChatCompletionsHandler()
        guardrail = MockPassThroughGuardrail(guardrail_name="test")

        # Mix of chunks - some with empty choices, some with valid choices
        # Stream hasn't finished (no finish_reason)
        chunk_empty = ModelResponseStream(
            id="chatcmpl-123",
            created=1234567890,
            model="gpt-4",
            object="chat.completion.chunk",
            choices=[],
        )

        chunk_valid = ModelResponseStream(
            id="chatcmpl-123",
            created=1234567890,
            model="gpt-4",
            object="chat.completion.chunk",
            choices=[
                StreamingChoices(
                    index=0,
                    delta=Delta(content="Hello"),
                    finish_reason=None,  # Stream not finished
                )
            ],
        )

        responses_so_far = [chunk_empty, chunk_valid]

        # This should not raise IndexError when checking has_stream_ended
        result = await handler.process_output_streaming_response(
            responses_so_far=responses_so_far,
            guardrail_to_apply=guardrail,
            litellm_logging_obj=None,
        )

        # Should return the responses
        assert result == responses_so_far


class TestGetStructuredMessages:
    """Test the get_structured_messages method."""

    def test_should_return_messages_from_chat_completions_request(self):
        """Test that messages are returned from a chat completions request."""
        handler = OpenAIChatCompletionsHandler()
        data = {
            "messages": [
                {"role": "system", "content": "You are helpful."},
                {"role": "user", "content": "Hello"},
            ]
        }
        result = handler.get_structured_messages(data)
        assert result is not None
        assert len(result) == 2
        assert result[0]["role"] == "system"
        assert result[1]["role"] == "user"

    def test_should_return_none_when_no_messages(self):
        """Test that None is returned when no messages key exists."""
        handler = OpenAIChatCompletionsHandler()
        data = {"model": "gpt-4"}
        result = handler.get_structured_messages(data)
        assert result is None

    def test_should_return_none_for_none_messages(self):
        """Test that None is returned when messages is explicitly None."""
        handler = OpenAIChatCompletionsHandler()
        data = {"messages": None}
        result = handler.get_structured_messages(data)
        assert result is None

    def test_should_handle_multimodal_content(self):
        """Test that messages with multimodal content are returned."""
        handler = OpenAIChatCompletionsHandler()
        data = {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "What's in this image?"},
                        {
                            "type": "image_url",
                            "image_url": {"url": "https://example.com/image.png"},
                        },
                    ],
                }
            ]
        }
        result = handler.get_structured_messages(data)
        assert result is not None
        assert len(result) == 1
        assert isinstance(result[0]["content"], list)


if __name__ == "__main__":
    # Run the tests
    pytest.main([__file__, "-v"])


class TestIncrementalScanRespectsSkipFlags:
    """PR #33278: skip_system_message_in_guardrail and skip_tool_message_in_guardrail
    are enforced while this handler builds inputs["texts"] (_extract_inputs early
    returns for system/tool roles), upstream of BedrockGuardrail's incremental path.
    Bypassing _select_messages_for_apply_guardrail therefore cannot resurrect skipped
    content on any turn, including a session's first turn where every segment is new.
    Verified live against a real Bedrock ApplyGuardrail before being encoded here.
    The flags are set as instance attributes, mirroring how guardrail_registry
    applies litellm_params to the callback (they are not constructor kwargs).
    """

    def _bedrock_guardrail(self):
        from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail

        guardrail = BedrockGuardrail(
            guardrail_name="bedrock-incremental-skip-flags",
            guardrailIdentifier="test-guardrail",
            guardrailVersion="DRAFT",
            default_on=True,
            only_scan_new_messages=True,
        )
        guardrail.skip_system_message_in_guardrail = True
        guardrail.skip_tool_message_in_guardrail = True
        return guardrail

    def _messages(self, followup=None):
        base = [
            {"role": "system", "content": "SYSTEM-PROMPT-must-not-be-scanned"},
            {"role": "user", "content": "Search for the weather in Paris"},
            {
                "role": "assistant",
                "content": "Let me look that up.",
                "tool_calls": [
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {"name": "search", "arguments": '{"query": "weather"}'},
                    }
                ],
            },
            {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-must-not-be-scanned"},
            {"role": "user", "content": "Thanks, summarize."},
        ]
        return base + (followup or [])

    @pytest.mark.asyncio
    async def test_first_turn_scans_no_system_or_tool_content(self):
        from unittest.mock import AsyncMock, patch

        handler = OpenAIChatCompletionsHandler()
        guardrail = self._bedrock_guardrail()
        data = {"messages": self._messages(), "litellm_session_id": "skip-flags-turn1"}
        with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
            mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
            await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
            assert mock_api.call_count == 1
            scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
            assert scanned == [
                "Search for the weather in Paris",
                "Let me look that up.",
                "Thanks, summarize.",
            ]
            assert not any("SYSTEM-PROMPT" in text for text in scanned)
            assert not any("TOOL-RESULT" in text for text in scanned)

    @pytest.mark.asyncio
    async def test_second_turn_scans_only_new_eligible_content(self):
        from unittest.mock import AsyncMock, patch

        handler = OpenAIChatCompletionsHandler()
        guardrail = self._bedrock_guardrail()
        session = "skip-flags-turn2"
        followup = [
            {"role": "assistant", "content": "It is sunny in Paris."},
            {"role": "user", "content": "And tomorrow?"},
        ]
        with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
            mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
            await handler.process_input_messages(
                data={"messages": self._messages(), "litellm_session_id": session},
                guardrail_to_apply=guardrail,
            )
            mock_api.reset_mock()
            await handler.process_input_messages(
                data={"messages": self._messages(followup), "litellm_session_id": session},
                guardrail_to_apply=guardrail,
            )
            assert mock_api.call_count == 1
            scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
            assert scanned == ["It is sunny in Paris.", "And tomorrow?"]


class StructuredRedactionGuardrail(CustomGuardrail):
    """Captures inputs and returns a new structured_messages list with a canary redacted."""

    def __init__(self):
        super().__init__(guardrail_name="structured-redaction")
        self.captured_inputs: Optional[GenericGuardrailAPIInputs] = None

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        self.captured_inputs = inputs
        structured = inputs.get("structured_messages") or []
        inputs["structured_messages"] = [
            {**m, "content": str(m.get("content", "")).replace("POISON", "[BLOCKED]")} for m in structured
        ]
        return inputs


class ToolSynthesizingGuardrail(CustomGuardrail):
    """Appends its own function tool to whatever tools it was given, like a
    retrieval/recovery guardrail that injects a tool the model can later call."""

    def __init__(self):
        super().__init__(guardrail_name="tool-synthesizing")

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        tools = list(inputs.get("tools") or [])
        tools.append(
            {
                "type": "function",
                "function": {"name": "injected_retrieve", "parameters": {"type": "object", "properties": {}}},
            }
        )
        inputs["tools"] = tools
        return inputs


class ToolNameCollidingGuardrail(CustomGuardrail):
    """Returns a tool reusing a request tool's name plus a genuinely new tool."""

    def __init__(self):
        super().__init__(guardrail_name="tool-name-colliding")

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        inputs["tools"] = [
            {
                "type": "function",
                "function": {
                    "name": "read_file",
                    "parameters": {"type": "object", "properties": {"hijacked": {"type": "string"}}},
                },
            },
            {
                "type": "function",
                "function": {"name": "injected_retrieve", "parameters": {"type": "object", "properties": {}}},
            },
        ]
        return inputs


class DuplicateToolReturningGuardrail(CustomGuardrail):
    """Returns the same synthesized tool name twice, second copy with a different schema."""

    def __init__(self):
        super().__init__(guardrail_name="duplicate-tool-returning")

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        inputs["tools"] = [
            {
                "type": "function",
                "function": {
                    "name": "injected_retrieve",
                    "parameters": {"type": "object", "properties": {"first": {"type": "string"}}},
                },
            },
            {
                "type": "function",
                "function": {
                    "name": "injected_retrieve",
                    "parameters": {"type": "object", "properties": {"second": {"type": "string"}}},
                },
            },
        ]
        return inputs


class TestScanOnlyToolResults:
    def _bedrock_guardrail(self):
        from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail

        guardrail = BedrockGuardrail(
            guardrail_name="bedrock-scan-only-tool-results",
            guardrailIdentifier="test-guardrail",
            guardrailVersion="DRAFT",
            default_on=True,
        )
        guardrail.scan_only_tool_results = True
        return guardrail

    @pytest.mark.asyncio
    async def test_only_tool_role_content_is_scanned(self):
        from unittest.mock import AsyncMock, patch

        handler = OpenAIChatCompletionsHandler()
        guardrail = self._bedrock_guardrail()
        data = {
            "messages": [
                {"role": "system", "content": "SYSTEM-PROMPT-not-scanned"},
                {"role": "user", "content": "USER-PROMPT-not-scanned"},
                {
                    "role": "assistant",
                    "content": "ASSISTANT-not-scanned",
                    "tool_calls": [
                        {
                            "id": "call_1",
                            "type": "function",
                            "function": {"name": "read_file", "arguments": '{"path": "report.html"}'},
                        }
                    ],
                },
                {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-scanned"},
            ]
        }
        with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
            mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
            await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
            assert mock_api.call_count == 1
            scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
            assert scanned == ["TOOL-RESULT-scanned"]

    @pytest.mark.asyncio
    async def test_legacy_function_role_results_are_scanned(self):
        from unittest.mock import AsyncMock, patch

        handler = OpenAIChatCompletionsHandler()
        guardrail = self._bedrock_guardrail()
        data = {
            "messages": [
                {"role": "user", "content": "USER-PROMPT-not-scanned"},
                {"role": "function", "name": "read_file", "content": "FUNCTION-RESULT-scanned"},
                {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-scanned"},
            ]
        }
        with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
            mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
            await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
            assert mock_api.call_count == 1
            scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
            assert scanned == ["FUNCTION-RESULT-scanned", "TOOL-RESULT-scanned"], (
                "a tool result sent with the legacy function role must not bypass the scoped scan"
            )

    @pytest.mark.parametrize("flag_value", [None, "false", 0, object()])
    @pytest.mark.asyncio
    async def test_scope_narrows_only_when_the_flag_is_actually_true(self, flag_value):
        from unittest.mock import AsyncMock, patch

        handler = OpenAIChatCompletionsHandler()
        guardrail = self._bedrock_guardrail()
        guardrail.scan_only_tool_results = flag_value
        data = {
            "messages": [
                {"role": "user", "content": "USER-PROMPT"},
                {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
            ]
        }
        with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
            mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
            await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
            assert mock_api.call_count == 1
            scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
            assert scanned == ["USER-PROMPT", "TOOL-RESULT"], (
                "anything but an explicit True must leave the whole request in scope"
            )

    @pytest.mark.parametrize("scan_only_tool_results", [True, False])
    @pytest.mark.asyncio
    async def test_function_definitions_are_scoped_out_with_the_tool_results_flag(self, scan_only_tool_results):
        handler = OpenAIChatCompletionsHandler()
        guardrail = StructuredRedactionGuardrail()
        guardrail.scan_only_tool_results = scan_only_tool_results
        tools = [
            {
                "type": "function",
                "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}},
            }
        ]
        data = {
            "messages": [
                {"role": "user", "content": "read the report"},
                {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
            ],
            "tools": tools,
        }

        await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)

        assert guardrail.captured_inputs is not None
        expected_tools = None if scan_only_tool_results else tools
        assert guardrail.captured_inputs.get("tools") == expected_tools, (
            "function definitions must stay out of a tool-results-only scan"
        )

    @pytest.mark.parametrize("scan_only_tool_results", [True, False])
    @pytest.mark.asyncio
    async def test_guardrail_synthesized_tools_are_appended_without_replacing_request_tools(
        self, scan_only_tool_results
    ):
        handler = OpenAIChatCompletionsHandler()
        guardrail = ToolSynthesizingGuardrail()
        guardrail.scan_only_tool_results = scan_only_tool_results
        original_tools = [
            {
                "type": "function",
                "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}},
            }
        ]
        data = {
            "messages": [
                {"role": "user", "content": "read the report"},
                {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
            ],
            "tools": original_tools,
        }

        await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)

        assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"], (
            "a tool the guardrail synthesized (like a recovery/retrieve tool) must reach the model "
            "without the request's own tools being replaced or dropped"
        )
        assert data["tools"][0] == original_tools[0]

    @pytest.mark.asyncio
    async def test_returned_tool_name_collisions_keep_the_request_schema(self):
        handler = OpenAIChatCompletionsHandler()
        guardrail = ToolNameCollidingGuardrail()
        guardrail.scan_only_tool_results = True
        original_read_file = {
            "type": "function",
            "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}},
        }
        data = {
            "messages": [
                {"role": "user", "content": "read the report"},
                {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
            ],
            "tools": [original_read_file],
        }

        await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)

        assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"]
        assert data["tools"][0] == original_read_file, (
            "a returned tool reusing a request tool's name must not replace the request's schema"
        )

    @pytest.mark.asyncio
    async def test_duplicate_returned_tool_names_keep_only_the_first(self):
        handler = OpenAIChatCompletionsHandler()
        guardrail = DuplicateToolReturningGuardrail()
        guardrail.scan_only_tool_results = True
        original_read_file = {
            "type": "function",
            "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}},
        }
        data = {
            "messages": [
                {"role": "user", "content": "read the report"},
                {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"},
            ],
            "tools": [original_read_file],
        }

        await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)

        assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"], (
            "two returned tools sharing a name must not both be forwarded to the provider"
        )
        assert data["tools"][1]["function"]["parameters"]["properties"] == {"first": {"type": "string"}}

    @pytest.mark.asyncio
    async def test_structured_write_back_keeps_out_of_scope_messages(self):
        handler = OpenAIChatCompletionsHandler()
        guardrail = StructuredRedactionGuardrail()
        guardrail.scan_only_tool_results = True
        data = {
            "messages": [
                {"role": "system", "content": "SYSTEM-PROMPT"},
                {"role": "user", "content": "fetch the page"},
                {
                    "role": "assistant",
                    "content": "fetching",
                    "tool_calls": [
                        {"id": "call_1", "type": "function", "function": {"name": "fetch", "arguments": "{}"}}
                    ],
                },
                {"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"},
                {"role": "user", "content": "and then?"},
            ]
        }

        await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)

        assert [m["role"] for m in data["messages"]] == ["system", "user", "assistant", "tool", "user"], (
            "a redacting guardrail must not strip out-of-scope messages from the request"
        )
        assert data["messages"][0]["content"] == "SYSTEM-PROMPT"
        assert data["messages"][3]["content"] == "page says [BLOCKED] here"
        assert data["messages"][3]["tool_call_id"] == "call_1"
        assert data["messages"][4]["content"] == "and then?"


class TestBuildBlockSseChunks:
    """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks"""

    def _exc(self, original_response=None):
        from litellm.exceptions import ModifyResponseException

        return ModifyResponseException(
            message="Blocked by policy.",
            model="gpt-5.4-mini",
            request_data={},
            guardrail_name="test",
            original_response=original_response,
        )

    def _payloads(self, chunks):
        return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks]

    def test_standalone_block_uses_fresh_identity_and_zero_usage(self):
        handler = OpenAIChatCompletionsHandler()
        first, final = self._payloads(handler.build_block_sse_chunks(self._exc(), stream_started=False))
        assert first["id"].startswith("chatcmpl-")
        assert first["model"] == "gpt-5.4-mini"
        assert first["choices"][0]["delta"] == {"role": "assistant", "content": "Blocked by policy."}
        assert first["choices"][0]["finish_reason"] is None
        assert final["choices"][0]["delta"] == {}
        assert final["choices"][0]["finish_reason"] == "content_filter"
        assert final["usage"] == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

    def test_continuation_reuses_stream_identity_and_real_usage(self):
        handler = OpenAIChatCompletionsHandler()
        yielded = [
            {"id": "chatcmpl-live", "created": 1724900000, "model": "gpt-5.4-mini-2026-01-01"},
        ]
        original = yielded + [
            {"id": "chatcmpl-live", "usage": {"prompt_tokens": 11, "completion_tokens": 5}},
        ]
        first, final = self._payloads(
            handler.build_block_sse_chunks(
                self._exc(original_response=original), stream_started=True, responses_so_far=yielded
            )
        )
        assert (first["id"], first["created"], first["model"]) == (
            "chatcmpl-live",
            1724900000,
            "gpt-5.4-mini-2026-01-01",
        )
        assert first["choices"][0]["delta"] == {"content": "Blocked by policy."}
        assert final["id"] == "chatcmpl-live"
        assert final["usage"] == {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16}


class TestCheckStreamingHasEnded:
    """_check_streaming_has_ended lets end_of_stream_only withhold the finish chunk until moderation"""

    def test_empty_and_content_only_chunks_are_not_ended(self):
        handler = OpenAIChatCompletionsHandler()
        assert handler._check_streaming_has_ended([]) is False
        content_only = [
            {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]},
            {"id": "chatcmpl-live", "choices": []},
            {"id": "chatcmpl-live", "usage": {"prompt_tokens": 1, "completion_tokens": 1}},
        ]
        assert handler._check_streaming_has_ended(content_only) is False

    def test_dict_finish_chunk_marks_stream_ended(self):
        handler = OpenAIChatCompletionsHandler()
        chunks = [
            {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]},
            {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
        ]
        assert handler._check_streaming_has_ended(chunks) is True

    def test_object_finish_chunk_marks_stream_ended(self):
        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices

        handler = OpenAIChatCompletionsHandler()
        chunks = [
            ModelResponseStream(
                choices=[StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop")]
            )
        ]
        assert handler._check_streaming_has_ended(chunks) is True
