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

"""Mid-stream termination tests — exercise ``_make_chat_stream_callback``
end-to-end through ``BaseCitationHandler._invoke_with_streaming`` to verify
that clicking Stop while the LLM is streaming actually interrupts the
loop, persists the partial chunks, and propagates ``ResearchTerminatedException``.

Uses a fake LLM that yields chunks one at a time; the test flips the
termination flag after the Nth chunk and asserts the rest never run.
"""

from unittest.mock import MagicMock, patch

import pytest

from local_deep_research.exceptions import ResearchTerminatedException


class _FakeStreamingLLM:
    """Minimal LLM with ``.stream()`` yielding one chunk at a time."""

    def __init__(self, chunks):
        self._chunks = chunks
        self.invoke_called = False

    def stream(self, _prompt):
        for c in self._chunks:
            yield c

    def invoke(self, _prompt):
        # If termination falls through to ``invoke()`` we want to know —
        # see the ``test_terminated_does_not_fall_back_to_invoke`` case.
        self.invoke_called = True
        return "FALLBACK"


def _build_callback_and_state(research_id):
    """Build the production stream_callback wired to a fake socket."""
    from local_deep_research.web.services.research_service import (
        _make_chat_stream_callback,
    )

    streaming_state = {
        "chunks_sent": 0,
        "chunks": [],
        "_bytes": 0,
        "_truncated": False,
    }
    sock = MagicMock()
    cb = _make_chat_stream_callback(research_id, streaming_state, sock)
    return cb, streaming_state, sock


class TestStreamCallbackBuffering:
    def test_buffers_chunks_for_partial_persistence(self):
        cb, state, _sock = _build_callback_and_state("rid-1")

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            return_value=False,
        ):
            for chunk in ["Hel", "lo, ", "wor", "ld"]:
                cb(chunk)

        assert state["chunks_sent"] == 4
        assert "".join(state["chunks"]) == "Hello, world"
        assert state["_truncated"] is False

    def test_buffer_caps_at_max_bytes(self):
        from local_deep_research.web.services import research_service

        cb, state, _sock = _build_callback_and_state("rid-cap")
        big_chunk = "x" * (research_service._MAX_PARTIAL_BUFFER_BYTES // 2 + 1)

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            return_value=False,
        ):
            cb(big_chunk)
            cb(big_chunk)  # second push should overflow the cap

        assert state["_truncated"] is True
        # The first chunk was buffered in full; the second was dropped from
        # the buffer but ``chunks_sent`` keeps counting.
        assert "".join(state["chunks"]) == big_chunk
        assert state["chunks_sent"] == 2

    def test_empty_chunks_ignored(self):
        cb, state, _sock = _build_callback_and_state("rid-empty")

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            return_value=False,
        ):
            cb("")
            cb(None)

        assert state["chunks_sent"] == 0
        assert state["chunks"] == []


class TestStreamCallbackTermination:
    def test_raises_when_termination_flag_set(self):
        cb, state, _sock = _build_callback_and_state("rid-x")

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            return_value=True,
        ):
            with pytest.raises(ResearchTerminatedException):
                cb("any-chunk")

        # The chunk was buffered before the raise — partial persistence
        # depends on this.
        assert state["chunks"] == ["any-chunk"]
        assert state["chunks_sent"] == 1

    def test_propagates_through_invoke_with_streaming(self):
        """The full path: fake LLM yields chunks, our callback raises mid-stream,
        ``BaseCitationHandler._invoke_with_streaming`` lets it propagate
        rather than falling back to ``llm.invoke()``."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        all_chunks = ["chunk1", "chunk2", "chunk3", "chunk4", "chunk5"]
        fake_llm = _FakeStreamingLLM(all_chunks)
        handler = StandardCitationHandler(fake_llm)

        cb, state, _sock = _build_callback_and_state("rid-mid")
        handler.set_stream_callback(cb)

        # Termination flips after chunk 2 — keep a counter.
        seen = {"count": 0}

        def termination_after_two(_rid):
            seen["count"] += 1
            return seen["count"] > 2

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            side_effect=termination_after_two,
        ):
            with pytest.raises(ResearchTerminatedException):
                handler._invoke_with_streaming("prompt")

        # Streamed exactly 3 chunks (the third triggered termination after
        # being buffered) — the remaining two never ran.
        assert state["chunks_sent"] == 3
        assert "".join(state["chunks"]) == "chunk1chunk2chunk3"
        # Critically: the LLM's ``invoke()`` fallback must NOT have run —
        # otherwise termination would silently complete the response.
        assert fake_llm.invoke_called is False

    def test_no_termination_completes_normally(self):
        """Sanity: with termination always False, streaming completes and
        returns the full joined text."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        chunks = ["a", "b", "c"]
        handler = StandardCitationHandler(_FakeStreamingLLM(chunks))
        cb, state, _sock = _build_callback_and_state("rid-ok")
        handler.set_stream_callback(cb)

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            return_value=False,
        ):
            result = handler._invoke_with_streaming("prompt")

        assert result == "abc"
        assert state["chunks_sent"] == 3
        assert "".join(state["chunks"]) == "abc"


