test_agui_protocol.py

"""AG-UI 协议处理器测试

测试 AGUIProtocolHandler 的各种功能。
"""

import json
from typing import cast

from fastapi.testclient import TestClient
import pytest

from agentrun.server import (
    AgentEvent,
    AgentRequest,
    AgentRunServer,
    AGUIProtocolHandler,
    EventType,
    ServerConfig,
)


class TestAGUIProtocolHandler:
    """测试 AGUIProtocolHandler"""

    def test_get_prefix_default(self):
        """测试默认前缀"""
        handler = AGUIProtocolHandler()
        assert handler.get_prefix() == "/ag-ui/agent"

    def test_get_prefix_custom(self):
        """测试自定义前缀"""
        from agentrun.server.model import AGUIProtocolConfig

        config = ServerConfig(agui=AGUIProtocolConfig(prefix="/custom/agui"))
        handler = AGUIProtocolHandler(config)
        assert handler.get_prefix() == "/custom/agui"


class TestAGUIProtocolEndpoints:
    """测试 AG-UI 协议端点"""

    def get_client(self, invoke_agent):
        server = AgentRunServer(invoke_agent=invoke_agent)
        return TestClient(server.as_fastapi_app())

    @pytest.mark.asyncio
    async def test_value_error_handling(self):
        """测试 ValueError 处理"""

        def invoke_agent(request: AgentRequest):
            raise ValueError("Test error")

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hello"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该包含 RUN_ERROR 事件
        types = []
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                types.append(data.get("type"))

        assert "RUN_ERROR" in types

    @pytest.mark.asyncio
    async def test_general_exception_handling(self):
        """测试一般异常处理（在 invoke_agent 中抛出异常）"""

        def invoke_agent(request: AgentRequest):
            raise RuntimeError("Internal error")

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hello"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该包含 RUN_ERROR 事件
        types = []
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                types.append(data.get("type"))

        assert "RUN_ERROR" in types

    @pytest.mark.asyncio
    async def test_exception_in_parse_request(self):
        """测试 parse_request 中的异常处理（覆盖 155-156 行）

        通过发送一个会导致 parse_request 抛出非 ValueError 异常的请求
        """
        from unittest.mock import AsyncMock, patch

        def invoke_agent(request: AgentRequest):
            return "Hello"

        client = self.get_client(invoke_agent)

        # 模拟 parse_request 抛出 RuntimeError
        with patch.object(
            AGUIProtocolHandler,
            "parse_request",
            new_callable=AsyncMock,
            side_effect=RuntimeError("Unexpected error"),
        ):
            response = client.post(
                "/ag-ui/agent",
                json={"messages": [{"role": "user", "content": "Hello"}]},
            )

            assert response.status_code == 200
            lines = [line async for line in response.aiter_lines() if line]

            # 应该包含 RUN_ERROR 事件
            types = []
            for line in lines:
                if line.startswith("data: "):
                    data = json.loads(line[6:])
                    types.append(data.get("type"))

            assert "RUN_ERROR" in types
            # 错误消息应该包含 "Internal error"
            for line in lines:
                if line.startswith("data: "):
                    data = json.loads(line[6:])
                    if data.get("type") == "RUN_ERROR":
                        assert "Internal error" in data.get("message", "")
                        break

    @pytest.mark.asyncio
    async def test_parse_messages_with_non_dict(self):
        """测试解析消息时跳过非字典项"""

        captured_request = {}

        def invoke_agent(request: AgentRequest):
            captured_request["messages"] = request.messages
            return "Done"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={
                "messages": [
                    "invalid_message",  # 非字典项，应该被跳过
                    {"role": "user", "content": "Hello"},
                ],
            },
        )

        assert response.status_code == 200
        assert len(captured_request["messages"]) == 1

    @pytest.mark.asyncio
    async def test_parse_messages_with_invalid_role(self):
        """测试解析消息时处理无效角色"""

        captured_request = {}

        def invoke_agent(request: AgentRequest):
            captured_request["messages"] = request.messages
            return "Done"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={
                "messages": [
                    {"role": "invalid_role", "content": "Hello"},  # 无效角色
                ],
            },
        )

        assert response.status_code == 200
        # 无效角色应该默认为 user
        from agentrun.server.model import MessageRole

        assert captured_request["messages"][0].role == MessageRole.USER

    @pytest.mark.asyncio
    async def test_parse_messages_with_tool_calls(self):
        """测试解析带有 toolCalls 的消息"""

        captured_request = {}

        def invoke_agent(request: AgentRequest):
            captured_request["messages"] = request.messages
            return "Done"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={
                "messages": [
                    {
                        "id": "msg-1",
                        "role": "assistant",
                        "content": None,
                        "toolCalls": [{
                            "id": "call_123",
                            "type": "function",
                            "function": {
                                "name": "get_weather",
                                "arguments": '{"city": "Beijing"}',
                            },
                        }],
                    },
                    {
                        "id": "msg-2",
                        "role": "tool",
                        "content": "Sunny",
                        "toolCallId": "call_123",
                    },
                ],
            },
        )

        assert response.status_code == 200
        assert len(captured_request["messages"]) == 2
        assert captured_request["messages"][0].tool_calls is not None
        assert captured_request["messages"][0].tool_calls[0].id == "call_123"
        assert captured_request["messages"][1].tool_call_id == "call_123"

    @pytest.mark.asyncio
    async def test_parse_tools(self):
        """测试解析工具列表"""

        captured_request = {}

        def invoke_agent(request: AgentRequest):
            captured_request["tools"] = request.tools
            return "Done"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={
                "messages": [{"role": "user", "content": "Hi"}],
                "tools": [{
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "description": "Get weather",
                    },
                }],
            },
        )

        assert response.status_code == 200
        assert captured_request["tools"] is not None
        assert len(captured_request["tools"]) == 1

    @pytest.mark.asyncio
    async def test_parse_tools_with_non_dict(self):
        """测试解析工具列表时跳过非字典项"""

        captured_request = {}

        def invoke_agent(request: AgentRequest):
            captured_request["tools"] = request.tools
            return "Done"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={
                "messages": [{"role": "user", "content": "Hi"}],
                "tools": [
                    "invalid_tool",
                    {"type": "function", "function": {"name": "valid"}},
                ],
            },
        )

        assert response.status_code == 200
        assert captured_request["tools"] is not None
        assert len(captured_request["tools"]) == 1

    @pytest.mark.asyncio
    async def test_parse_tools_empty_after_filter(self):
        """测试解析工具列表后为空时返回 None"""

        captured_request = {}

        def invoke_agent(request: AgentRequest):
            captured_request["tools"] = request.tools
            return "Done"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={
                "messages": [{"role": "user", "content": "Hi"}],
                "tools": ["invalid1", "invalid2"],
            },
        )

        assert response.status_code == 200
        assert captured_request["tools"] is None

    @pytest.mark.asyncio
    async def test_state_delta_event(self):
        """测试 STATE delta 事件"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.STATE,
                data={"delta": [{"op": "add", "path": "/count", "value": 1}]},
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 查找 STATE_DELTA 事件
        found_delta = False
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "STATE_DELTA":
                    found_delta = True
                    assert data["delta"] == [
                        {"op": "add", "path": "/count", "value": 1}
                    ]

        assert found_delta

    @pytest.mark.asyncio
    async def test_state_snapshot_fallback(self):
        """测试 STATE 事件没有 snapshot 或 delta 时的回退"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.STATE,
                data={
                    "custom_key": "custom_value"
                },  # 既不是 snapshot 也不是 delta
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该作为 STATE_SNAPSHOT 处理
        found_snapshot = False
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "STATE_SNAPSHOT":
                    found_snapshot = True
                    assert data["snapshot"]["custom_key"] == "custom_value"

        assert found_snapshot

    @pytest.mark.asyncio
    async def test_unknown_event_type(self):
        """测试未知事件类型转换为 CUSTOM"""

        # 直接测试 _process_event_with_boundaries 方法
        # 由于我们无法直接发送未知事件类型，我们测试 CUSTOM 事件的正常处理
        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.CUSTOM,
                data={"name": "unknown_event", "value": {"data": "test"}},
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 查找 CUSTOM 事件
        found_custom = False
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "CUSTOM":
                    found_custom = True
                    assert data["name"] == "unknown_event"

        assert found_custom

    @pytest.mark.asyncio
    async def test_addition_merge_overrides(self):
        """测试 addition 默认合并覆盖字段"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.TEXT,
                data={"delta": "Hello"},
                addition={"custom": "value", "delta": "overwritten"},
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 查找 TEXT_MESSAGE_CONTENT 事件
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "TEXT_MESSAGE_CONTENT":
                    assert "custom" in data
                    assert data["delta"] == "overwritten"
                    break

    @pytest.mark.asyncio
    async def test_addition_protocol_only_mode(self):
        """测试 addition PROTOCOL_ONLY 模式"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.TEXT,
                data={"delta": "Hello"},
                addition={
                    "delta": "overwritten",  # 已存在的字段会被覆盖
                    "new_field": "ignored",  # 新字段会被忽略
                },
                addition_merge_options={"no_new_field": True},
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 查找 TEXT_MESSAGE_CONTENT 事件
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "TEXT_MESSAGE_CONTENT":
                    assert data["delta"] == "overwritten"
                    assert "new_field" not in data
                    break

    @pytest.mark.asyncio
    async def test_raw_event_with_newline(self):
        """测试 RAW 事件自动添加换行符"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.RAW,
                data={"raw": '{"custom": "data"}'},  # 没有换行符
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        # RAW 事件应该被正确处理
        content = response.text
        assert '{"custom": "data"}' in content

    @pytest.mark.asyncio
    async def test_raw_event_with_trailing_newlines(self):
        """测试 RAW 事件带有尾随换行符时的处理"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.RAW,
                data={"raw": '{"custom": "data"}\n'},  # 只有一个换行符
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        content = response.text
        assert '{"custom": "data"}' in content

    @pytest.mark.asyncio
    async def test_raw_event_already_has_double_newline(self):
        """测试 RAW 事件已经有双换行符时不再添加"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.RAW,
                data={"raw": '{"custom": "data"}\n\n'},  # 已经有双换行符
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        content = response.text
        assert '{"custom": "data"}' in content

    @pytest.mark.asyncio
    async def test_raw_event_empty(self):
        """测试空 RAW 事件"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.RAW,
                data={"raw": ""},  # 空内容
            )
            yield "Hello"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该正常完成，包含 Hello 文本
        types = []
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                types.append(data.get("type"))

        assert "TEXT_MESSAGE_CONTENT" in types

    @pytest.mark.asyncio
    async def test_thread_id_and_run_id_from_request(self):
        """测试使用请求中的 threadId 和 runId"""

        async def invoke_agent(request: AgentRequest):
            yield "Hello"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={
                "messages": [{"role": "user", "content": "Hi"}],
                "threadId": "custom-thread-123",
                "runId": "custom-run-456",
            },
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 检查 RUN_STARTED 事件
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "RUN_STARTED":
                    assert data["threadId"] == "custom-thread-123"
                    assert data["runId"] == "custom-run-456"
                    break

    @pytest.mark.asyncio
    async def test_new_text_message_after_tool_result(self):
        """测试工具结果后的新文本消息有新的 messageId"""

        async def invoke_agent(request: AgentRequest):
            yield "First message"
            yield AgentEvent(
                event=EventType.TOOL_CALL_CHUNK,
                data={"id": "tc-1", "name": "tool", "args_delta": "{}"},
            )
            yield AgentEvent(
                event=EventType.TOOL_RESULT,
                data={"id": "tc-1", "result": "done"},
            )
            # 这应该触发 TEXT_MESSAGE_END 然后 TEXT_MESSAGE_START（新消息）
            yield "Second message"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 收集所有 messageId
        message_ids = []
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") in [
                    "TEXT_MESSAGE_START",
                    "TEXT_MESSAGE_CONTENT",
                    "TEXT_MESSAGE_END",
                ]:
                    message_ids.append(data.get("messageId"))

        # 第一段和第二段文本应该有相同的 messageId（AG-UI 允许并行）
        # 实际上在当前实现中，工具调用后的文本是同一个消息的延续
        assert len(message_ids) >= 2


