ai-video-fullstack/backend/tests/test_conversation_history.py

from __future__ import annotations

import asyncio
import unittest
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

from services.conversation_history import ConversationRecorder


class ConversationRecorderTest(unittest.IsolatedAsyncioTestCase):
    async def test_background_messages_keep_order_and_ignore_streaming_deltas(self):
        recorder = ConversationRecorder("conv_test")
        recorder.record_transport_message = AsyncMock()
        events = [
            {"type": "transcript", "content": "user"},
            {"type": "assistant-text-end", "content": "assistant"},
        ]
        for event in events:
            recorder.record_transport_message_later(event)
        recorder.record_transport_message_later({"type": "assistant-text-delta"})
        await recorder.flush_pending_messages()
        self.assertEqual(
            [call.args[0] for call in recorder.record_transport_message.await_args_list],
            events,
        )
        self.assertIsNone(recorder._message_writer)

    async def test_slow_database_does_not_block_audio(self):
        from pipecat.frames.frames import EndFrame, OutputAudioRawFrame, OutputTransportMessageUrgentFrame
        from pipecat.pipeline.pipeline import Pipeline
        from pipecat.pipeline.worker import PipelineWorker
        from pipecat.processors.frame_processor import FrameProcessor
        from pipecat.workers.runner import WorkerRunner
        from services.pipecat.processors import ConversationHistoryProcessor

        entered, release, audio_received = asyncio.Event(), asyncio.Event(), asyncio.Event()
        recorder = ConversationRecorder("conv_test")

        async def slow_write(message):
            entered.set()
            await release.wait()

        recorder.record_transport_message = slow_write

        class AudioSink(FrameProcessor):
            async def process_frame(self, frame, direction):
                await super().process_frame(frame, direction)
                if isinstance(frame, OutputAudioRawFrame):
                    audio_received.set()
                await self.push_frame(frame, direction)

        worker = PipelineWorker(
            Pipeline([ConversationHistoryProcessor(recorder), AudioSink()]), enable_rtvi=False,
        )
        runner = WorkerRunner(handle_sigint=False)
        await runner.add_workers(worker)
        task = asyncio.create_task(runner.run())
        try:
            await worker.queue_frame(OutputTransportMessageUrgentFrame(message={"type": "transcript"}))
            await asyncio.wait_for(entered.wait(), 3)
            await worker.queue_frame(OutputAudioRawFrame(audio=b"\0" * 320, sample_rate=16000, num_channels=1))
            await asyncio.wait_for(audio_received.wait(), 1)
            self.assertFalse(release.is_set())
        finally:
            release.set()
            await recorder.flush_pending_messages()
            await worker.queue_frame(EndFrame())
            await asyncio.wait_for(task, 3)

    async def test_history_buffer_is_bounded_and_slow_writer_is_cancelled(self):
        started = asyncio.Event()
        cancelled = asyncio.Event()
        with patch("services.conversation_history.MAX_PENDING_HISTORY_MESSAGES", 2):
            recorder = ConversationRecorder("conv_test")

        async def blocked_write(message):
            started.set()
            try:
                await asyncio.Event().wait()
            finally:
                cancelled.set()

        recorder.record_transport_message = blocked_write
        recorder.record_transport_message_later({"type": "transcript"})
        await started.wait()
        for _ in range(4):
            recorder.record_transport_message_later({"type": "transcript"})
        self.assertEqual(recorder._pending_messages.qsize(), 2)
        with patch("services.conversation_history.HISTORY_FLUSH_TIMEOUT_SECONDS", 0.01):
            await recorder.flush_pending_messages()
        self.assertTrue(cancelled.is_set())
        self.assertTrue(recorder._pending_messages.empty())
        recorder.record_transport_message_later({"type": "transcript"})
        self.assertIsNone(recorder._message_writer)

    async def test_completed_conversation_queues_enabled_analysis(self):
        conversation = SimpleNamespace(
            status="active",
            ended_at=None,
            message_count=0,
            analysis_status="none",
            analysis_data={},
            analysis_error="",
        )

        class FakeSession:
            async def __aenter__(self):
                return self

            async def __aexit__(self, *_args):
                return None

            async def get(self, _model, _session_id):
                return conversation

            async def commit(self):
                return None

        plan = {
            "enabled": True,
            "model_resource_id": "model_001",
            "fields": [{"name": "customer_name", "type": "string"}],
        }
        recorder = ConversationRecorder("conv_test", plan)
        recorder._sequence = 1
        with patch(
            "services.conversation_history.SessionLocal",
            return_value=FakeSession(),
        ):
            await recorder._finish(status="completed")

        self.assertEqual(conversation.analysis_status, "pending")
        self.assertEqual(conversation.analysis_data, {"plan": plan})

    async def test_finish_waits_for_database_cleanup_when_cancelled(self):
        recorder = ConversationRecorder("conv_test")
        started = asyncio.Event()
        release = asyncio.Event()
        completed = asyncio.Event()

        async def finish_database_write(*, status: str):
            self.assertEqual(status, "failed")
            started.set()
            await release.wait()
            completed.set()

        recorder._finish = finish_database_write
        task = asyncio.create_task(recorder.finish(status="failed"))
        await started.wait()
        task.cancel()
        await asyncio.sleep(0)

        self.assertFalse(task.done())
        release.set()
        with self.assertRaises(asyncio.CancelledError):
            await task
        self.assertTrue(completed.is_set())

    async def test_fixed_reply_transcript_keeps_workflow_metadata(self):
        recorder = ConversationRecorder("conv_test")
        recorder._append = AsyncMock()

        await recorder.record_transport_message(
            {
                "type": "transcript",
                "role": "assistant",
                "content": "请稍等，我正在处理。",
                "timestamp": "2026-07-14T10:00:00+08:00",
                "source": "workflow-fixed-reply",
                "nodeId": "agent_service",
            }
        )

        recorder._append.assert_awaited_once_with(
            "assistant",
            "请稍等，我正在处理。",
            "2026-07-14T10:00:00+08:00",
            {
                "source": "workflow-fixed-reply",
                "node_id": "agent_service",
            },
        )

    async def test_workflow_trace_is_persisted_without_counting_as_message(self):
        conversation = SimpleNamespace(
            extra={"workflow": {"revision": "sha256:test"}},
            message_count=3,
        )

        class FakeSession:
            committed = False

            async def __aenter__(self):
                return self

            async def __aexit__(self, *_args):
                return None

            async def get(self, _model, _session_id):
                return conversation

            async def commit(self):
                self.committed = True

        session = FakeSession()
        recorder = ConversationRecorder("conv_test")
        event = {
            "type": "workflow-event",
            "eventId": "wfe_test",
            "event": "node_entered",
            "timestamp": "2026-08-01T10:00:00+08:00",
            "workflowRevision": "sha256:test",
            "transitionId": 0,
            "nodeId": "start",
        }

        with patch(
            "services.conversation_history.SessionLocal",
            return_value=session,
        ):
            await recorder.record_transport_message(event)
            await recorder.record_transport_message(event)

        self.assertTrue(session.committed)
        self.assertEqual(conversation.message_count, 3)
        self.assertEqual(len(conversation.extra["workflowTrace"]), 1)
        saved = conversation.extra["workflowTrace"][0]
        self.assertEqual(saved["sessionId"], "conv_test")
        self.assertEqual(saved["sequence"], 1)

    async def test_internal_message_stage_transport_trace_is_not_persisted_twice(self):
        recorder = ConversationRecorder("conv_test")
        event = {
            "type": "workflow-event",
            "eventId": "wfe_internal",
            "event": "client_tool_completed",
            "functionName": "show_message",
            "internalStage": True,
        }

        with patch("services.conversation_history.SessionLocal") as session_local:
            await recorder.record_transport_message(event)

        session_local.assert_not_called()
        self.assertEqual(recorder._trace_sequence, 0)

    async def test_image_is_persisted_as_message_with_artifact(self):
        conversation = SimpleNamespace(message_count=0)

        class FakeSession:
            def __init__(self):
                self.added = []
                self.committed = False

            async def __aenter__(self):
                return self

            async def __aexit__(self, *_args):
                return None

            def add(self, value):
                self.added.append(value)

            async def flush(self):
                return None

            async def get(self, _model, _session_id):
                return conversation

            async def commit(self):
                self.committed = True

        session = FakeSession()
        recorder = ConversationRecorder("conv_test")
        with (
            patch("services.conversation_history.SessionLocal", return_value=session),
            patch("services.conversation_history.put_object") as put_object,
        ):
            await recorder._record_image(
                b"jpeg-data",
                input_id="input_photo",
                timestamp="2026-08-05T10:00:00+08:00",
                mime_type="image/jpeg",
                content="帮我看看",
                source="uploaded_asset",
            )

        self.assertTrue(session.committed)
        self.assertEqual(len(session.added), 2)
        message, artifact = session.added
        self.assertEqual(message.content_type, "image")
        self.assertEqual(message.role, "user")
        self.assertEqual(message.content, "帮我看看")
        self.assertEqual(message.extra["input_id"], "input_photo")
        self.assertEqual(message.extra["source"], "uploaded_asset")
        self.assertEqual(artifact.message_id, message.id)
        self.assertEqual(artifact.kind, "image")
        self.assertEqual(artifact.size_bytes, len(b"jpeg-data"))
        self.assertEqual(conversation.message_count, 1)
        put_object.assert_called_once()


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