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

"""
Tests for send_message error-handling and concurrency-cap behavior.

Specifically covers:

* atomic write of user-message + research row in one transaction;
  no orphan ChatMessage when the research insert raises.
* ValueError("not found") from a session-deleted race surfaces
  as HTTP 404, not 500.
* DuplicateResearchError from start_research_process is caught
  explicitly, both side-effect rows are rolled back, response is 409.
* global per-user concurrent-research cap enforced for chat-triggered
  research (via UserActiveResearch + app.max_concurrent_researches).

tests/chat/test_chat_concurrency_guard.py covers the per-session 409
invariant; these tests cover the remaining error branches.
"""

import json
from unittest.mock import patch


def _create_session(client, query="Test"):
    resp = client.post(
        "/api/chat/sessions",
        json={"initial_query": query},
        content_type="application/json",
    )
    assert resp.status_code == 200, resp.data
    return json.loads(resp.data)["session_id"]


class TestSessionDeletedRace:
    """When the session row vanishes between the existence-check and the
    UPDATE...RETURNING inside insert_message_in_db, the route must return
    HTTP 404, not a generic 500.
    """

    def test_deleted_session_after_existence_check_returns_404(
        self, authenticated_client
    ):
        session_id = _create_session(authenticated_client)

        # Simulate the race: get_session sees the session, but by the
        # time insert_message_in_db runs the UPDATE...RETURNING the
        # row is gone. We patch the helper to raise the same
        # "not found" ValueError the real code raises when
        # scalar_one_or_none() returns None.
        with patch(
            "local_deep_research.chat.service.ChatService.insert_message_in_db",
            side_effect=ValueError(f"Chat session {session_id} not found"),
        ):
            resp = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={"content": "hello", "trigger_research": True},
                content_type="application/json",
            )

        assert resp.status_code == 404, resp.data
        body = json.loads(resp.data)
        assert body["success"] is False
        assert "not found" in body["error"].lower()


class TestDuplicateResearchErrorCleanup:
    """start_research_process can raise DuplicateResearchError,
    which inherits from Exception and is NOT in ROUTE_EXCEPTIONS. It is
    caught explicitly, returns 409, and rolls back the user
    message + research row this request just committed (per the
    DuplicateResearchError docstring's 'do not touch
    UserActiveResearch' rule)."""

    def test_duplicate_research_returns_409_not_500(self, authenticated_client):
        from local_deep_research.exceptions import DuplicateResearchError

        session_id = _create_session(authenticated_client)

        with patch(
            "local_deep_research.chat.routes.start_research_process",
            side_effect=DuplicateResearchError("simulated"),
        ):
            resp = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={"content": "first", "trigger_research": True},
                content_type="application/json",
            )

        assert resp.status_code == 409, resp.data
        body = json.loads(resp.data)
        assert body["success"] is False
        assert "in progress" in body["error"].lower()

    def test_duplicate_research_does_not_orphan_user_message(
        self, authenticated_client
    ):
        """After a DuplicateResearchError-induced 409, the chat must
        not contain the rejected user message (i.e. cleanup ran)."""
        from local_deep_research.exceptions import DuplicateResearchError

        session_id = _create_session(authenticated_client)

        with patch(
            "local_deep_research.chat.routes.start_research_process",
            side_effect=DuplicateResearchError("simulated"),
        ):
            blocked = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={
                    "content": "should-be-cleaned-up",
                    "trigger_research": True,
                },
                content_type="application/json",
            )
            assert blocked.status_code == 409

        # Read messages back; the rejected user message must not be
        # in the durable message list. The ChatSession.message_count
        # would also be off-by-one if the cleanup had failed.
        list_resp = authenticated_client.get(
            f"/api/chat/sessions/{session_id}/messages"
        )
        assert list_resp.status_code == 200
        msgs = json.loads(list_resp.data)["messages"]
        user_contents = [m["content"] for m in msgs if m["role"] == "user"]
        assert "should-be-cleaned-up" not in user_contents