class TestCarryBufferFlush:
    """Regression guard: carry-buffer flush before the is_final sentinel.

    When inline-citation hyperlinking is active, a chunk ending in a
    partial bracket token like ``[12`` is held in the carry buffer until
    the closing ``]`` arrives. If the LLM stream ends while a partial
    token is still held, the completion finalizer
    (``_save_chat_message_and_context``) flushes it via
    ``streaming_state['_flush_carry']`` before emitting the ``is_final``
    sentinel — otherwise the trailing text is silently dropped from the
    client's accumulated view.
    """

    def _build_with_citations(self, research_id):
        from local_deep_research.web.services.research_service import (
            _make_chat_stream_callback,
        )

        streaming_state = {
            "chunks_sent": 0,
            "chunks": [],
            "_bytes": 0,
            "_truncated": False,
        }
        sock = MagicMock()
        # Non-empty sources so the carry path activates; a pass-through
        # formatter lets us assert on the raw carry behaviour without
        # depending on CitationFormatter's substitution details.
        formatter = MagicMock()
        formatter.apply_inline_hyperlinks.side_effect = lambda text, sources: (
            text
        )
        cb = _make_chat_stream_callback(
            research_id,
            streaming_state,
            sock,
            source_resolver=lambda: [{"url": "https://example.com"}],
            formatter=formatter,
        )
        return cb, streaming_state, sock

    def _emitted_chunks(self, sock):
        return "".join(
            call.args[2]["chunk"]
            for call in sock.emit_to_subscribers.call_args_list
            if call.args and call.args[0] == "response_chunk"
        )

    def test_flush_carry_is_registered(self):
        """The completion finalizer reaches the carry via this hook."""
        _cb, state, _sock = self._build_with_citations("rid-flush-reg")
        assert callable(state.get("_flush_carry"))

    def test_partial_bracket_held_then_flushed(self):
        cb, state, sock = self._build_with_citations("rid-flush")

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            return_value=False,
        ):
            # Stream ends mid-token: "[12" must be held back, not emitted.
            cb("See the source [12")

        # The partial token is withheld from the live socket emit...
        assert "[12" not in self._emitted_chunks(sock)
        # ...but the safe prefix was emitted.
        assert "See the source" in self._emitted_chunks(sock)

        # The finalizer's flush releases the held fragment so the tail
        # isn't lost when the stream ends without a closing ']'.
        assert state["_flush_carry"]() == "[12"

        # Idempotent: a second flush (e.g. duplicate finalizer call)
        # returns nothing rather than re-emitting.
        assert state["_flush_carry"]() == ""

    def test_flush_empty_when_no_partial_token(self):
        """A clean stream (no trailing partial bracket) leaves nothing in
        carry, so the finalizer's flush is a no-op."""
        cb, state, _sock = self._build_with_citations("rid-clean")

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            return_value=False,
        ):
            cb("A complete sentence with no open bracket.")

        assert state["_flush_carry"]() == ""

    def test_carry_caps_and_flushes_on_overflow(self):
        """Regression guard: carry-buffer overflow handling.

        A never-closing '[' followed by a long digit run must NOT grow
        the carry buffer without bound. Past _MAX_CARRY_BYTES the
        fragment is flushed raw to the socket and the carry is cleared,
        so a hostile/misbehaving LLM can't consume unbounded memory.
        """
        from local_deep_research.web.services import research_service

        cb, state, sock = self._build_with_citations("rid-cap")
        # A bracket opener followed by more digits than the carry cap.
        long_token = "[" + "1" * (research_service._MAX_CARRY_BYTES + 50)

        with patch(
            "local_deep_research.web.routes.globals.is_termination_requested",
            return_value=False,
        ):
            cb(long_token)

        # The oversized fragment was flushed to the socket, not held...
        assert long_token in self._emitted_chunks(sock)
        # ...and the carry buffer is empty afterwards (nothing to grow).
        assert state["_flush_carry"]() == ""