DocsGPT/tests/api/answer/test_base_routes.py

"""Unit tests for docsgpt/api/answer/routes/base.py — BaseAnswerResource.

Additional coverage beyond tests/api/answer/routes/test_base.py:
  - _prepare_tool_calls_for_logging: truncation, non-dict items
  - complete_stream: tool_calls, thoughts, structured output, metadata,
    isNoneDoc, GeneratorExit handling, compression metadata
  - process_response_stream: structured answer, incomplete stream
  - error_stream_generate: format
  - check_usage: string boolean parsing ("True" strings)
"""

import json
import uuid
from contextlib import contextmanager
from unittest.mock import MagicMock, patch

import pytest


@pytest.mark.unit
class TestPrepareToolCallsForLogging:
    pass

    def test_empty_list(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            assert resource._prepare_tool_calls_for_logging([]) == []

    def test_none_returns_empty(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            assert resource._prepare_tool_calls_for_logging(None) == []

    def test_truncates_long_result(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            tool_calls = [{"result": "x" * 20000}]
            prepared = resource._prepare_tool_calls_for_logging(tool_calls, max_chars=100)
            assert len(prepared[0]["result"]) == 100

    def test_truncates_result_full(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            tool_calls = [{"result_full": "y" * 20000}]
            prepared = resource._prepare_tool_calls_for_logging(tool_calls, max_chars=50)
            assert len(prepared[0]["result_full"]) == 50

    def test_non_dict_items_wrapped(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            tool_calls = ["string_item", 42]
            prepared = resource._prepare_tool_calls_for_logging(tool_calls)
            assert prepared[0] == {"result": "string_item"}
            assert prepared[1] == {"result": "42"}

    def test_preserves_short_results(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            tool_calls = [{"tool_name": "search", "result": "short text"}]
            prepared = resource._prepare_tool_calls_for_logging(tool_calls)
            assert prepared[0]["result"] == "short text"
            assert prepared[0]["tool_name"] == "search"


@pytest.mark.unit
class TestCompleteStreamToolCalls:
    pass

    def test_streams_tool_calls(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {"answer": "Using tool..."},
                    {"tool_calls": [{"name": "search", "result": "found"}]},
                ]
            )

            stream = list(
                resource.complete_stream(
                    question="Search for X",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=False,
                )
            )
            tool_chunks = [s for s in stream if '"type": "tool_calls"' in s]
            assert len(tool_chunks) == 1

    def test_streams_thought_events(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {"thought": "Let me think..."},
                    {"answer": "Here is the answer"},
                ]
            )

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=False,
                )
            )
            thought_chunks = [s for s in stream if '"type": "thought"' in s]
            assert len(thought_chunks) == 1
            assert "Let me think" in thought_chunks[0]


@pytest.mark.unit
class TestCompleteStreamStructuredOutput:
    pass

    def test_streams_structured_answer(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {
                        "answer": '{"key": "value"}',
                        "structured": True,
                        "schema": {"type": "object"},
                    },
                ]
            )

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=False,
                )
            )
            structured_chunks = [
                s for s in stream if '"type": "structured_answer"' in s
            ]
            assert len(structured_chunks) == 1


@pytest.mark.unit
class TestCompleteStreamMetadata:
    pass

    def test_metadata_collected(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {"metadata": {"search_query": "test"}},
                    {"answer": "result"},
                ]
            )

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=False,
                )
            )
            # Should not crash, metadata handled silently
            answer_chunks = [s for s in stream if '"type": "answer"' in s]
            assert len(answer_chunks) == 1


@pytest.mark.unit
class TestCompleteStreamIsNoneDoc:
    pass

    def test_isNoneDoc_sets_source_to_none(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {"answer": "answer"},
                    {"sources": [{"text": "doc", "source": "real_source"}]},
                ]
            )

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    isNoneDoc=True,
                    should_persist=False,
                )
            )
            # Verify stream completes without error
            end_chunks = [s for s in stream if '"type": "end"' in s]
            assert len(end_chunks) == 1


@pytest.mark.unit
class TestCompleteStreamErrorType:
    pass

    def test_error_type_event_sanitized(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {"type": "error", "error": "API key invalid: sk-xxx"},
                ]
            )

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=False,
                )
            )
            error_chunks = [s for s in stream if '"type": "error"' in s]
            assert len(error_chunks) == 1

    def test_non_error_type_event_passed_through(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {"type": "custom_event", "data": "value"},
                ]
            )

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=False,
                )
            )
            custom_chunks = [s for s in stream if '"type": "custom_event"' in s]
            assert len(custom_chunks) == 1