class TestGlobalPerUserCap:
    """send_message must respect app.max_concurrent_researches the
    same way research_routes.start_research does; otherwise multiple
    chat tabs would let a user bypass the cap.

    The end-to-end "saturate UserActiveResearch then verify 429"
    scenario requires injecting rows into the per-user encrypted DB
    from outside the request context, which is awkward. Instead we
    patch the SettingsManager to return cap=0 so any non-negative
    active_count trips the 429 — this exercises the same code path
    (the cap-comparison in send_message) without the encrypted-DB
    plumbing.
    """

    def test_chat_send_blocks_at_429_when_cap_exceeded(
        self, authenticated_client
    ):
        session_id = _create_session(authenticated_client)

        with (
            patch(
                "local_deep_research.settings.manager.SettingsManager.get_setting",
                return_value=0,
            ),
            patch("local_deep_research.chat.routes.start_research_process"),
        ):
            resp = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={"content": "blocked-by-cap", "trigger_research": True},
                content_type="application/json",
            )

        assert resp.status_code == 429, resp.data
        body = json.loads(resp.data)
        assert body["success"] is False
        assert "concurrent research limit" in body["error"].lower()

    def test_chat_send_inserts_active_research_so_cap_counts_it(
        self, authenticated_client
    ):
        """A chat research send must insert a UserActiveResearch row so it
        counts toward the per-user cap. With cap=1, a first send (session A)
        succeeds and a second send on a DIFFERENT session (B) is then blocked
        at 429 — which only happens if the first send actually recorded an
        active-research row."""

        def _cap_one(key, default=None, *args, **kwargs):
            # Force the per-user cap to 1; every other setting falls back to
            # the caller-provided default so the send path is otherwise normal.
            if key == "app.max_concurrent_researches":
                return 1
            return default

        session_a = _create_session(authenticated_client)
        session_b = _create_session(authenticated_client)

        with (
            patch(
                "local_deep_research.settings.manager.SettingsManager.get_setting",
                side_effect=_cap_one,
            ),
            # No-op the spawn so the first research stays "in progress"
            # (its UserActiveResearch row is never cleaned up by completion).
            patch("local_deep_research.chat.routes.start_research_process"),
            # Treat the first research's (stubbed-spawn) thread as alive.
            # Two independent cleanup paths would otherwise delete the first
            # send's UserActiveResearch row before the second send's cap check
            # sees it (dropping active_count back to 0 → flaky 200 instead of
            # 429). Because the spawn is stubbed, the row has no real thread /
            # no _active_research entry, so it looks "done" to both paths:
            #
            #   1. reclaim_stale_user_active_research (run inline by
            #      chat.routes.send_message) — flips IN_PROGRESS rows whose
            #      thread is dead, gated only by a 30s grace window on
            #      started_at. Patching is_research_thread_alive -> True
            #      short-circuits before the grace check, so a >30s GC stall
            #      between the two sequential posts can no longer trip it.
            #
            #   2. cleanup_completed_research before_request middleware
            #      (web/auth/cleanup_middleware.py) — sampled on ~1% of
            #      requests, deletes UserActiveResearch rows whose
            #      research_id is not in _active_research. Patching
            #      is_research_active -> True makes the middleware skip the
            #      row, so the 1% sample can no longer trip it.
            #
            # Both patches model the production invariant (a live research's
            # row is left alone) without the timing/lottery dependencies that
            # surfaced under heavy CI load on the 36k-test xdist lane.
            patch(
                "local_deep_research.web.routes.globals.is_research_thread_alive",
                return_value=True,
            ),
            patch(
                "local_deep_research.web.routes.globals.is_research_active",
                return_value=True,
            ),
        ):
            first = authenticated_client.post(
                f"/api/chat/sessions/{session_a}/messages",
                json={"content": "first research", "trigger_research": True},
                content_type="application/json",
            )
            second = authenticated_client.post(
                f"/api/chat/sessions/{session_b}/messages",
                json={"content": "second research", "trigger_research": True},
                content_type="application/json",
            )

        assert first.status_code == 200, first.data
        assert second.status_code == 429, second.data
        body = json.loads(second.data)
        assert body["success"] is False
        assert "concurrent research limit" in body["error"].lower()


class TestSystemAtCapacityErrorCleanup:
    """start_research_process can also raise
    SystemAtCapacityError when the global semaphore is exhausted AFTER
    the pre-spawn cap check passes (race against a parallel request).
    The branch is symmetric with DuplicateResearchError — both must
    return 429-class status and clean up the user-message + research
    rows the request just committed.
    """

    def test_capacity_returns_429_not_500(self, authenticated_client):
        from local_deep_research.exceptions import SystemAtCapacityError

        session_id = _create_session(authenticated_client)

        with patch(
            "local_deep_research.chat.routes.start_research_process",
            side_effect=SystemAtCapacityError("simulated"),
        ):
            resp = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={
                    "content": "blocked-by-capacity",
                    "trigger_research": True,
                },
                content_type="application/json",
            )

        assert resp.status_code == 429, resp.data
        body = json.loads(resp.data)
        assert body["success"] is False
        # The route surfaces a capacity-related error string; assert on
        # "capacity"/"busy"/"try again" so the exact wording can evolve.
        msg = body["error"].lower()
        assert "capacity" in msg or "busy" in msg or "try again" in msg

    def test_capacity_does_not_orphan_user_message(self, authenticated_client):
        """After a SystemAtCapacityError-induced 429, the rejected
        user message must not appear in the chat's durable message
        list. Mirrors TestDuplicateResearchErrorCleanup — if the
        ``_cleanup_chat_send_rows("capacity")`` branch regresses, the
        user-message and ResearchHistory rows become orphans and
        ``ChatSession.message_count`` is off-by-one.
        """
        from local_deep_research.exceptions import SystemAtCapacityError

        session_id = _create_session(authenticated_client)

        with patch(
            "local_deep_research.chat.routes.start_research_process",
            side_effect=SystemAtCapacityError("simulated"),
        ):
            blocked = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={
                    "content": "should-be-cleaned-up-on-capacity",
                    "trigger_research": True,
                },
                content_type="application/json",
            )
            assert blocked.status_code == 429

        list_resp = authenticated_client.get(
            f"/api/chat/sessions/{session_id}/messages"
        )
        assert list_resp.status_code == 200
        msgs = json.loads(list_resp.data)["messages"]
        user_contents = [m["content"] for m in msgs if m["role"] == "user"]
        assert "should-be-cleaned-up-on-capacity" not in user_contents