import pytest

import litellm
from litellm.types.utils import ModelResponse


@pytest.mark.asyncio
async def test_acompletion_mcp_auto_exec(monkeypatch):
    from types import SimpleNamespace

    from litellm.responses.mcp.litellm_proxy_mcp_handler import (
        LiteLLM_Proxy_MCP_Handler,
    )
    from litellm.responses.utils import ResponsesAPIRequestUtils

    dummy_tool = SimpleNamespace(
        name="local_search",
        description="search",
        inputSchema={"type": "object", "properties": {}},
    )

    async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs):
        return [dummy_tool], {"local_search": "local"}

    async def fake_execute(**kwargs):
        fake_execute.called = True  # type: ignore[attr-defined]
        tool_calls = kwargs.get("tool_calls") or []
        assert tool_calls, "tool calls should be present during auto execution"
        call_entry = tool_calls[0]
        call_id = call_entry.get("id") or call_entry.get("call_id") or "call"
        return [
            {
                "tool_call_id": call_id,
                "result": "executed",
                "name": call_entry.get("name", "local_search"),
            }
        ]

    fake_execute.called = False  # type: ignore[attr-defined]

    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_process_mcp_tools_without_openai_transform",
        fake_process,
    )
    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_execute_tool_calls",
        fake_execute,
    )
    monkeypatch.setattr(
        ResponsesAPIRequestUtils,
        "extract_mcp_headers_from_request",
        staticmethod(lambda secret_fields, tools: (None, None, None, None)),
    )

    response = await litellm.acompletion(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": "hello"}],
        tools=[
            {
                "type": "mcp",
                "server_url": "litellm_proxy/mcp/local",
                "server_label": "local",
                "require_approval": "never",
            }
        ],
        mock_response="Final answer",
        mock_tool_calls=[
            {
                "id": "call-1",
                "type": "function",
                "function": {"name": "local_search", "arguments": "{}"},
            }
        ],
    )

    assert isinstance(response, ModelResponse)
    assert response.choices[0].message.content == "Final answer"
    assert fake_execute.called is True  # type: ignore[attr-defined]


@pytest.mark.asyncio
async def test_acompletion_mcp_respects_manual_approval(monkeypatch):
    from types import SimpleNamespace

    from litellm.responses.mcp.litellm_proxy_mcp_handler import (
        LiteLLM_Proxy_MCP_Handler,
    )
    from litellm.responses.utils import ResponsesAPIRequestUtils

    dummy_tool = SimpleNamespace(
        name="local_search",
        description="search",
        inputSchema={"type": "object", "properties": {}},
    )

    async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs):
        return [dummy_tool], {"local_search": "local"}

    async def fake_execute(**kwargs):
        pytest.fail("auto execution should not run when approval is required")

    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_process_mcp_tools_without_openai_transform",
        fake_process,
    )
    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_execute_tool_calls",
        fake_execute,
    )
    monkeypatch.setattr(
        ResponsesAPIRequestUtils,
        "extract_mcp_headers_from_request",
        staticmethod(lambda secret_fields, tools: (None, None, None, None)),
    )

    response = await litellm.acompletion(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": "hello"}],
        tools=[
            {
                "type": "mcp",
                "server_url": "litellm_proxy/mcp/local",
                "server_label": "local",
                "require_approval": "manual",
            }
        ],
        mock_response="Pending tool",
        mock_tool_calls=[
            {
                "id": "call-2",
                "type": "function",
                "function": {"name": "local_search", "arguments": "{}"},
            }
        ],
    )

    assert isinstance(response, ModelResponse)
    tool_calls = response.choices[0].message.tool_calls
    assert tool_calls is not None and len(tool_calls) == 1


