learningcircuit--local-deep-research/tests/chat/test_chat_termination.py

"""Tests for chat-mode research termination — partial-content persistence,
final socket emit, and idempotency under mid-stream interrupts.

These tests target the helper directly with mocked ChatService and
SocketIOService. End-to-end termination through the worker thread is
covered by the chat e2e suite.
"""

from unittest.mock import patch, MagicMock


SESSION_ID = "00000000-0000-0000-0000-000000000001"
RESEARCH_ID = "00000000-0000-0000-0000-000000000002"
USERNAME = "alice"


def _import_helper():
    from local_deep_research.web.services.research_service import (
        _save_partial_chat_message_on_terminate,
        _STOPPED_BEFORE_OUTPUT_MARKER,
        _STOPPED_FOOTER,
    )

    return (
        _save_partial_chat_message_on_terminate,
        _STOPPED_BEFORE_OUTPUT_MARKER,
        _STOPPED_FOOTER,
    )


class TestPartialPersistOnTerminate:
    """`_save_partial_chat_message_on_terminate` behaviour."""

    def test_persists_partial_with_stopped_footer(self):
        helper, _marker, footer = _import_helper()

        with (
            patch("local_deep_research.chat.service.ChatService") as mock_svc,
            patch(
                "local_deep_research.web.services.research_service.SocketIOService"
            ) as mock_sock,
        ):
            instance = MagicMock()
            mock_svc.return_value = instance

            helper(
                SESSION_ID,
                RESEARCH_ID,
                USERNAME,
                "partial answer text",
            )

            mock_svc.assert_called_once_with(USERNAME)
            instance.add_message.assert_called_once()
            kwargs = instance.add_message.call_args.kwargs
            assert kwargs["session_id"] == SESSION_ID
            assert kwargs["role"] == "assistant"
            assert kwargs["message_type"] == "response"
            assert kwargs["research_id"] == RESEARCH_ID
            assert kwargs["content"].startswith("partial answer text")
            assert footer.strip() in kwargs["content"]

            # Socket.IO final-chunk emit must fire so the bubble drops the
            # streaming class on the client.
            mock_sock.return_value.emit_to_subscribers.assert_called_once()
            event_args = (
                mock_sock.return_value.emit_to_subscribers.call_args.args
            )
            assert event_args[0] == "response_chunk"
            assert event_args[1] == RESEARCH_ID
            payload = event_args[2]
            assert payload["is_final"] is True

    def test_persists_marker_when_no_chunks_streamed(self):
        helper, marker, _footer = _import_helper()

        with (
            patch("local_deep_research.chat.service.ChatService") as mock_svc,
            patch(
                "local_deep_research.web.services.research_service.SocketIOService"
            ),
        ):
            instance = MagicMock()
            mock_svc.return_value = instance

            helper(SESSION_ID, RESEARCH_ID, USERNAME, "")

            assert instance.add_message.call_args.kwargs["content"] == marker

    def test_truncated_marker_appended_when_buffer_capped(self):
        helper, _marker, _footer = _import_helper()

        with (
            patch("local_deep_research.chat.service.ChatService") as mock_svc,
            patch(
                "local_deep_research.web.services.research_service.SocketIOService"
            ),
        ):
            instance = MagicMock()
            mock_svc.return_value = instance

            helper(
                SESSION_ID,
                RESEARCH_ID,
                USERNAME,
                "a long but truncated stream",
                truncated=True,
            )

            content = instance.add_message.call_args.kwargs["content"]
            assert "truncated" in content.lower()

    def test_skips_silently_when_no_chat_session_id(self):
        helper, _marker, _footer = _import_helper()

        with (
            patch("local_deep_research.chat.service.ChatService") as mock_svc,
            patch(
                "local_deep_research.web.services.research_service.SocketIOService"
            ) as mock_sock,
        ):
            helper(None, RESEARCH_ID, USERNAME, "anything")

            mock_svc.assert_not_called()
            mock_sock.return_value.emit_to_subscribers.assert_not_called()

    def test_idempotent_via_streaming_state_flag(self):
        """Two calls with the same streaming_state must persist exactly one
        chat row — guards against the in-callback path and the outer
        except handler both firing on a mid-stream interrupt."""
        helper, _marker, _footer = _import_helper()

        state = {"chunks": [], "_persisted": False}

        with (
            patch("local_deep_research.chat.service.ChatService") as mock_svc,
            patch(
                "local_deep_research.web.services.research_service.SocketIOService"
            ) as mock_sock,
        ):
            instance = MagicMock()
            mock_svc.return_value = instance

            helper(
                SESSION_ID,
                RESEARCH_ID,
                USERNAME,
                "first call",
                streaming_state=state,
            )
            helper(
                SESSION_ID,
                RESEARCH_ID,
                USERNAME,
                "second call",
                streaming_state=state,
            )

            assert state["_persisted"] is True
            assert instance.add_message.call_count == 1
            assert mock_sock.return_value.emit_to_subscribers.call_count == 1

    def test_swallows_chat_service_failures(self):
        """Persistence failures must NOT crash — termination cleanup is
        best-effort."""
        helper, _marker, _footer = _import_helper()

        with (
            patch("local_deep_research.chat.service.ChatService") as mock_svc,
            patch(
                "local_deep_research.web.services.research_service.SocketIOService"
            ),
        ):
            mock_svc.side_effect = RuntimeError("DB exploded")

            # Should not raise.
            helper(SESSION_ID, RESEARCH_ID, USERNAME, "anything")