@pytest.mark.unit
class TestProcessResponseStreamExtended:
    pass

    def test_handles_structured_answer(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            stream = [
                f'data: {json.dumps({"type": "structured_answer", "answer": "{}", "structured": True, "schema": None})}\n\n',
                f'data: {json.dumps({"type": "id", "id": str(uuid.uuid4())})}\n\n',
                f'data: {json.dumps({"type": "end"})}\n\n',
            ]
            result = resource.process_response_stream(iter(stream))
            assert result["answer"] == "{}"
            assert result.get("extra", {}).get("structured") is True

    def test_handles_tool_calls_event(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            stream = [
                f'data: {json.dumps({"type": "answer", "answer": "result"})}\n\n',
                f'data: {json.dumps({"type": "tool_calls", "tool_calls": [{"name": "t1"}]})}\n\n',
                f'data: {json.dumps({"type": "id", "id": "conv1"})}\n\n',
                f'data: {json.dumps({"type": "end"})}\n\n',
            ]
            result = resource.process_response_stream(iter(stream))
            assert result["tool_calls"] == [{"name": "t1"}]

    def test_incomplete_stream(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            stream = [
                f'data: {json.dumps({"type": "answer", "answer": "partial"})}\n\n',
            ]
            result = resource.process_response_stream(iter(stream))
            assert result["error"] == "Stream ended unexpectedly"

    def test_handles_thought_event(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            stream = [
                f'data: {json.dumps({"type": "thought", "thought": "thinking..."})}\n\n',
                f'data: {json.dumps({"type": "end"})}\n\n',
            ]
            result = resource.process_response_stream(iter(stream))
            assert result["thought"] == "thinking..."

    def test_handles_id_prefixed_chunks(self, mock_mongo_db, flask_app):
        """``complete_stream`` emits ``id: <seq>\\n`` before each
        ``data:`` line so reconnects can resume. The non-streaming
        ``/api/answer`` consumer must skip the ``id:`` header (and the
        informational ``message_id`` event) without breaking JSON
        decoding.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            stream = [
                'id: 0\n'
                f'data: {json.dumps({"type": "message_id", "message_id": "abc"})}\n\n',
                'id: 1\n'
                f'data: {json.dumps({"type": "answer", "answer": "Hello "})}\n\n',
                'id: 2\n'
                f'data: {json.dumps({"type": "answer", "answer": "world"})}\n\n',
                'id: 3\n'
                f'data: {json.dumps({"type": "id", "id": "conv-1"})}\n\n',
                'id: 4\n'
                f'data: {json.dumps({"type": "end"})}\n\n',
            ]
            result = resource.process_response_stream(iter(stream))
            assert result["error"] is None
            assert result["answer"] == "Hello world"
            assert result["conversation_id"] == "conv-1"

    def test_skips_keepalive_comment_lines(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            stream = [
                ': keepalive\n\n',
                'id: 0\n'
                f'data: {json.dumps({"type": "answer", "answer": "ok"})}\n\n',
                'id: 1\n'
                f'data: {json.dumps({"type": "end"})}\n\n',
            ]
            result = resource.process_response_stream(iter(stream))
            assert result["answer"] == "ok"
            assert result["error"] is None


@pytest.mark.unit
class TestCheckUsageStringBooleans:
    pass

@pytest.mark.unit
class TestCompleteStreamCompressionMetadata:
    """Cover lines 307-319 (compression metadata persistence in complete_stream)."""

    def test_compression_metadata_persisted(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {"answer": "compressed answer"},
                ]
            )
            mock_agent.compression_metadata = {"ratio": 2.5}
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []

            resource.conversation_service = MagicMock()
            resource.conversation_service.save_conversation.return_value = "conv123"
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv123",
                "message_id": "msg123",
                "request_id": "req123",
            }

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=True,
                    model_id="gpt-4",
                )
            )

            # Verify compression metadata was persisted
            resource.conversation_service.update_compression_metadata.assert_called_once_with(
                "conv123", {"ratio": 2.5}
            )
            resource.conversation_service.append_compression_message.assert_called_once()
            assert mock_agent.compression_saved is True
            end_chunks = [s for s in stream if '"type": "end"' in s]
            assert len(end_chunks) == 1

    def test_compression_metadata_error_handled(self, mock_mongo_db, flask_app):
        """Cover lines 318-322: compression metadata persistence error."""
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter([{"answer": "answer"}])
            mock_agent.compression_metadata = {"ratio": 2.5}
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []

            resource.conversation_service = MagicMock()
            resource.conversation_service.save_conversation.return_value = "conv123"
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv123",
                "message_id": "msg123",
                "request_id": "req123",
            }
            resource.conversation_service.update_compression_metadata.side_effect = (
                Exception("db error")
            )

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=True,
                    model_id="gpt-4",
                )
            )

            # Stream should still complete despite compression error
            end_chunks = [s for s in stream if '"type": "end"' in s]
            assert len(end_chunks) == 1


@pytest.mark.unit
class TestCompleteStreamLogTruncation:
    """Cover line 354: log data truncation for long values."""

    def test_long_response_truncated_in_log(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            long_answer = "x" * 20000
            mock_agent.gen.return_value = iter([{"answer": long_answer}])
            mock_agent.tool_calls = []

            stream = list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id=None,
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=False,
                )
            )

            end_chunks = [s for s in stream if '"type": "end"' in s]
            assert len(end_chunks) == 1


@pytest.mark.unit
class TestCompleteStreamGeneratorExit:
    """Cover lines 360-416 (GeneratorExit handling in complete_stream)."""

    def test_generator_exit_saves_partial_response(self, mock_mongo_db, flask_app):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()

            def gen_with_answers():
                yield {"answer": "partial"}
                yield {"answer": " answer"}
                # Simulating a long stream that gets interrupted
                yield {"answer": " more"}

            mock_agent.gen.return_value = gen_with_answers()
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []

            resource.conversation_service = MagicMock()
            resource.conversation_service.save_conversation.return_value = "conv1"
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv1",
                "message_id": "msg1",
                "request_id": "req1",
            }

            gen = resource.complete_stream(
                question="Q",
                agent=mock_agent,
                conversation_id="conv1",
                user_api_key=None,
                decoded_token={"sub": "u"},
                should_persist=True,
                model_id="gpt-4",
            )

            # Drain the early ``message_id`` event before reading the
            # ``partial`` chunk that this test is asserting on.
            next(gen)
            chunk = next(gen)
            assert "partial" in chunk
            gen.close()  # This triggers GeneratorExit

    def test_generator_exit_with_compression_metadata(self, mock_mongo_db, flask_app):
        """Cover lines 393-411: GeneratorExit with compression metadata."""
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()

            def gen_answers():
                yield {"answer": "partial answer"}

            mock_agent.gen.return_value = gen_answers()
            mock_agent.compression_metadata = {"ratio": 3.0}
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []

            resource.conversation_service = MagicMock()
            resource.conversation_service.save_conversation.return_value = "conv1"
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv1",
                "message_id": "msg1",
                "request_id": "req1",
            }

            gen = resource.complete_stream(
                question="Q",
                agent=mock_agent,
                conversation_id="conv1",
                user_api_key=None,
                decoded_token={"sub": "u"},
                should_persist=True,
                model_id="gpt-4",
                isNoneDoc=True,
            )

            # Skip past the early ``message_id`` event.
            next(gen)
            next(gen)
            gen.close()

    def test_generator_exit_save_error_handled(self, mock_mongo_db, flask_app):
        """Cover lines 412-415: exception during partial save."""
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()

            def gen_answers():
                yield {"answer": "partial"}

            mock_agent.gen.return_value = gen_answers()
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []

            resource.conversation_service = MagicMock()
            resource.conversation_service.save_conversation.side_effect = Exception(
                "save error"
            )
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv1",
                "message_id": "msg1",
                "request_id": "req1",
            }

            gen = resource.complete_stream(
                question="Q",
                agent=mock_agent,
                conversation_id="conv1",
                user_api_key=None,
                decoded_token={"sub": "u"},
                should_persist=True,
                model_id="gpt-4",
            )

            # Skip past the early ``message_id`` event.
            next(gen)
            next(gen)
            gen.close()  # Should not crash even with save error

    def test_generator_exit_before_any_response_journals_error_not_end(
        self, mock_mongo_db, flask_app,
    ):
        """A client disconnect right after the early ``message_id`` frame
        leaves ``response_full`` empty, so finalize never runs. The abort
        handler must journal ``error`` (not ``end``) and flip the row to
        ``failed`` — otherwise a reconnecting client sees a terminal
        ``end`` for a row whose DB status is still non-terminal and the
        UI parks on a blank successful answer.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()

            # An agent that's primed but never yields anything before the
            # client disconnects — keeps ``response_full`` empty.
            def gen_never_yields():
                if False:
                    yield  # pragma: no cover
                return

            mock_agent.gen.return_value = gen_never_yields()
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []

            resource.conversation_service = MagicMock()
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv1",
                "message_id": "msg1",
                "request_id": "req1",
            }

            journaled: list[tuple] = []

            def _capture_record(message_id, sequence_no, event_type, payload):
                journaled.append((message_id, sequence_no, event_type, payload))
                return True

            with patch(
                "docsgpt.api.answer.routes.base.record_event",
                side_effect=_capture_record,
            ):
                gen = resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id="conv1",
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=True,
                    model_id="gpt-4",
                )
                # Drain the early ``message_id`` event, then close before
                # the agent yields anything.
                next(gen)
                gen.close()

            # The early message_id frame got journaled (seq 0); the abort
            # handler must follow with an ``error`` event (NOT ``end``).
            terminal_events = [
                (et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
            ]
            assert len(terminal_events) == 1, (
                f"expected exactly one terminal journal write, got {terminal_events}"
            )
            assert terminal_events[0][0] == "error", (
                f"expected ``error`` terminal but got ``end``: {terminal_events}"
            )
            payload = terminal_events[0][1]
            assert payload.get("type") == "error"
            assert payload.get("code") == "client_disconnect"

            # And the DB row should have been flipped to ``failed`` via
            # finalize_message. The mocked service records the call.
            finalize_calls = (
                resource.conversation_service.finalize_message.call_args_list
            )
            assert len(finalize_calls) == 1
            assert finalize_calls[0].kwargs.get("status") == "failed"

    def test_generator_exit_after_response_still_journals_end(
        self, mock_mongo_db, flask_app,
    ):
        """Regression guard: a disconnect AFTER partial response was
        produced and ``finalize_message`` succeeded must still journal
        ``end`` (the row matches ``complete``). Only the empty-response
        branch flips to ``error``.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.storage.db.repositories.conversations import (
            MessageUpdateOutcome,
        )

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()

            def gen_answers():
                yield {"answer": "partial"}

            mock_agent.gen.return_value = gen_answers()
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []

            resource.conversation_service = MagicMock()
            resource.conversation_service.finalize_message.return_value = (
                MessageUpdateOutcome.UPDATED
            )
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv1",
                "message_id": "msg1",
                "request_id": "req1",
            }

            journaled: list[tuple] = []

            def _capture_record(message_id, sequence_no, event_type, payload):
                journaled.append((message_id, sequence_no, event_type, payload))
                return True

            with patch(
                "docsgpt.api.answer.routes.base.record_event",
                side_effect=_capture_record,
            ):
                gen = resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id="conv1",
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=True,
                    model_id="gpt-4",
                )
                next(gen)  # message_id frame
                next(gen)  # answer frame (consumes ``partial``)
                gen.close()

            terminal_events = [
                (et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
            ]
            assert terminal_events and terminal_events[0][0] == "end", (
                f"finalize succeeded but abort journaled {terminal_events}"
            )

    def test_generator_exit_after_normal_finalize_already_complete_journals_end(
        self, mock_mongo_db, flask_app,
    ):
        """Regression for the race where the normal-path finalize wins
        against a client disconnect.

        Trace: agent finishes, ``complete_stream`` runs the normal-path
        ``finalize_message`` at base.py:632 and flips the row to
        ``complete``. The client TCP-resets before the ``end`` frame can
        be journaled. The GeneratorExit handler calls ``finalize_message``
        again — and the repository, gated by ``only_if_non_terminal``,
        reports ``ALREADY_COMPLETE`` because the row is already at the
        target state. The abort handler must journal ``end``, not
        ``error``: the DB says ``complete``, the reconnecting client
        must see the same.

        Without the fix the abort handler treated ``ALREADY_COMPLETE``
        as a failure and journaled ``error`` — a reconnect would then
        replay a successful completion as a failed answer.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.storage.db.repositories.conversations import (
            MessageUpdateOutcome,
        )

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()

            def gen_answers():
                yield {"answer": "partial"}

            mock_agent.gen.return_value = gen_answers()
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []

            resource.conversation_service = MagicMock()
            # Normal-path call returns UPDATED; abort-handler call sees
            # the row already terminal and returns ALREADY_COMPLETE.
            resource.conversation_service.finalize_message.side_effect = [
                MessageUpdateOutcome.UPDATED,
                MessageUpdateOutcome.ALREADY_COMPLETE,
            ]
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv1",
                "message_id": "msg1",
                "request_id": "req1",
            }

            journaled: list[tuple] = []

            def _capture_record(message_id, sequence_no, event_type, payload):
                journaled.append((message_id, sequence_no, event_type, payload))
                return True

            with patch(
                "docsgpt.api.answer.routes.base.record_event",
                side_effect=_capture_record,
            ):
                gen = resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id="conv1",
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=True,
                    model_id="gpt-4",
                )
                next(gen)  # message_id frame
                next(gen)  # answer frame
                # Pull the ``id`` frame so the normal-path finalize at
                # line 632 actually runs (it sits between the answer
                # frame yield and the id frame yield).
                next(gen)  # id frame — runs normal-path finalize first
                gen.close()  # GeneratorExit at the id frame yield

            terminal_events = [
                (et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
            ]
            assert len(terminal_events) == 1, (
                f"expected exactly one terminal journal write, got "
                f"{terminal_events}"
            )
            assert terminal_events[0][0] == "end", (
                f"row was already ``complete`` but abort journaled "
                f"{terminal_events[0]} — reconnect would surface a "
                f"successful answer as failed"
            )

            # Both finalize_message calls were made: the normal-path
            # one and the abort-handler one. Asserting on side_effect
            # consumption ensures the test really exercised both
            # branches.
            assert (
                resource.conversation_service.finalize_message.call_count == 2
            )


@contextmanager
def _patch_db_session(conn):
    @contextmanager
    def _yield():
        yield conn

    with patch(
        "docsgpt.api.answer.services.conversation_service.db_session",
        _yield,
    ), patch(
        "docsgpt.api.answer.services.conversation_service.db_readonly",
        _yield,
    ), patch(
        # ``record_event`` opens its own short-lived ``db_session`` for
        # cross-connection visibility. In tests we route it back to the
        # same ``pg_conn`` so the journal write can see the message row
        # the conversation_service just wrote in this transaction.
        "docsgpt.streaming.message_journal.db_session",
        _yield,
    ), patch(
        # ``complete_stream`` reads ``latest_sequence_no`` via
        # ``db_readonly`` to seed continuation runs. Same patch reason
        # as the journal — keep the read on the same pg_conn so it sees
        # uncommitted writes from this transaction.
        "docsgpt.api.answer.routes.base.db_readonly",
        _yield,
    ), patch(
        # The terminal ``stream_answer`` user_logs write opens its own
        # ``db_session``. Left unpatched it is a *second* connection that
        # blocks on the uncommitted ``users`` row this transaction just
        # inserted (via the ``ensure_user_exists`` trigger) until the
        # statement timeout fires — ~30s per test, swallowed by the
        # caller's except, so it only ever showed up as slowness.
        "docsgpt.api.answer.routes.base.db_session",
        _yield,
    ):
        yield


def _extract_sse_data(chunk: str) -> str:
    """Pull the ``data:`` payload from an SSE record, ignoring any
    ``id:`` header introduced by the journal wiring.
    """
    for line in chunk.split("\n"):
        if line.startswith("data:"):
            return line[len("data:") :].lstrip()
    return ""


@pytest.mark.unit
class TestCompleteStreamWalAcceptance:
    """Acceptance for the WAL pre-persist behaviour: when the LLM raises
    immediately, the user question is still queryable from PG with
    status='failed' and a meaningful error in metadata."""

    def test_failed_llm_persists_question_with_failed_status(
        self, pg_conn, flask_app,
    ):
        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.storage.db.repositories.conversations import (
            ConversationsRepository,
        )

        with flask_app.app_context():
            resource = BaseAnswerResource()

            mock_agent = MagicMock()
            mock_agent.gen.side_effect = RuntimeError("LLM upstream failed")

            with _patch_db_session(pg_conn):
                stream = list(
                    resource.complete_stream(
                        question="why does the WAL matter?",
                        agent=mock_agent,
                        conversation_id=None,
                        user_api_key=None,
                        decoded_token={"sub": "u-acceptance"},
                        should_persist=True,
                        model_id="gpt-4",
                    )
                )
            error_chunks = [s for s in stream if '"type": "error"' in s]
            assert len(error_chunks) == 1

            from sqlalchemy import text as sql_text
            convs = pg_conn.execute(
                sql_text("SELECT id FROM conversations WHERE user_id = :u"),
                {"u": "u-acceptance"},
            ).fetchall()
            assert len(convs) == 1
            conv_id = str(convs[0][0])
            msgs = ConversationsRepository(pg_conn).get_messages(conv_id)
            assert len(msgs) == 1
            assert msgs[0]["prompt"] == "why does the WAL matter?"
            assert msgs[0]["status"] == "failed"
            assert "RuntimeError" in msgs[0]["metadata"]["error"]
            assert "LLM upstream failed" in msgs[0]["metadata"]["error"]

    def test_workflow_node_error_persists_as_failed_not_blank_complete(
        self, pg_conn, flask_app,
    ):
        """A workflow node failure must not land as a blank ``complete`` row.

        The engine reports node failures by *yielding* ``{"type": "error"}``
        rather than raising, so the generator returns normally and the turn
        used to be finalized ``complete`` with an empty response. Live, the
        client shows an error bubble; on reload the row mapped to an empty
        answer with no error and no retry affordance — the user saw a blank
        message and re-sent the prompt. Reproduces the 2026-08-01 report.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.storage.db.repositories.conversations import (
            ConversationsRepository,
        )

        with flask_app.app_context():
            resource = BaseAnswerResource()

            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [{"type": "error", "error": "No LLM class found for type foundry"}]
            )

            with _patch_db_session(pg_conn):
                stream = list(
                    resource.complete_stream(
                        question="hello",
                        agent=mock_agent,
                        conversation_id=None,
                        user_api_key=None,
                        decoded_token={"sub": "u-wf-error"},
                        should_persist=True,
                        model_id="gpt-4",
                    )
                )
            assert len([s for s in stream if '"type": "error"' in s]) == 1

            from sqlalchemy import text as sql_text
            convs = pg_conn.execute(
                sql_text("SELECT id FROM conversations WHERE user_id = :u"),
                {"u": "u-wf-error"},
            ).fetchall()
            assert len(convs) == 1
            msgs = ConversationsRepository(pg_conn).get_messages(str(convs[0][0]))
            assert len(msgs) == 1
            assert msgs[0]["prompt"] == "hello"
            assert msgs[0]["status"] == "failed", (
                "a turn that produced no answer and emitted an error must be "
                "failed, so history renders the error and a retry button"
            )
            assert "foundry" in msgs[0]["metadata"]["error"]

    def test_error_after_partial_answer_keeps_the_answer(
        self, pg_conn, flask_app,
    ):
        """A late error must not discard text the user already received.

        Only the *blank* turn is a failure. If a workflow produced output and
        then a downstream node failed, the row stays ``complete`` so the
        partial answer still renders; the error was already surfaced live.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.storage.db.repositories.conversations import (
            ConversationsRepository,
        )

        with flask_app.app_context():
            resource = BaseAnswerResource()

            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [
                    {"answer": "partial result"},
                    {"type": "error", "error": "node 2 blew up"},
                ]
            )

            with _patch_db_session(pg_conn):
                list(
                    resource.complete_stream(
                        question="run the workflow",
                        agent=mock_agent,
                        conversation_id=None,
                        user_api_key=None,
                        decoded_token={"sub": "u-wf-partial"},
                        should_persist=True,
                        model_id="gpt-4",
                    )
                )

            from sqlalchemy import text as sql_text
            convs = pg_conn.execute(
                sql_text("SELECT id FROM conversations WHERE user_id = :u"),
                {"u": "u-wf-partial"},
            ).fetchall()
            msgs = ConversationsRepository(pg_conn).get_messages(str(convs[0][0]))
            assert msgs[0]["status"] == "complete"
            assert msgs[0]["response"] == "partial result"
            # Still recorded, so the failure is greppable in review queries.
            assert "node 2 blew up" in msgs[0]["metadata"]["error"]

    def test_workflow_error_fails_the_row_on_the_non_wal_path_too(
        self, pg_conn, flask_app,
    ):
        """The same guarantee when no placeholder row was reserved.

        ``save_conversation`` has its own insert path (used when the WAL
        reservation failed, or on a continuation carrying no
        ``reserved_message_id``). It took no status, so the row landed on the
        column default ``complete`` — leaving exactly the blank bubble this
        changeset removes on the other branch.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.storage.db.repositories.conversations import (
            ConversationsRepository,
        )

        with flask_app.app_context():
            resource = BaseAnswerResource()

            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [{"type": "error", "error": "CEL error in node 'Build reply'"}]
            )

            with _patch_db_session(pg_conn), patch.object(
                resource.conversation_service,
                "save_user_question",
                side_effect=RuntimeError("WAL reservation unavailable"),
            ):
                list(
                    resource.complete_stream(
                        question="hello",
                        agent=mock_agent,
                        conversation_id=None,
                        user_api_key=None,
                        decoded_token={"sub": "u-wf-nonwal"},
                        should_persist=True,
                        model_id="gpt-4",
                    )
                )

            from sqlalchemy import text as sql_text
            convs = pg_conn.execute(
                sql_text("SELECT id FROM conversations WHERE user_id = :u"),
                {"u": "u-wf-nonwal"},
            ).fetchall()
            assert len(convs) == 1
            msgs = ConversationsRepository(pg_conn).get_messages(str(convs[0][0]))
            assert len(msgs) == 1
            assert msgs[0]["status"] == "failed"
            assert "CEL error" in msgs[0]["metadata"]["error"]

    def test_tool_approval_event_only_fires_when_state_saved(
        self, pg_conn, flask_app,
    ):
        """A `tool.approval.required` notification with no resumable
        ``pending_tool_state`` row would deep-link the user to a 404.
        Gate the publish on save_state actually committing.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            # Drive the agent into the paused branch with a single
            # ``tool_calls_pending`` event.
            mock_agent.gen.return_value = iter(
                [
                    {
                        "type": "tool_calls_pending",
                        "data": {"pending_tool_calls": [{"call_id": "c1"}]},
                    }
                ]
            )
            mock_agent._pending_continuation = {
                "messages": [],
                "tools_dict": {},
                "pending_tool_calls": [{"call_id": "c1"}],
            }
            mock_agent.tool_calls = []
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False

            published: list = []

            def _capture(*args, **kwargs):
                published.append(args)

            with _patch_db_session(pg_conn), patch(
                "docsgpt.api.answer.routes.base.publish_user_event",
                side_effect=_capture,
            ), patch(
                "docsgpt.api.answer.services.continuation_service."
                "ContinuationService.save_state",
                side_effect=RuntimeError("PG outage"),
            ):
                list(
                    resource.complete_stream(
                        question="run my tool",
                        agent=mock_agent,
                        conversation_id=None,
                        user_api_key=None,
                        decoded_token={"sub": "u-tap"},
                        should_persist=True,
                        model_id="gpt-4",
                    )
                )

            # No tool.approval.required publish when save_state failed.
            event_types = [a[1] for a in published if len(a) >= 2]
            assert "tool.approval.required" not in event_types

    def test_continuation_seeds_sequence_no_from_journal_high_water_mark(
        self, pg_conn, flask_app,
    ):
        """A resumed (tool-actions) stream must continue numbering past
        the original run's max ``sequence_no``. Otherwise the second
        invocation collides on the duplicate-PK and silently drops
        every journal write past the resume point.

        We pre-seed the journal with synthetic rows simulating an
        original run, then invoke ``complete_stream`` with
        ``_continuation`` set and assert seq numbering picks up past
        the high-water mark.
        """
        import uuid as _uuid

        from sqlalchemy import text as sql_text

        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.storage.db.repositories.message_events import (
            MessageEventsRepository,
        )

        # Pre-seed: insert a parent conversation + message + a few
        # journal rows so latest_sequence_no returns 7.
        user_id = "u-resume"
        conv_id = _uuid.uuid4()
        message_id = _uuid.uuid4()
        pg_conn.execute(
            sql_text("INSERT INTO users (user_id) VALUES (:u)"),
            {"u": user_id},
        )
        pg_conn.execute(
            sql_text(
                "INSERT INTO conversations (id, user_id, name) "
                "VALUES (:id, :u, 'pre-seed')"
            ),
            {"id": conv_id, "u": user_id},
        )
        pg_conn.execute(
            sql_text(
                "INSERT INTO conversation_messages "
                "(id, conversation_id, user_id, position, prompt) "
                "VALUES (:id, :c, :u, 0, 'q')"
            ),
            {"id": message_id, "c": conv_id, "u": user_id},
        )
        repo = MessageEventsRepository(pg_conn)
        for seq in range(8):  # rows 0..7
            repo.record(str(message_id), seq, "answer", {"chunk": str(seq)})
        original_max = repo.latest_sequence_no(str(message_id))
        assert original_max == 7

        with flask_app.app_context():
            resource = BaseAnswerResource()
            cont_agent = MagicMock()
            cont_agent.gen_continuation.return_value = iter(
                [{"answer": "d"}, {"answer": "e"}]
            )
            cont_agent.tool_calls = []
            cont_agent.compression_metadata = None
            cont_agent.compression_saved = False
            cont_agent.tool_executor = None

            with _patch_db_session(pg_conn):
                list(
                    resource.complete_stream(
                        question="",
                        agent=cont_agent,
                        conversation_id=str(conv_id),
                        user_api_key=None,
                        decoded_token={"sub": user_id},
                        should_persist=True,
                        model_id="gpt-4",
                        _continuation={
                            "messages": [],
                            "tools_dict": {},
                            "pending_tool_calls": [],
                            "tool_actions": [],
                            "reserved_message_id": str(message_id),
                            "request_id": "req-resume",
                        },
                    )
                )

        new_max = MessageEventsRepository(pg_conn).latest_sequence_no(
            str(message_id)
        )
        # Continuation extended the journal — the new high-water mark
        # is strictly greater than the seeded ``original_max=7``,
        # confirming the allocator picked up past the resume point.
        assert new_max is not None and new_max > original_max

    def test_request_id_consistent_across_sse_event_and_wal_row(
        self, pg_conn, flask_app,
    ):
        """The early ``message_id`` SSE event reports the same
        ``request_id`` that ``save_user_question`` writes onto the WAL
        row, so client-side correlation, ``token_usage`` joins, and
        ``count_in_range``'s DISTINCT all line up.
        """
        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.storage.db.repositories.conversations import (
            ConversationsRepository,
        )

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter([{"answer": "ok"}])
            mock_agent.tool_calls = []
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False

            with _patch_db_session(pg_conn):
                stream = list(
                    resource.complete_stream(
                        question="hello",
                        agent=mock_agent,
                        conversation_id=None,
                        user_api_key=None,
                        decoded_token={"sub": "u-request-id"},
                        should_persist=True,
                        model_id="gpt-4",
                    )
                )

            sse_events = [
                json.loads(_extract_sse_data(s))
                for s in stream
                if "data:" in s
            ]
            early_events = [e for e in sse_events if e.get("type") == "message_id"]
            assert len(early_events) == 1
            sse_request_id = early_events[0]["request_id"]
            assert sse_request_id

            from sqlalchemy import text as sql_text
            convs = pg_conn.execute(
                sql_text("SELECT id FROM conversations WHERE user_id = :u"),
                {"u": "u-request-id"},
            ).fetchall()
            assert len(convs) == 1
            msgs = ConversationsRepository(pg_conn).get_messages(str(convs[0][0]))
            assert len(msgs) == 1
            assert msgs[0]["request_id"] == sse_request_id


@pytest.mark.unit
class TestStreamingHeartbeatSeed:
    """Regression guard: the reserved row must carry a fresh
    ``last_heartbeat_at`` from generation start so the reconciler doesn't fall
    back to ``timestamp`` (creation time) on slow LLM cold-starts or while a
    reasoning model streams only ``thought`` chunks. The heartbeat is seeded
    once before the first chunk and re-stamped when the row flips to
    ``streaming``; the ``pending → streaming`` status transition itself still
    fires exactly once on the first ``answer``/``sources`` chunk.
    """

    def test_heartbeat_seeded_at_generation_start_and_on_first_chunk(
        self, mock_mongo_db, flask_app,
    ):
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [{"answer": "a"}, {"answer": "b"}]
            )
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []
            mock_agent.tool_executor = None

            resource.conversation_service = MagicMock()
            resource.conversation_service.save_conversation.return_value = "conv1"
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv1",
                "message_id": "msg1",
                "request_id": "req1",
            }

            list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id="conv1",
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=True,
                    model_id="gpt-4",
                )
            )

            # update_message_status flips the row to ``streaming`` exactly
            # once (idempotent via streaming_marked) on the first answer chunk.
            status_calls = [
                c for c in resource.conversation_service
                .update_message_status.call_args_list
                if c.args[1] == "streaming"
            ]
            assert len(status_calls) == 1
            assert status_calls[0].args[0] == "msg1"

            # Heartbeat is stamped twice for this 2-answer-chunk stream: once at
            # the generation-start seed (before the first chunk, while still
            # ``pending``) and once when the first answer chunk flips the row to
            # ``streaming``. The second answer chunk does NOT re-stamp — the
            # per-chunk ``_heartbeat_streaming`` pump is throttled by
            # STREAM_HEARTBEAT_INTERVAL and the two chunks fall inside one
            # interval under real ``time.monotonic``.
            hb_calls = (
                resource.conversation_service.heartbeat_message.call_args_list
            )
            assert len(hb_calls) == 2
            assert all(c.args[0] == "msg1" for c in hb_calls)

    def test_heartbeat_seed_skipped_without_reserved_message_id(
        self, mock_mongo_db, flask_app,
    ):
        """No DB-backed message row → no heartbeat call (and no error)."""
        from docsgpt.api.answer.routes.base import BaseAnswerResource

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter([{"answer": "a"}])
            mock_agent.compression_metadata = None
            mock_agent.compression_saved = False
            mock_agent.tool_calls = []
            mock_agent.tool_executor = None

            resource.conversation_service = MagicMock()
            resource.conversation_service.save_conversation.return_value = "conv1"
            # save_user_question returns no message_id → reservation absent
            resource.conversation_service.save_user_question.return_value = {
                "conversation_id": "conv1",
                "message_id": None,
                "request_id": "req1",
            }

            list(
                resource.complete_stream(
                    question="Q",
                    agent=mock_agent,
                    conversation_id="conv1",
                    user_api_key=None,
                    decoded_token={"sub": "u"},
                    should_persist=True,
                    model_id="gpt-4",
                )
            )

            resource.conversation_service.heartbeat_message.assert_not_called()