@pytest.mark.asyncio
async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch):
    """
    Test that litellm.completion with stream=True and MCP tools does not raise
    RuntimeError: Timeout context manager should be used inside a task.

    This test ensures that the fix in ba43f742ab86d51b7da63077b85b39d0ac808d30
    prevents event loop nesting issues when using MCP tools with streaming.

    The fix changes completion() to return a coroutine from acompletion_with_mcp,
    which acompletion() then awaits, avoiding event loop nesting.
    """
    from types import SimpleNamespace
    from unittest.mock import patch

    from litellm.responses.mcp.litellm_proxy_mcp_handler import (
        LiteLLM_Proxy_MCP_Handler,
    )
    from litellm.responses.utils import ResponsesAPIRequestUtils
    from litellm.utils import CustomStreamWrapper

    dummy_tool = SimpleNamespace(
        name="local_search",
        description="search",
        inputSchema={"type": "object", "properties": {}},
    )

    async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs):
        return [dummy_tool], {"local_search": "local"}

    async def fake_execute(**kwargs):
        fake_execute.called = True  # type: ignore[attr-defined]
        tool_calls = kwargs.get("tool_calls") or []
        assert tool_calls, "tool calls should be present during auto execution"
        call_entry = tool_calls[0]
        call_id = call_entry.get("id") or call_entry.get("call_id") or "call"
        return [
            {
                "tool_call_id": call_id,
                "result": "executed",
                "name": call_entry.get("name", "local_search"),
            }
        ]

    fake_execute.called = False  # type: ignore[attr-defined]

    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_process_mcp_tools_without_openai_transform",
        fake_process,
    )
    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_execute_tool_calls",
        fake_execute,
    )
    monkeypatch.setattr(
        ResponsesAPIRequestUtils,
        "extract_mcp_headers_from_request",
        staticmethod(lambda secret_fields, tools: (None, None, None, None)),
    )

    # Create a mock streaming response
    from unittest.mock import MagicMock, AsyncMock

    logging_obj = MagicMock()
    logging_obj.model_call_details = {}
    logging_obj.async_failure_handler = AsyncMock()

    class MockStreamingResponse(CustomStreamWrapper):
        def __init__(self):
            super().__init__(
                completion_stream=None,
                model="gpt-4o-mini",
                logging_obj=logging_obj,
            )
            self.chunks = [
                type(
                    "Chunk",
                    (),
                    {
                        "choices": [
                            type(
                                "Choice",
                                (),
                                {"delta": type("Delta", (), {"content": "Final"})()},
                            )()
                        ]
                    },
                )(),
                type(
                    "Chunk",
                    (),
                    {
                        "choices": [
                            type(
                                "Choice",
                                (),
                                {"delta": type("Delta", (), {"content": " answer"})()},
                            )()
                        ]
                    },
                )(),
            ]
            self._index = 0

        def __iter__(self):
            return self

        def __next__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopIteration

        def __aiter__(self):
            return self

        async def __anext__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopAsyncIteration

    # Track calls to acompletion
    acompletion_calls = []

    # Create mock streaming response for initial call
    from unittest.mock import MagicMock, AsyncMock

    logging_obj = MagicMock()
    logging_obj.model_call_details = {}
    logging_obj.async_failure_handler = AsyncMock()

    from litellm.types.utils import (
        ModelResponseStream,
        StreamingChoices,
        Delta,
        ChatCompletionDeltaToolCall,
        Function,
    )

    # Create initial streaming chunks with tool_calls
    # Add tool_calls to the final chunk so stream_chunk_builder can extract them
    tool_calls = [
        ChatCompletionDeltaToolCall(
            id="call-1",
            type="function",
            function=Function(name="local_search", arguments="{}"),
            index=0,
        )
    ]

    initial_chunks = [
        ModelResponseStream(
            id="test-1",
            model="gpt-4o-mini",
            created=1234567890,
            object="chat.completion.chunk",
            choices=[
                StreamingChoices(
                    index=0,
                    delta=Delta(
                        content="",
                        role="assistant",
                        tool_calls=tool_calls,
                    ),
                    finish_reason="tool_calls",
                )
            ],
        )
    ]

    class InitialStreamingResponse(CustomStreamWrapper):
        def __init__(self):
            super().__init__(
                completion_stream=None,
                model="gpt-4o-mini",
                logging_obj=logging_obj,
            )
            self.chunks = initial_chunks
            self._index = 0

        def __iter__(self):
            return self

        def __next__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopIteration

        def __aiter__(self):
            return self

        async def __anext__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopAsyncIteration

    async def mock_acompletion(**kwargs):
        acompletion_calls.append(kwargs)
        # With new implementation, first call should be streaming
        if kwargs.get("stream", False):
            # Check if this is the follow-up call
            messages = kwargs.get("messages", [])
            is_follow_up = any(
                msg.get("role") == "tool"
                or (isinstance(msg, dict) and "tool_call_id" in str(msg))
                for msg in messages
            )
            if is_follow_up:
                # Follow-up call (streaming)
                return MockStreamingResponse()
            else:
                # Initial call (streaming)
                return InitialStreamingResponse()
        # Non-streaming call should not happen with new implementation, but handle it
        return ModelResponse(
            id="test-1",
            model="gpt-4o-mini",
            choices=[
                {
                    "message": {
                        "role": "assistant",
                        "tool_calls": [
                            {
                                "id": "call-1",
                                "type": "function",
                                "function": {"name": "local_search", "arguments": "{}"},
                            }
                        ],
                    },
                    "finish_reason": "tool_calls",
                }
            ],
            created=0,
            object="chat.completion",
        )

    with patch("litellm.acompletion", side_effect=mock_acompletion):
        # This should not raise RuntimeError: Timeout context manager should be used inside a task
        # completion() returns a coroutine when MCP tools are present, which acompletion() awaits
        response = litellm.completion(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "hello"}],
            tools=[
                {
                    "type": "mcp",
                    "server_url": "litellm_proxy/mcp/local",
                    "server_label": "local",
                    "require_approval": "never",
                }
            ],
            stream=True,
            mock_response="Final answer",
            mock_tool_calls=[
                {
                    "id": "call-1",
                    "type": "function",
                    "function": {"name": "local_search", "arguments": "{}"},
                }
            ],
        )

        # completion() returns a coroutine when MCP tools are present
        import asyncio

        assert asyncio.iscoroutine(
            response
        ), "completion() should return a coroutine when MCP tools are present"

        # Await the coroutine (this is what acompletion() does internally)
        # This should not raise RuntimeError: Timeout context manager should be used inside a task
        result = await response

        # Verify response is a streaming response
        assert isinstance(result, CustomStreamWrapper) or hasattr(result, "__iter__")

        # Consume the stream to ensure it works (run in separate thread to avoid event loop conflict)
        from concurrent.futures import ThreadPoolExecutor

        def consume_stream():
            return list(result)

        with ThreadPoolExecutor(max_workers=1) as executor:
            chunks = executor.submit(consume_stream).result()
        assert len(chunks) > 0, "Should have received streaming chunks"

        # Verify tool execution was called
        assert fake_execute.called is True  # type: ignore[attr-defined]

        # Verify acompletion was called (should be called by acompletion_with_mcp)
        assert len(acompletion_calls) >= 1, "acompletion should be called"