class TestAGUIProtocolErrorStream:
    """测试 AG-UI 协议错误流"""

    def get_client(self, invoke_agent):
        server = AgentRunServer(invoke_agent=invoke_agent)
        return TestClient(server.as_fastapi_app())

    @pytest.mark.asyncio
    async def test_error_stream_on_json_parse_error(self):
        """测试 JSON 解析错误时的错误流"""

        def invoke_agent(request: AgentRequest):
            return "Hello"

        client = self.get_client(invoke_agent)
        # 发送无效的 JSON
        response = client.post(
            "/ag-ui/agent",
            content="invalid json",
            headers={"Content-Type": "application/json"},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该包含 RUN_STARTED 和 RUN_ERROR
        types = []
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                types.append(data.get("type"))

        assert "RUN_STARTED" in types
        assert "RUN_ERROR" in types

    @pytest.mark.asyncio
    async def test_error_stream_format(self):
        """测试错误流的格式"""

        def invoke_agent(request: AgentRequest):
            raise ValueError("Test validation error")

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 检查错误事件的格式
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "RUN_ERROR":
                    # 错误事件应该包含 message
                    assert "message" in data
                    break


class TestAGUIProtocolApplyAddition:
    """测试 _apply_addition 方法"""

    def test_apply_addition_default_merge(self):
        """默认合并应覆盖已有字段"""
        handler = AGUIProtocolHandler()

        event_data = {"delta": "Hello", "type": "TEXT_MESSAGE_CONTENT"}
        addition = {"delta": "overwritten", "new_field": "added"}

        result = handler._apply_addition(
            event_data.copy(),
            addition.copy(),
        )

        assert result["delta"] == "overwritten"
        assert result["new_field"] == "added"
        assert result["type"] == "TEXT_MESSAGE_CONTENT"

    def test_apply_addition_merge_options_none(self):
        """显式传入 merge_options=None 仍按默认合并"""
        handler = AGUIProtocolHandler()

        event_data = {"delta": "Hello", "type": "TEXT_MESSAGE_CONTENT"}
        addition = {"delta": "overwritten", "new_field": "added"}

        result = handler._apply_addition(
            event_data.copy(),
            addition.copy(),
        )

        assert result["delta"] == "overwritten"
        assert result["new_field"] == "added"

    def test_apply_addition_protocol_only_mode(self):
        """测试 PROTOCOL_ONLY 模式（覆盖 616->620 分支）"""
        handler = AGUIProtocolHandler()

        event_data = {"delta": "Hello", "type": "TEXT_MESSAGE_CONTENT"}
        addition = {"delta": "overwritten", "new_field": "ignored"}

        result = handler._apply_addition(
            event_data.copy(),
            addition.copy(),
            {"no_new_field": True},
        )

        # delta 被覆盖
        assert result["delta"] == "overwritten"
        # new_field 不存在（被忽略）
        assert "new_field" not in result
        # type 保持不变
        assert result["type"] == "TEXT_MESSAGE_CONTENT"


class TestAGUIProtocolConvertMessages:
    """测试消息转换功能"""

    def test_convert_messages_for_snapshot(self):
        """测试 _convert_messages_for_snapshot 方法"""
        handler = AGUIProtocolHandler()

        messages = [
            {"role": "user", "content": "Hello", "id": "msg-1"},
            {"role": "assistant", "content": "Hi there", "id": "msg-2"},
            {"role": "system", "content": "You are helpful", "id": "msg-3"},
            {
                "role": "tool",
                "content": "Result",
                "id": "msg-4",
                "tool_call_id": "tc-1",
            },
        ]

        result = handler._convert_messages_for_snapshot(messages)

        assert len(result) == 4
        assert result[0].role == "user"
        assert result[1].role == "assistant"
        assert result[2].role == "system"
        assert result[3].role == "tool"

    def test_convert_messages_with_non_dict(self):
        """测试 _convert_messages_for_snapshot 跳过非字典项"""
        handler = AGUIProtocolHandler()

        messages = [
            "invalid",
            {"role": "user", "content": "Hello", "id": "msg-1"},
            123,
            {"role": "assistant", "content": "Hi", "id": "msg-2"},
        ]

        result = handler._convert_messages_for_snapshot(messages)

        assert len(result) == 2

    def test_convert_messages_with_missing_id(self):
        """测试 _convert_messages_for_snapshot 自动生成 id"""
        handler = AGUIProtocolHandler()

        messages = [
            {"role": "user", "content": "Hello"},  # 没有 id
        ]

        result = handler._convert_messages_for_snapshot(messages)

        assert len(result) == 1
        assert result[0].id is not None
        assert len(result[0].id) > 0

    def test_convert_messages_with_unknown_role(self):
        """测试 _convert_messages_for_snapshot 处理未知角色"""
        handler = AGUIProtocolHandler()

        messages = [
            {"role": "unknown", "content": "Hello", "id": "msg-1"},
            {"role": "user", "content": "World", "id": "msg-2"},
        ]

        result = handler._convert_messages_for_snapshot(messages)

        # 未知角色的消息会被跳过
        assert len(result) == 1
        assert result[0].role == "user"


class TestAGUIProtocolUnknownEvent:
    """测试未知事件类型处理"""

    def get_client(self, invoke_agent):
        server = AgentRunServer(invoke_agent=invoke_agent)
        return TestClient(server.as_fastapi_app())

    @pytest.mark.asyncio
    async def test_tool_result_event(self):
        """测试 TOOL_RESULT 事件（使用 content 字段）"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.TOOL_CALL_CHUNK,
                data={"id": "tc-1", "name": "tool", "args_delta": "{}"},
            )
            yield AgentEvent(
                event=EventType.TOOL_RESULT,
                data={
                    "id": "tc-1",
                    "content": "Tool content result",
                },  # 使用 content
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 查找 TOOL_CALL_RESULT 事件
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "TOOL_CALL_RESULT":
                    assert data["content"] == "Tool content result"
                    break


class TestAGUIProtocolTextMessageRestart:
    """测试文本消息重新开始"""

    def get_client(self, invoke_agent):
        server = AgentRunServer(invoke_agent=invoke_agent)
        return TestClient(server.as_fastapi_app())

    @pytest.mark.asyncio
    async def test_text_message_restart_after_end(self):
        """测试文本消息结束后重新开始新消息

        当 text_state["ended"] 为 True 时，新的 TEXT 事件应该开始新消息
        """

        async def invoke_agent(request: AgentRequest):
            # 第一段文本
            yield "First"
            # 工具调用会导致文本消息结束
            yield AgentEvent(
                event=EventType.TOOL_CALL_CHUNK,
                data={"id": "tc-1", "name": "tool", "args_delta": "{}"},
            )
            yield AgentEvent(
                event=EventType.TOOL_RESULT,
                data={"id": "tc-1", "result": "done"},
            )
            # 第二段文本应该开始新消息
            yield "Second"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 统计 TEXT_MESSAGE_START 事件数量
        start_count = 0
        message_ids = set()
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "TEXT_MESSAGE_START":
                    start_count += 1
                    message_ids.add(data.get("messageId"))

        # 应该只有一个 TEXT_MESSAGE_START（AG-UI 允许并行）
        # 但如果实现了重新开始逻辑，可能有两个
        assert start_count >= 1

    @pytest.mark.asyncio
    async def test_state_delta_event(self):
        """测试 STATE delta 事件"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.STATE,
                data={
                    "delta": [{"op": "replace", "path": "/count", "value": 10}]
                },
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 查找 STATE_DELTA 事件
        found = False
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                if data.get("type") == "STATE_DELTA":
                    found = True
                    assert data["delta"][0]["op"] == "replace"
                    break

        assert found


class TestAGUIProtocolExceptionHandling:
    """测试异常处理"""

    def get_client(self, invoke_agent):
        server = AgentRunServer(invoke_agent=invoke_agent)
        return TestClient(server.as_fastapi_app())

    @pytest.mark.asyncio
    async def test_general_exception_returns_error_stream(self):
        """测试一般异常返回错误流"""

        def invoke_agent(request: AgentRequest):
            raise RuntimeError("Unexpected error")

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该包含 RUN_ERROR
        types = [
            json.loads(line[6:]).get("type")
            for line in lines
            if line.startswith("data: ")
        ]
        assert "RUN_ERROR" in types


class TestAGUIProtocolUnknownEventType:
    """测试未知事件类型处理（覆盖第 531 行）

    TOOL_CALL 事件没有在 _process_event_with_boundaries 中被显式处理，
    所以它会走到第 531 行的 else 分支，被转换为 CUSTOM 事件。
    """

    def get_client(self, invoke_agent):
        server = AgentRunServer(invoke_agent=invoke_agent)
        return TestClient(server.as_fastapi_app())

    def test_process_event_with_boundaries_unknown_event(self):
        """直接测试 _process_event_with_boundaries 处理未知事件类型

        TOOL_CALL 事件没有在 _process_event_with_boundaries 中被显式处理，
        所以它会走到 else 分支，被转换为 CUSTOM 事件。
        """
        handler = AGUIProtocolHandler()

        # 创建 TOOL_CALL 事件（在协议层没有被显式处理）
        event = AgentEvent(
            event=EventType.TOOL_CALL,
            data={"id": "tc-1", "name": "test_tool", "args": '{"x": 1}'},
        )

        context = {"thread_id": "test-thread", "run_id": "test-run"}

        # 创建 StreamStateMachine 对象
        from agentrun.server.agui_protocol import StreamStateMachine

        state = StreamStateMachine()

        # 调用方法
        results = list(
            handler._process_event_with_boundaries(event, context, state)
        )

        # TOOL_CALL 现在被实际处理，生成 TOOL_CALL_START 和 TOOL_CALL_ARGS 事件
        assert len(results) == 2
        # 解析第一个 SSE 数据 (TOOL_CALL_START)
        sse_data_1 = results[0]
        assert sse_data_1.startswith("data: ")
        data_1 = json.loads(sse_data_1[6:].strip())
        assert data_1["type"] == "TOOL_CALL_START"
        assert data_1["toolCallId"] == "tc-1"
        assert data_1["toolCallName"] == "test_tool"

        # 解析第二个 SSE 数据 (TOOL_CALL_ARGS)
        sse_data_2 = results[1]
        assert sse_data_2.startswith("data: ")
        data_2 = json.loads(sse_data_2[6:].strip())
        assert data_2["type"] == "TOOL_CALL_ARGS"
        assert data_2["toolCallId"] == "tc-1"
        assert data_2["delta"] == '{"x": 1}'

    @pytest.mark.asyncio
    async def test_tool_call_event_expanded_by_invoker(self):
        """测试 TOOL_CALL 事件被 invoker 展开为 TOOL_CALL_CHUNK

        通过端到端测试验证 TOOL_CALL 被正确处理。
        """

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.TOOL_CALL,
                data={"id": "tc-1", "name": "test_tool", "args": '{"x": 1}'},
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # TOOL_CALL 会被 invoker 展开为 TOOL_CALL_CHUNK
        # 所以应该有 TOOL_CALL_START 和 TOOL_CALL_ARGS 事件
        types = []
        for line in lines:
            if line.startswith("data: "):
                data = json.loads(line[6:])
                types.append(data.get("type"))

        assert "TOOL_CALL_START" in types
        assert "TOOL_CALL_ARGS" in types


class TestAGUIProtocolToolCallBranches:
    """测试工具调用的各种分支"""

    def get_client(self, invoke_agent):
        server = AgentRunServer(invoke_agent=invoke_agent)
        return TestClient(server.as_fastapi_app())

    @pytest.mark.asyncio
    async def test_tool_result_for_already_ended_tool(self):
        """测试对已结束的工具调用发送结果"""

        async def invoke_agent(request: AgentRequest):
            # 工具调用
            yield AgentEvent(
                event=EventType.TOOL_CALL_CHUNK,
                data={"id": "tc-1", "name": "tool", "args_delta": "{}"},
            )
            # 第一个结果（会结束工具调用）
            yield AgentEvent(
                event=EventType.TOOL_RESULT,
                data={"id": "tc-1", "result": "result1"},
            )
            # 第二个结果（工具调用已结束）
            yield AgentEvent(
                event=EventType.TOOL_RESULT,
                data={"id": "tc-1", "result": "result2"},
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该有两个 TOOL_CALL_RESULT
        result_count = sum(
            1
            for line in lines
            if line.startswith("data: ") and "TOOL_CALL_RESULT" in line
        )
        assert result_count == 2

    @pytest.mark.asyncio
    async def test_empty_sse_data_filtered(self):
        """测试空 SSE 数据被过滤"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.RAW,
                data={"raw": ""},  # 空内容
            )
            yield "Hello"

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该正常完成
        types = [
            json.loads(line[6:]).get("type")
            for line in lines
            if line.startswith("data: ")
        ]
        assert "RUN_FINISHED" in types

    @pytest.mark.asyncio
    async def test_tool_call_with_empty_id(self):
        """测试空 tool_id 的工具调用"""

        async def invoke_agent(request: AgentRequest):
            yield AgentEvent(
                event=EventType.TOOL_CALL_CHUNK,
                data={"id": "", "name": "tool", "args_delta": "{}"},  # 空 id
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200

    @pytest.mark.asyncio
    async def test_tool_result_without_prior_start(self):
        """测试没有先发送 TOOL_CALL_CHUNK 的 TOOL_RESULT"""

        async def invoke_agent(request: AgentRequest):
            # 直接发送 TOOL_RESULT，没有 TOOL_CALL_CHUNK
            yield AgentEvent(
                event=EventType.TOOL_RESULT,
                data={"id": "tc-orphan", "result": "orphan result"},
            )

        client = self.get_client(invoke_agent)
        response = client.post(
            "/ag-ui/agent",
            json={"messages": [{"role": "user", "content": "Hi"}]},
        )

        assert response.status_code == 200
        lines = [line async for line in response.aiter_lines() if line]

        # 应该自动补充 TOOL_CALL_START 和 TOOL_CALL_END
        types = [
            json.loads(line[6:]).get("type")
            for line in lines
            if line.startswith("data: ")
        ]
        assert "TOOL_CALL_START" in types
        assert "TOOL_CALL_END" in types
        assert "TOOL_CALL_RESULT" in types