wyndham-ARR/tests/test_agent_events.py

from __future__ import annotations

import unittest

from agent_integration.client import OpenAgentEvent
from agent_integration.events import AgentStreamNormalizer


class AgentStreamNormalizerTests(unittest.TestCase):
    def test_live_langgraph_shape_exposes_only_visible_answer(self):
        normalizer = AgentStreamNormalizer()
        raw_events = [
            OpenAgentEvent(
                event="metadata",
                data={"run_id": "run-001", "thread_id": "internal-thread"},
                event_id="1",
            ),
            OpenAgentEvent(
                event="values",
                data={
                    "messages": [{"type": "human", "content": "hello"}],
                    "thread_data": {"workspace_path": "/internal/secret/path"},
                },
            ),
            OpenAgentEvent(
                event="messages",
                data=[
                    {
                        "type": "AIMessageChunk",
                        "content": "",
                        "additional_kwargs": {"reasoning_content": "INTERNAL_REASONING"},
                    },
                    {"langgraph_node": "model"},
                ],
            ),
            OpenAgentEvent(
                event="messages",
                data=[
                    {"type": "AIMessageChunk", "content": "INTERNAL_TITLE"},
                    {"langgraph_node": "TitleMiddleware.after_model", "tags": ["middleware:title"]},
                ],
            ),
            OpenAgentEvent(
                event="messages",
                data=[
                    {"type": "AIMessageChunk", "content": "连接"},
                    {"langgraph_node": "model"},
                ],
                event_id="2",
            ),
            OpenAgentEvent(
                event="messages",
                data=[
                    {"type": "AIMessageChunk", "content": "成功"},
                    {"langgraph_node": "model"},
                ],
                event_id="3",
            ),
            OpenAgentEvent(
                event="messages",
                data=[
                    {
                        "type": "ai",
                        "content": "连接成功",
                        "additional_kwargs": {"reasoning_content": "INTERNAL_REASONING"},
                    },
                    {"langgraph_node": "model"},
                ],
            ),
            OpenAgentEvent(
                event="values",
                data={
                    "messages": [
                        {"type": "human", "content": "hello"},
                        {
                            "type": "ai",
                            "content": "连接成功",
                            "additional_kwargs": {"reasoning_content": "INTERNAL_REASONING"},
                        },
                    ],
                    "thread_data": {"workspace_path": "/internal/secret/path"},
                },
            ),
            OpenAgentEvent(
                event="error",
                data={
                    "code": "open_agent_final_content_missing",
                    "message": "Open agent run completed without final content",
                    "retryable": True,
                    "run_id": "run-001",
                    "request_id": "request-001",
                    "internal_trace": "DO_NOT_EXPOSE",
                },
            ),
            OpenAgentEvent(event="end", data=None, event_id="4"),
        ]

        public_events = []
        for event in raw_events:
            public_events.extend(normalizer.feed(event))
        public_events.extend(normalizer.finish())

        self.assertEqual(
            [event.event for event in public_events],
            [
                "run.started",
                "message.delta",
                "message.delta",
                "message.completed",
                "run.warning",
                "run.end",
            ],
        )
        self.assertEqual(public_events[0].data, {"run_id": "run-001"})
        self.assertEqual(public_events[1].data, {"content": "连接"})
        self.assertEqual(public_events[2].data, {"content": "成功"})
        self.assertEqual(
            public_events[3].data,
            {"content": "连接成功", "streamed": True},
        )
        self.assertEqual(public_events[4].event, "run.warning")
        self.assertEqual(public_events[5].data["status"], "completed_with_warning")
        self.assertFalse(normalizer.fatal_error)
        self.assertEqual(normalizer.final_content, "连接成功")
        rendered = repr([event.to_dict() for event in public_events])
        self.assertNotIn("INTERNAL_REASONING", rendered)
        self.assertNotIn("INTERNAL_TITLE", rendered)
        self.assertNotIn("secret/path", rendered)
        self.assertNotIn("DO_NOT_EXPOSE", rendered)

    def test_final_content_missing_without_visible_content_is_fatal(self):
        normalizer = AgentStreamNormalizer()
        normalizer.feed(
            OpenAgentEvent(
                event="error",
                data={
                    "code": "open_agent_final_content_missing",
                    "message": "missing",
                    "retryable": True,
                },
            )
        )
        public_events = normalizer.feed(OpenAgentEvent(event="end", data=None))

        self.assertTrue(normalizer.fatal_error)
        self.assertEqual([event.event for event in public_events], ["run.error", "run.end"])
        self.assertEqual(public_events[-1].data["status"], "failed")

    def test_values_final_message_is_fallback_when_no_deltas_arrive(self):
        normalizer = AgentStreamNormalizer()
        normalizer.feed(
            OpenAgentEvent(
                event="values",
                data={"messages": [{"type": "ai", "content": "完整答案"}]},
            )
        )
        public_events = normalizer.finish()

        self.assertEqual(public_events[0].event, "message.completed")
        self.assertEqual(
            public_events[0].data,
            {"content": "完整答案", "streamed": False},
        )
        self.assertFalse(normalizer.fatal_error)

    def test_documented_message_delta_shape_remains_supported(self):
        normalizer = AgentStreamNormalizer()
        public_events = normalizer.feed(
            OpenAgentEvent(event="message.delta", data={"content": "你"})
        )
        normalizer.feed(
            OpenAgentEvent(event="run.completed", data={"status": "completed"})
        )
        public_events.extend(normalizer.finish())

        self.assertEqual(public_events[0].data, {"content": "你"})
        self.assertEqual(public_events[1].data, {"content": "你", "streamed": True})
        self.assertEqual(public_events[2].data["status"], "completed")


if __name__ == "__main__":
    unittest.main(verbosity=2)