class TestCompleteStreamSupersededLookupIsGuarded:
    """A DB failure in the supersede lookup must not error a finished turn.

    The ``NOT_FOUND`` branch used to only log. It now consults
    ``was_superseded`` to tell a replaced turn from an orphaned answer — over
    its own connection, inside the streaming ``try``. Unguarded, a blip there
    (or a deploy where the app runs ahead of migration 0030 and the
    ``superseded_messages`` table does not exist yet) escapes to the generic
    handler: the client gets a terminal ``error`` after the whole answer, the
    row is finalized a second time, a continuation's resume claim is released,
    and the ``answer_persist_failed`` alert this branch exists to raise is
    swallowed.
    """

    def test_lookup_failure_falls_back_instead_of_erroring_the_stream(
        self, pg_conn, flask_app, caplog,
    ):
        import logging

        from docsgpt.api.answer.routes.base import BaseAnswerResource
        from docsgpt.api.answer.services.conversation_service import (
            ConversationService,
        )
        from docsgpt.storage.db.repositories.conversations import (
            MessageUpdateOutcome,
        )

        with flask_app.app_context():
            resource = BaseAnswerResource()
            mock_agent = MagicMock()
            mock_agent.gen.return_value = iter(
                [{"type": "answer", "answer": "the answer"}]
            )

            with _patch_db_session(pg_conn), patch.object(
                ConversationService,
                "finalize_message",
                return_value=MessageUpdateOutcome.NOT_FOUND,
            ), patch.object(
                ConversationService,
                "was_superseded",
                side_effect=RuntimeError("server closed the connection"),
            ), caplog.at_level(logging.ERROR):
                stream = list(
                    resource.complete_stream(
                        question="does the guard hold?",
                        agent=mock_agent,
                        conversation_id=None,
                        user_api_key=None,
                        decoded_token={"sub": "u-superseded-guard"},
                        should_persist=True,
                        model_id="gpt-4",
                    )
                )

        # The answer was already streamed; the row is gone either way. What
        # must not happen is retracting it with a terminal error frame.
        assert [s for s in stream if '"type": "error"' in s] == []
        # And the alert must still fire, on the orphaned-answer path.
        assert any(
            "answer_persist_failed" in r.getMessage() for r in caplog.records
        )