@pytest.mark.asyncio
async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
    """
    Test that MCP metadata is added correctly to streaming chunks:
    - mcp_list_tools should be in the first chunk
    - mcp_tool_calls and mcp_call_results should be in the final chunk of initial response
    - Follow-up response should be streamed after initial response
    """
    from types import SimpleNamespace
    from unittest.mock import patch

    from litellm.responses.mcp.litellm_proxy_mcp_handler import (
        LiteLLM_Proxy_MCP_Handler,
    )
    from litellm.responses.utils import ResponsesAPIRequestUtils
    from litellm.utils import CustomStreamWrapper
    from litellm.types.utils import (
        ModelResponseStream,
        StreamingChoices,
        Delta,
        ChatCompletionDeltaToolCall,
        Function,
    )
    from litellm.litellm_core_utils.litellm_logging import Logging

    dummy_tool = SimpleNamespace(
        name="local_search",
        description="search",
        inputSchema={"type": "object", "properties": {}},
    )

    async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs):
        return [dummy_tool], {"local_search": "local"}

    async def fake_execute(**kwargs):
        tool_calls = kwargs.get("tool_calls") or []
        call_entry = tool_calls[0]
        call_id = call_entry.get("id") or call_entry.get("call_id") or "call"
        return [
            {
                "tool_call_id": call_id,
                "result": "executed",
                "name": call_entry.get("name", "local_search"),
            }
        ]

    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_process_mcp_tools_without_openai_transform",
        fake_process,
    )
    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_execute_tool_calls",
        fake_execute,
    )
    monkeypatch.setattr(
        ResponsesAPIRequestUtils,
        "extract_mcp_headers_from_request",
        staticmethod(lambda secret_fields, tools: (None, None, None, None)),
    )

    # Create mock streaming chunks
    def create_chunk(content, finish_reason=None, tool_calls=None):
        delta = Delta(
            content=content,
            role="assistant",
        )
        if tool_calls:
            delta.tool_calls = tool_calls
        return ModelResponseStream(
            id="test-stream",
            model="gpt-4o-mini",
            created=1234567890,
            object="chat.completion.chunk",
            choices=[
                StreamingChoices(
                    index=0,
                    delta=delta,
                    finish_reason=finish_reason,
                )
            ],
        )

    # Create initial streaming chunks with tool_calls
    # Add tool_calls to the final chunk so stream_chunk_builder can extract them
    tool_calls = [
        ChatCompletionDeltaToolCall(
            id="call-1",
            type="function",
            function=Function(name="local_search", arguments="{}"),
            index=0,
        )
    ]
    initial_chunks = [
        create_chunk(
            "", finish_reason="tool_calls", tool_calls=tool_calls
        ),  # Final chunk with tool_calls
    ]

    # Create follow-up streaming chunks
    follow_up_chunks = [
        create_chunk("Hello"),
        create_chunk(" world"),
        create_chunk("!", finish_reason="stop"),  # Final chunk
    ]

    # Create a proper CustomStreamWrapper with logging_obj
    from unittest.mock import MagicMock, AsyncMock

    logging_obj = MagicMock()
    logging_obj.model_call_details = {}
    logging_obj.async_failure_handler = AsyncMock()

    class InitialStreamingResponse(CustomStreamWrapper):
        def __init__(self):
            super().__init__(
                completion_stream=None,
                model="gpt-4o-mini",
                logging_obj=logging_obj,
            )
            self.chunks = initial_chunks
            self._index = 0

        def __iter__(self):
            return self

        def __next__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopIteration

        def __aiter__(self):
            return self

        async def __anext__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopAsyncIteration

    class FollowUpStreamingResponse(CustomStreamWrapper):
        def __init__(self):
            super().__init__(
                completion_stream=None,
                model="gpt-4o-mini",
                logging_obj=logging_obj,
            )
            self.chunks = follow_up_chunks
            self._index = 0

        def __iter__(self):
            return self

        def __next__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopIteration

        def __aiter__(self):
            return self

        async def __anext__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopAsyncIteration

    # Track calls to acompletion
    acompletion_calls = []

    async def mock_acompletion(**kwargs):
        acompletion_calls.append(kwargs)
        # With new implementation, first call should be streaming
        if kwargs.get("stream", False):
            # Check if this is the follow-up call (has tool results in messages)
            messages = kwargs.get("messages", [])
            is_follow_up = any(
                msg.get("role") == "tool"
                or (isinstance(msg, dict) and "tool_call_id" in str(msg))
                for msg in messages
            )

            if is_follow_up:
                # Follow-up call - return follow-up chunks
                return FollowUpStreamingResponse()
            else:
                # Initial streaming call - return chunks with tool_calls
                return InitialStreamingResponse()
        # Non-streaming call should not happen with new implementation, but handle it
        return ModelResponse(
            id="test-1",
            model="gpt-4o-mini",
            choices=[
                {
                    "message": {
                        "role": "assistant",
                        "tool_calls": [
                            {
                                "id": "call-1",
                                "type": "function",
                                "function": {"name": "local_search", "arguments": "{}"},
                            }
                        ],
                    },
                    "finish_reason": "tool_calls",
                }
            ],
            created=0,
            object="chat.completion",
        )

    with patch("litellm.acompletion", side_effect=mock_acompletion):
        response = litellm.completion(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "hello"}],
            tools=[
                {
                    "type": "mcp",
                    "server_url": "litellm_proxy/mcp/local",
                    "server_label": "local",
                    "require_approval": "never",
                }
            ],
            stream=True,
            mock_response="Final answer",
            mock_tool_calls=[
                {
                    "id": "call-1",
                    "type": "function",
                    "function": {"name": "local_search", "arguments": "{}"},
                }
            ],
        )

        import asyncio

        assert asyncio.iscoroutine(response)
        result = await response

        assert isinstance(result, CustomStreamWrapper)

        # Consume the stream and check chunks (run in separate thread to avoid event loop conflict)
        from concurrent.futures import ThreadPoolExecutor

        def consume_stream():
            return list(result)

        with ThreadPoolExecutor(max_workers=1) as executor:
            all_chunks = executor.submit(consume_stream).result()
        assert len(all_chunks) > 0, "Should have received streaming chunks"

        # Find chunks from initial response (with tool_calls finish_reason)
        initial_chunks_list = []
        follow_up_chunks_list = []
        for chunk in all_chunks:
            if hasattr(chunk, "choices") and chunk.choices:
                choice = chunk.choices[0]
                if (
                    hasattr(choice, "finish_reason")
                    and choice.finish_reason == "tool_calls"
                ):
                    initial_chunks_list.append(chunk)
                elif (
                    hasattr(choice, "finish_reason") and choice.finish_reason == "stop"
                ):
                    follow_up_chunks_list.append(chunk)
                elif (
                    not hasattr(choice, "finish_reason") or choice.finish_reason is None
                ):
                    # Chunks without finish_reason could be from either stream
                    # Check if we've seen tool_calls yet
                    if initial_chunks_list:
                        follow_up_chunks_list.append(chunk)
                    else:
                        initial_chunks_list.append(chunk)

        # Verify initial response chunks
        assert len(initial_chunks_list) > 0, "Should have initial response chunks"

        # Find the final chunk from initial response (with tool_calls finish_reason)
        initial_final_chunk = None
        for chunk in initial_chunks_list:
            if hasattr(chunk, "choices") and chunk.choices:
                choice = chunk.choices[0]
                if (
                    hasattr(choice, "finish_reason")
                    and choice.finish_reason == "tool_calls"
                ):
                    initial_final_chunk = chunk
                    break

        if initial_final_chunk is None and initial_chunks_list:
            initial_final_chunk = initial_chunks_list[-1]

        assert (
            initial_final_chunk is not None
        ), "Should have a final chunk from initial response"

        # Verify mcp_list_tools is in the first chunk of initial response
        first_chunk = initial_chunks_list[0] if initial_chunks_list else None
        assert first_chunk is not None, "Should have a first chunk"
        if hasattr(first_chunk, "choices") and first_chunk.choices:
            choice = first_chunk.choices[0]
            if hasattr(choice, "delta") and choice.delta:
                provider_fields = getattr(
                    choice.delta, "provider_specific_fields", None
                )
                assert (
                    provider_fields is not None
                ), "First chunk should have provider_specific_fields"
                assert (
                    "mcp_list_tools" in provider_fields
                ), "First chunk should have mcp_list_tools"

        # Verify mcp_tool_calls and mcp_call_results are in the final chunk of initial response
        if hasattr(initial_final_chunk, "choices") and initial_final_chunk.choices:
            choice = initial_final_chunk.choices[0]
            if hasattr(choice, "delta") and choice.delta:
                provider_fields = getattr(
                    choice.delta, "provider_specific_fields", None
                )
                assert (
                    provider_fields is not None
                ), "Final chunk should have provider_specific_fields"
                assert "mcp_tool_calls" in provider_fields, "Should have mcp_tool_calls"
                assert (
                    "mcp_call_results" in provider_fields
                ), "Should have mcp_call_results"

        # Verify follow-up response chunks are present
        assert len(follow_up_chunks_list) > 0, "Should have follow-up response chunks"