class TestNoDuplicateRowOnLateTermination:
    """Corner case: termination flag flips between the success-path response
    write (``_save_chat_message_and_context``) and the trailing
    ``progress_callback("Research completed successfully", 100)``.

    The trailing callback's termination check would otherwise call
    ``_save_partial_chat_message_on_terminate`` and write a SECOND
    response row. The fix: ``_save_chat_message_and_context`` sets
    ``streaming_state["_persisted"] = True`` after writing the success
    row, and the helper short-circuits on that flag.
    """

    def test_late_terminate_does_not_write_duplicate_row(self):
        """Simulate the production sequence: success-path writes the row,
        then a late terminate fires. Only one row total."""
        from local_deep_research.web.services.research_service import (
            _save_chat_message_and_context,
            _save_partial_chat_message_on_terminate,
        )

        streaming_state = {
            "chunks_sent": 3,
            "chunks": ["abc", "def", "ghi"],
            "_bytes": 9,
            "_truncated": False,
        }

        with (
            patch("local_deep_research.chat.service.ChatService") as mock_svc,
            patch(
                "local_deep_research.chat.context.ChatContextManager"
            ) as mock_ctx,
            patch(
                "local_deep_research.web.services.research_service.SocketIOService"
            ),
        ):
            success_chat = MagicMock()
            mock_svc.return_value = success_chat
            mock_ctx.return_value.extract_context_updates.return_value = {}

            # 1) Success path runs first — writes the response row.
            sock = MagicMock()
            _save_chat_message_and_context(
                SESSION_ID,
                RESEARCH_ID,
                USERNAME,
                "the full answer",
                streaming_enabled=True,
                streaming_state=streaming_state,
                socket_service=sock,
            )

            # Sanity: success-path wrote exactly one row + flipped the flag.
            assert success_chat.add_message.call_count == 1
            assert streaming_state["_persisted"] is True

            # 2) Trailing progress_callback detects termination and calls
            #    the helper. The flag must short-circuit it.
            _save_partial_chat_message_on_terminate(
                SESSION_ID,
                RESEARCH_ID,
                USERNAME,
                "".join(streaming_state["chunks"]),
                streaming_state=streaming_state,
            )

            # Still exactly one add_message — no duplicate row.
            assert success_chat.add_message.call_count == 1

    def test_late_terminate_emits_no_duplicate_final_chunk(self):
        """The success path emits its own ``is_final`` chunk; the helper
        must not emit a second one when short-circuited."""
        from local_deep_research.web.services.research_service import (
            _save_chat_message_and_context,
            _save_partial_chat_message_on_terminate,
        )

        streaming_state = {
            "chunks_sent": 1,
            "chunks": ["x"],
            "_bytes": 1,
            "_truncated": False,
        }

        success_sock = MagicMock()

        with (
            patch("local_deep_research.chat.service.ChatService"),
            patch("local_deep_research.chat.context.ChatContextManager"),
            patch(
                "local_deep_research.web.services.research_service.SocketIOService"
            ) as mock_terminate_sock,
        ):
            _save_chat_message_and_context(
                SESSION_ID,
                RESEARCH_ID,
                USERNAME,
                "answer",
                streaming_enabled=True,
                streaming_state=streaming_state,
                socket_service=success_sock,
            )

            # Success path emitted its own final chunk.
            assert success_sock.emit_to_subscribers.call_count == 1

            _save_partial_chat_message_on_terminate(
                SESSION_ID,
                RESEARCH_ID,
                USERNAME,
                "x",
                streaming_state=streaming_state,
            )

            # The helper's SocketIOService factory must NOT have been used —
            # the short-circuit fired before any emit.
            mock_terminate_sock.return_value.emit_to_subscribers.assert_not_called()