@pytest.mark.asyncio
async def test_mcp_streaming_metadata_ordering(monkeypatch):
    """
    Test that MCP metadata appears in the correct order:
    - mcp_list_tools should appear in the first chunk (before tool_calls)
    - mcp_tool_calls and mcp_call_results should appear in the final chunk of initial response
    - Follow-up response should be streamed after initial response completes
    """
    from types import SimpleNamespace
    from unittest.mock import patch

    from litellm.responses.mcp.litellm_proxy_mcp_handler import (
        LiteLLM_Proxy_MCP_Handler,
    )
    from litellm.responses.utils import ResponsesAPIRequestUtils
    from litellm.utils import CustomStreamWrapper
    from litellm.types.utils import (
        ModelResponseStream,
        StreamingChoices,
        Delta,
        ChatCompletionDeltaToolCall,
        Function,
    )

    dummy_tool = SimpleNamespace(
        name="local_search",
        description="search",
        inputSchema={"type": "object", "properties": {}},
    )

    async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy, **kwargs):
        return [dummy_tool], {"local_search": "local"}

    async def fake_execute(**kwargs):
        tool_calls = kwargs.get("tool_calls") or []
        call_entry = tool_calls[0]
        call_id = call_entry.get("id") or call_entry.get("call_id") or "call"
        return [
            {
                "tool_call_id": call_id,
                "result": "executed",
                "name": call_entry.get("name", "local_search"),
            }
        ]

    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_process_mcp_tools_without_openai_transform",
        fake_process,
    )
    monkeypatch.setattr(
        LiteLLM_Proxy_MCP_Handler,
        "_execute_tool_calls",
        fake_execute,
    )
    monkeypatch.setattr(
        ResponsesAPIRequestUtils,
        "extract_mcp_headers_from_request",
        staticmethod(lambda secret_fields, tools: (None, None, None, None)),
    )

    # Create mock streaming chunks
    def create_chunk(content, finish_reason=None, tool_calls=None):
        delta = Delta(
            content=content,
            role="assistant",
        )
        if tool_calls:
            delta.tool_calls = tool_calls
        return ModelResponseStream(
            id="test-stream",
            model="gpt-4o-mini",
            created=1234567890,
            object="chat.completion.chunk",
            choices=[
                StreamingChoices(
                    index=0,
                    delta=delta,
                    finish_reason=finish_reason,
                )
            ],
        )

    # Create initial streaming chunks with tool_calls
    # Add tool_calls to the final chunk so stream_chunk_builder can extract them
    tool_calls = [
        ChatCompletionDeltaToolCall(
            id="call-1",
            type="function",
            function=Function(name="local_search", arguments="{}"),
            index=0,
        )
    ]
    initial_chunks = [
        create_chunk(
            "", finish_reason="tool_calls", tool_calls=tool_calls
        ),  # Final chunk with tool_calls
    ]

    # Create follow-up streaming chunks
    follow_up_chunks = [
        create_chunk("Hello"),
        create_chunk(" world"),
        create_chunk("!", finish_reason="stop"),  # Final chunk
    ]

    # Create a proper CustomStreamWrapper with logging_obj
    from unittest.mock import MagicMock, AsyncMock

    logging_obj = MagicMock()
    logging_obj.model_call_details = {}
    logging_obj.async_failure_handler = AsyncMock()

    class InitialStreamingResponse(CustomStreamWrapper):
        def __init__(self):
            super().__init__(
                completion_stream=None,
                model="gpt-4o-mini",
                logging_obj=logging_obj,
            )
            self.chunks = initial_chunks
            self._index = 0

        def __iter__(self):
            return self

        def __next__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopIteration

        def __aiter__(self):
            return self

        async def __anext__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopAsyncIteration

    class FollowUpStreamingResponse(CustomStreamWrapper):
        def __init__(self):
            super().__init__(
                completion_stream=None,
                model="gpt-4o-mini",
                logging_obj=logging_obj,
            )
            self.chunks = follow_up_chunks
            self._index = 0

        def __iter__(self):
            return self

        def __next__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopIteration

        def __aiter__(self):
            return self

        async def __anext__(self):
            if self._index < len(self.chunks):
                chunk = self.chunks[self._index]
                self._index += 1
                # Add mcp_list_tools to first chunk if present
                if not self.sent_first_chunk:
                    chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
                    self.sent_first_chunk = True
                return chunk
            raise StopAsyncIteration

    # Track calls to acompletion
    acompletion_calls = []

    async def mock_acompletion(**kwargs):
        acompletion_calls.append(kwargs)
        # With new implementation, first call should be streaming
        if kwargs.get("stream", False):
            # Check if this is the follow-up call (has tool results in messages)
            messages = kwargs.get("messages", [])
            is_follow_up = any(
                msg.get("role") == "tool"
                or (isinstance(msg, dict) and "tool_call_id" in str(msg))
                for msg in messages
            )

            if is_follow_up:
                # Follow-up call - return follow-up chunks
                return FollowUpStreamingResponse()
            else:
                # Initial streaming call - return chunks with tool_calls
                return InitialStreamingResponse()
        # Non-streaming call should not happen with new implementation
        pytest.fail("Non-streaming call should not happen with new implementation")

    with patch("litellm.acompletion", side_effect=mock_acompletion):
        response = litellm.completion(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "hello"}],
            tools=[
                {
                    "type": "mcp",
                    "server_url": "litellm_proxy/mcp/local",
                    "server_label": "local",
                    "require_approval": "never",
                }
            ],
            stream=True,
            mock_response="Final answer",
            mock_tool_calls=[
                {
                    "id": "call-1",
                    "type": "function",
                    "function": {"name": "local_search", "arguments": "{}"},
                }
            ],
        )

        import asyncio

        assert asyncio.iscoroutine(response)
        result = await response

        assert isinstance(result, CustomStreamWrapper)

        # Consume the stream and verify order (run in separate thread to avoid event loop conflict)
        from concurrent.futures import ThreadPoolExecutor

        def consume_stream():
            return list(result)

        with ThreadPoolExecutor(max_workers=1) as executor:
            all_chunks = executor.submit(consume_stream).result()
        assert len(all_chunks) > 0, "Should have received streaming chunks"

        # Track when we see each type of metadata
        mcp_list_tools_seen = False
        mcp_tool_calls_seen = False
        mcp_call_results_seen = False
        tool_calls_finish_reason_seen = False
        follow_up_content_seen = False

        for i, chunk in enumerate(all_chunks):
            if hasattr(chunk, "choices") and chunk.choices:
                choice = chunk.choices[0]
                if hasattr(choice, "delta") and choice.delta:
                    provider_fields = getattr(
                        choice.delta, "provider_specific_fields", None
                    )
                    if provider_fields:
                        if "mcp_list_tools" in provider_fields:
                            mcp_list_tools_seen = True
                            # mcp_list_tools should appear before tool_calls finish_reason
                            assert (
                                not tool_calls_finish_reason_seen
                            ), "mcp_list_tools should appear before tool_calls finish_reason"
                        if "mcp_tool_calls" in provider_fields:
                            mcp_tool_calls_seen = True
                        if "mcp_call_results" in provider_fields:
                            mcp_call_results_seen = True

                if (
                    hasattr(choice, "finish_reason")
                    and choice.finish_reason == "tool_calls"
                ):
                    tool_calls_finish_reason_seen = True
                    # mcp_tool_calls and mcp_call_results should be in the same chunk as tool_calls finish_reason
                    if hasattr(choice, "delta") and choice.delta:
                        provider_fields = getattr(
                            choice.delta, "provider_specific_fields", None
                        )
                        assert provider_fields is not None
                        assert (
                            "mcp_tool_calls" in provider_fields
                        ), "mcp_tool_calls should be in the chunk with tool_calls finish_reason"
                        assert (
                            "mcp_call_results" in provider_fields
                        ), "mcp_call_results should be in the chunk with tool_calls finish_reason"

                if hasattr(choice, "delta") and choice.delta and choice.delta.content:
                    content = choice.delta.content
                    if content and (
                        "Hello" in content or "world" in content or "!" in content
                    ):
                        follow_up_content_seen = True
                        # Follow-up content should appear after tool_calls finish_reason
                        assert (
                            tool_calls_finish_reason_seen
                        ), "Follow-up content should appear after tool_calls finish_reason"

        # Verify all metadata was seen
        assert mcp_list_tools_seen, "Should have seen mcp_list_tools"
        assert mcp_tool_calls_seen, "Should have seen mcp_tool_calls"
        assert mcp_call_results_seen, "Should have seen mcp_call_results"
        assert (
            tool_calls_finish_reason_seen
        ), "Should have seen tool_calls finish_reason"
        assert follow_up_content_seen, "Should have seen follow-up content"
