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

# allow: no-sut-import — black-box HTTP test; drives real routes through the Flask test client
"""Tests for per-attempt delete + retry endpoints.

Issue #4659: when a chat-triggered research fails (OoM, crash, dead LLM
endpoint), the user message stays in the thread and there's no in-chat
way to retry without copy-pasting. The new endpoints

    DELETE /api/chat/sessions/<sid>/attempts/<rid>
    POST   /api/chat/sessions/<sid>/attempts/<rid>/retry

let the client delete or re-run a single attempt atomically.

Coverage in this file:

- Service-level (``ChatService.delete_attempt`` /
  ``get_original_attempt_query``): happy path, missing, in-progress
  refusal, stale-IN_PROGRESS reclaim, message_count decrement,
  research_meta fallback.
- HTTP-level: happy paths, 404 / 409 mappings, CSRF enforcement,
  concurrent-message 409 race via the partial unique index.
"""

import json
import uuid
from datetime import datetime, UTC
from contextlib import contextmanager
from unittest.mock import patch

import pytest

from src.local_deep_research.chat.service import (
    AttemptInProgress,
    AttemptNotFound,
    ChatService,
)
from src.local_deep_research.constants import ResearchStatus
from src.local_deep_research.database.models import (
    ChatMessage,
    ChatProgressStep,
    ChatRole,
    ChatSession,
    ChatSessionStatus,
    ResearchHistory,
)


# ---------------------------------------------------------------------------
# Seed helpers — used by both service-level and HTTP-level tests.
# ---------------------------------------------------------------------------


@contextmanager
def _service_db(SessionLocal):
    """Redirect ``ChatService.get_user_db_session`` to the test DB.

    ChatService opens the per-user encrypted DB by username; the test
    DB from ``setup_database_for_all_tests`` is a plain SQLite file.
    Same pattern as ``tests/chat/test_chat_delete_terminates_research.py``.
    """
    with patch(
        "src.local_deep_research.chat.service.get_user_db_session"
    ) as ctx:
        ctx.return_value.__enter__.return_value = SessionLocal()
        ctx.return_value.__exit__.return_value = False
        yield


def _seed_attempt(
    db,
    username: str,
    session_id: str,
    research_id: str,
    *,
    research_status: str = ResearchStatus.FAILED.value,
    user_content: str = "What is quantum computing?",
    assistant_content: str | None = "Quantum computing uses qubits...",
    include_research_meta: bool = True,
    user_message_id: str | None = None,
):
    """Insert a full attempt (user msg + research + optional assistant msg).

    Returns ``(user_message_id, assistant_message_id_or_None)``.
    """
    # Look up the session to find the next sequence_number.
    session = db.query(ChatSession).filter_by(id=session_id).first()
    next_seq = (session.message_count if session else 0) + 1

    user_msg_id = user_message_id or str(uuid.uuid4())
    # Mirror production: _spawn_chat_research inserts the user message with
    # research_id=NULL and links it to the attempt only via
    # research_meta.submission.message_id. Only legacy rows that predate the
    # research_meta field carried research_id on the user message — model
    # those when include_research_meta is False so the fallback lookup path
    # in get_original_attempt_query / delete_attempt stays exercised.
    user_research_id = None if include_research_meta else research_id
    db.add(
        ChatMessage(
            id=user_msg_id,
            session_id=session_id,
            research_id=user_research_id,
            role=ChatRole.USER.value,
            message_type="query",
            content=user_content,
            sequence_number=next_seq,
        )
    )

    research_meta = None
    if include_research_meta:
        research_meta = {
            "submission": {
                "chat_session_id": session_id,
                "message_id": user_msg_id,
                "research_mode": "quick",
            }
        }

    db.add(
        ResearchHistory(
            id=research_id,
            query=user_content,
            mode="quick",
            status=research_status,
            created_at=datetime.now(UTC).isoformat(),
            chat_session_id=session_id,
            research_meta=research_meta,
        )
    )

    assistant_msg_id = None
    if assistant_content is not None:
        assistant_msg_id = str(uuid.uuid4())
        db.add(
            ChatMessage(
                id=assistant_msg_id,
                session_id=session_id,
                research_id=research_id,
                role=ChatRole.ASSISTANT.value,
                message_type="response",
                content=assistant_content,
                sequence_number=next_seq + 1,
            )
        )

    # A progress step tied to this research.
    db.add(
        ChatProgressStep(
            id=str(uuid.uuid4()),
            research_id=research_id,
            session_id=session_id,
            phase="search",
            content="Searching...",
            sequence_number=1,
        )
    )

    # Bump session message_count to match what we inserted.
    if session:
        session.message_count = session.message_count + (
            2 if assistant_content is not None else 1
        )

    db.commit()
    return user_msg_id, assistant_msg_id


def _make_session(db, username: str, title: str = "test chat") -> str:
    session_id = str(uuid.uuid4())
    db.add(
        ChatSession(
            id=session_id,
            title=title,
            status=ChatSessionStatus.ACTIVE.value,
            message_count=0,
        )
    )
    db.commit()
    return session_id


# ===========================================================================
# Service-level: ChatService.delete_attempt
# ===========================================================================


class TestDeleteAttemptService:
    """Direct ChatService.delete_attempt tests."""

    def test_delete_attempt_happy_path(self, setup_database_for_all_tests):
        """Deleting a FAILED attempt removes the user msg, assistant msg,
        progress step, and research row; decrements message_count by 2."""
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_id = _make_session(db, username)
            research_id = str(uuid.uuid4())
            user_msg_id, assistant_msg_id = _seed_attempt(
                db,
                username,
                session_id,
                research_id,
                research_status=ResearchStatus.FAILED.value,
            )
            # Sanity: 2 messages (user + assistant) + 1 step + 1 research
            # exist. Count by session_id, NOT research_id — the user message
            # is seeded with research_id=NULL (production behaviour), so a
            # research_id filter would only see the assistant row.
            assert (
                db.query(ChatMessage).filter_by(session_id=session_id).count()
                == 2
            )
            # Regression guard for #4659: the user message must be linked
            # to the attempt only via research_meta, with research_id NULL.
            user_row = db.query(ChatMessage).filter_by(id=user_msg_id).first()
            assert user_row is not None
            assert user_row.research_id is None
            assert (
                db.query(ChatProgressStep)
                .filter_by(research_id=research_id)
                .count()
                == 1
            )

        with _service_db(SessionLocal):
            result = service.delete_attempt(session_id, research_id)
        assert result is True

        with SessionLocal() as db:
            # All attempt rows gone — including the NULL-research_id user
            # message (the #4659 bug left it orphaned in the thread).
            assert (
                db.query(ChatMessage).filter_by(session_id=session_id).count()
                == 0
            )
            assert (
                db.query(ChatMessage).filter_by(id=user_msg_id).first() is None
            )
            assert (
                db.query(ChatProgressStep)
                .filter_by(research_id=research_id)
                .count()
                == 0
            )
            assert (
                db.query(ResearchHistory).filter_by(id=research_id).first()
                is None
            )
            # Session still exists and message_count was decremented.
            session = db.query(ChatSession).filter_by(id=session_id).first()
            assert session is not None
            assert session.message_count == 0

    def test_delete_attempt_failed_no_assistant_removes_user_msg(
        self, setup_database_for_all_tests
    ):
        """Primary #4659 case: a FAILED attempt with NO assistant reply.

        The user message is the only chat row and is stored with
        research_id=NULL (linked via research_meta.submission.message_id).
        Before the fix, delete_attempt filtered solely on research_id, so it
        removed nothing user-visible and never decremented message_count —
        leaving an orphaned user bubble. Delete must remove the user message
        and decrement the counter by exactly 1.
        """
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_id = _make_session(db, username)
            research_id = str(uuid.uuid4())
            user_msg_id, assistant_msg_id = _seed_attempt(
                db,
                username,
                session_id,
                research_id,
                research_status=ResearchStatus.FAILED.value,
                assistant_content=None,
            )
            assert assistant_msg_id is None
            # Only the user message exists, and it has NULL research_id.
            assert (
                db.query(ChatMessage).filter_by(session_id=session_id).count()
                == 1
            )
            user_row = db.query(ChatMessage).filter_by(id=user_msg_id).first()
            assert user_row is not None and user_row.research_id is None
            assert (
                db.query(ChatSession)
                .filter_by(id=session_id)
                .first()
                .message_count
                == 1
            )

        with _service_db(SessionLocal):
            assert service.delete_attempt(session_id, research_id) is True

        with SessionLocal() as db:
            assert (
                db.query(ChatMessage).filter_by(id=user_msg_id).first() is None
            )
            assert (
                db.query(ChatMessage).filter_by(session_id=session_id).count()
                == 0
            )
            assert (
                db.query(ResearchHistory).filter_by(id=research_id).first()
                is None
            )
            # The counter was decremented by exactly 1 (the orphaned user
            # bubble), not left drifting at 1.
            assert (
                db.query(ChatSession)
                .filter_by(id=session_id)
                .first()
                .message_count
                == 0
            )

    def test_delete_attempt_missing(self, setup_database_for_all_tests):
        """Unknown research_id → AttemptNotFound."""
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_id = _make_session(db, username)

        with _service_db(SessionLocal):
            with pytest.raises(AttemptNotFound):
                service.delete_attempt(session_id, str(uuid.uuid4()))

    def test_delete_attempt_wrong_session(self, setup_database_for_all_tests):
        """Research exists but belongs to a different session_id →
        AttemptNotFound (scoped lookup refuses to return cross-session
        rows)."""
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_a = _make_session(db, username, title="A")
            session_b = _make_session(db, username, title="B")
            research_id = str(uuid.uuid4())
            _seed_attempt(
                db,
                username,
                session_a,
                research_id,
                research_status=ResearchStatus.FAILED.value,
            )

        # Asking for the research under session_b must 404.
        with _service_db(SessionLocal):
            with pytest.raises(AttemptNotFound):
                service.delete_attempt(session_b, research_id)

    def test_delete_attempt_in_progress_with_live_thread(
        self, setup_database_for_all_tests
    ):
        """IN_PROGRESS + alive thread → AttemptInProgress (409 mapping)
        and the worker is flagged for termination."""
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_id = _make_session(db, username)
            research_id = str(uuid.uuid4())
            _seed_attempt(
                db,
                username,
                session_id,
                research_id,
                research_status=ResearchStatus.IN_PROGRESS.value,
            )

        flagged = []
        with (
            patch(
                "src.local_deep_research.chat.service.set_termination_flag",
                side_effect=lambda rid: flagged.append(rid),
            ),
            patch(
                "src.local_deep_research.chat.service.is_research_thread_alive",
                return_value=True,
            ),
            _service_db(SessionLocal),
        ):
            with pytest.raises(AttemptInProgress):
                service.delete_attempt(session_id, research_id)

        # Worker was flagged before the raise.
        assert flagged == [research_id]

        # Row still exists — we refused to hard-delete.
        with SessionLocal() as db:
            assert (
                db.query(ResearchHistory).filter_by(id=research_id).first()
                is not None
            )

    def test_delete_attempt_stale_in_progress_reclaims(
        self, setup_database_for_all_tests
    ):
        """IN_PROGRESS but thread is dead → reclaimed and deleted (200),
        not 409."""
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_id = _make_session(db, username)
            research_id = str(uuid.uuid4())
            _seed_attempt(
                db,
                username,
                session_id,
                research_id,
                research_status=ResearchStatus.IN_PROGRESS.value,
            )

        # is_research_thread_alive returns False → sweep reclaims.
        with (
            patch(
                "src.local_deep_research.chat.service.is_research_thread_alive",
                return_value=False,
            ),
            _service_db(SessionLocal),
        ):
            assert service.delete_attempt(session_id, research_id) is True

        with SessionLocal() as db:
            assert (
                db.query(ResearchHistory).filter_by(id=research_id).first()
                is None
            )


# ===========================================================================
# Service-level: ChatService.get_original_attempt_query
# ===========================================================================


class TestGetOriginalAttemptQuery:
    """Direct ChatService.get_original_attempt_query tests."""

    def test_uses_research_meta_message_id_fast_path(
        self, setup_database_for_all_tests
    ):
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_id = _make_session(db, username)
            research_id = str(uuid.uuid4())
            _seed_attempt(
                db,
                username,
                session_id,
                research_id,
                user_content="Probe question",
                include_research_meta=True,
            )

        with _service_db(SessionLocal):
            result = service.get_original_attempt_query(session_id, research_id)
        assert result == "Probe question"

    def test_falls_back_to_chatmessage_query_when_meta_missing(
        self, setup_database_for_all_tests
    ):
        """Older rows predate research_meta; the fallback lookup by
        ChatMessage.research_id + role='user' must succeed."""
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_id = _make_session(db, username)
            research_id = str(uuid.uuid4())
            _seed_attempt(
                db,
                username,
                session_id,
                research_id,
                user_content="Legacy query",
                include_research_meta=False,
            )

        with _service_db(SessionLocal):
            result = service.get_original_attempt_query(session_id, research_id)
        assert result == "Legacy query"

    def test_missing_research_raises(self, setup_database_for_all_tests):
        SessionLocal = setup_database_for_all_tests
        username = f"alice_{uuid.uuid4().hex[:8]}"
        service = ChatService(username=username)

        with SessionLocal() as db:
            session_id = _make_session(db, username)

        with _service_db(SessionLocal):
            with pytest.raises(AttemptNotFound):
                service.get_original_attempt_query(
                    session_id, str(uuid.uuid4())
                )


# ===========================================================================
# HTTP-level: DELETE /api/chat/sessions/<sid>/attempts/<rid>
# ===========================================================================


class TestDeleteAttemptRoute:
    """HTTP tests for the delete-attempt endpoint."""

    def test_delete_happy_path(self, authenticated_client):
        """DELETE removes the attempt and returns 200."""
        create_resp = authenticated_client.post(
            "/api/chat/sessions",
            json={"initial_query": "test"},
            content_type="application/json",
        )
        session_id = json.loads(create_resp.data)["session_id"]

        # Drive a real send_message with mocked spawn so we get a real
        # ResearchHistory + ChatMessage + step row to delete.
        with patch("local_deep_research.chat.routes.start_research_process"):
            send_resp = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={"content": "what is X?", "trigger_research": True},
                content_type="application/json",
            )
        assert send_resp.status_code == 200
        research_id = json.loads(send_resp.data)["research_id"]
        assert research_id is not None

        # Sanity: research exists, message_count is at least 1.
        get_resp = authenticated_client.get(f"/api/chat/sessions/{session_id}")
        assert get_resp.status_code == 200
        assert json.loads(get_resp.data)["session"]["message_count"] >= 1

        # Now delete the attempt.
        del_resp = authenticated_client.delete(
            f"/api/chat/sessions/{session_id}/attempts/{research_id}"
        )
        assert del_resp.status_code == 200
        assert json.loads(del_resp.data)["success"] is True

    def test_delete_missing_research_returns_404(self, authenticated_client):
        create_resp = authenticated_client.post(
            "/api/chat/sessions",
            json={"initial_query": "test"},
            content_type="application/json",
        )
        session_id = json.loads(create_resp.data)["session_id"]

        del_resp = authenticated_client.delete(
            f"/api/chat/sessions/{session_id}/attempts/{uuid.uuid4()}"
        )
        assert del_resp.status_code == 404

    def test_delete_missing_session_returns_404(self, authenticated_client):
        """Unknown session_id → 404 (research can't belong to it)."""
        del_resp = authenticated_client.delete(
            f"/api/chat/sessions/{uuid.uuid4()}/attempts/{uuid.uuid4()}"
        )
        assert del_resp.status_code == 404

    def test_delete_in_progress_returns_409(self, authenticated_client):
        """IN_PROGRESS research with a live thread → 409 with
        active_research_id."""
        create_resp = authenticated_client.post(
            "/api/chat/sessions",
            json={"initial_query": "test"},
            content_type="application/json",
        )
        session_id = json.loads(create_resp.data)["session_id"]

        with patch("local_deep_research.chat.routes.start_research_process"):
            send_resp = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={"content": "what is X?", "trigger_research": True},
                content_type="application/json",
            )
        research_id = json.loads(send_resp.data)["research_id"]

        # Patch is_research_thread_alive to True so delete_attempt
        # treats the just-spawned research as live.
        with patch(
            "local_deep_research.chat.service.is_research_thread_alive",
            return_value=True,
        ):
            del_resp = authenticated_client.delete(
                f"/api/chat/sessions/{session_id}/attempts/{research_id}"
            )

        assert del_resp.status_code == 409
        body = json.loads(del_resp.data)
        assert body["success"] is False
        assert body["active_research_id"] == research_id


# ===========================================================================
# HTTP-level: POST /api/chat/sessions/<sid>/attempts/<rid>/retry
# ===========================================================================


class TestRetryAttemptRoute:
    """HTTP tests for the retry-attempt endpoint."""

    def test_retry_happy_path_returns_new_research_id(
        self, authenticated_client
    ):
        """Retry of a FAILED attempt: returns 200 + a fresh research_id
        + message_id."""
        create_resp = authenticated_client.post(
            "/api/chat/sessions",
            json={"initial_query": "test"},
            content_type="application/json",
        )
        session_id = json.loads(create_resp.data)["session_id"]

        with patch("local_deep_research.chat.routes.start_research_process"):
            send_resp = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/messages",
                json={"content": "what is X?", "trigger_research": True},
                content_type="application/json",
            )
        old_research_id = json.loads(send_resp.data)["research_id"]

        # The retry route's stale-row sweep only fires for rows older
        # than the grace window (30s); the just-spawned research is
        # brand-new so the sweep skips it and the per-session guard
        # would refuse with 409. Patch the cap check to short-circuit
        # so we can exercise the actual delete + re-spawn path. The
        # cap-check itself is covered by the dedicated 409 test below.
        # Also patch is_research_thread_alive=False so delete_attempt's
        # own IN_PROGRESS check reclaims the row instead of refusing.
        with (
            patch(
                "local_deep_research.chat.routes._enforce_chat_session_research_slot",
                return_value=None,
            ),
            patch(
                "local_deep_research.chat.service.is_research_thread_alive",
                return_value=False,
            ),
            patch("local_deep_research.chat.routes.start_research_process"),
        ):
            retry_resp = authenticated_client.post(
                f"/api/chat/sessions/{session_id}/attempts/{old_research_id}/retry"
            )

        assert retry_resp.status_code == 200, retry_resp.data
        body = json.loads(retry_resp.data)
        assert body["success"] is True
        new_research_id = body["research_id"]
        new_message_id = body["message_id"]
        assert new_research_id is not None
        assert new_message_id is not None
        # The retry produced a DIFFERENT research id (the old one was
        # deleted; new UUID allocated).
        assert new_research_id != old_research_id

    def test_retry_missing_research_returns_404(self, authenticated_client):
        create_resp = authenticated_client.post(
            "/api/chat/sessions",
            json={"initial_query": "test"},
            content_type="application/json",
        )
        session_id = json.loads(create_resp.data)["session_id"]

        retry_resp = authenticated_client.post(
            f"/api/chat/sessions/{session_id}/attempts/{uuid.uuid4()}/retry"
        )
        assert retry_resp.status_code == 404

    def test_retry_on_missing_session_returns_404(self, authenticated_client):
        retry_resp = authenticated_client.post(
            f"/api/chat/sessions/{uuid.uuid4()}/attempts/{uuid.uuid4()}/retry"
        )
        assert retry_resp.status_code == 404

    def test_retry_archived_session_returns_409(self, authenticated_client):
        """Archived sessions are read-only; retry must refuse."""
        create_resp = authenticated_client.post(
            "/api/chat/sessions",
            json={"initial_query": "test"},
            content_type="application/json",
        )
        session_id = json.loads(create_resp.data)["session_id"]

        # Archive it.
        archived = authenticated_client.patch(
            f"/api/chat/sessions/{session_id}",
            json={"status": "archived"},
            content_type="application/json",
        )
        assert archived.status_code == 200

        retry_resp = authenticated_client.post(
            f"/api/chat/sessions/{session_id}/attempts/{uuid.uuid4()}/retry"
        )
        assert retry_resp.status_code == 409


# ===========================================================================
# CSRF enforcement
# ===========================================================================


class TestAttemptEndpointsCSRF:
    """DELETE and POST retry must reject requests without a CSRF token."""

    def test_delete_without_csrf_token_rejected(
        self, csrf_authenticated_client
    ):
        client, token = csrf_authenticated_client

        create_resp = client.post(
            "/api/chat/sessions",
            json={"initial_query": "Test"},
            content_type="application/json",
            headers={"X-CSRFToken": token},
        )
        assert create_resp.status_code == 200
        session_id = json.loads(create_resp.data)["session_id"]

        # DELETE without token — must be rejected.
        response = client.delete(
            f"/api/chat/sessions/{session_id}/attempts/{uuid.uuid4()}"
        )
        assert response.status_code == 400

    def test_retry_without_csrf_token_rejected(self, csrf_authenticated_client):
        client, token = csrf_authenticated_client

        create_resp = client.post(
            "/api/chat/sessions",
            json={"initial_query": "Test"},
            content_type="application/json",
            headers={"X-CSRFToken": token},
        )
        assert create_resp.status_code == 200
        session_id = json.loads(create_resp.data)["session_id"]

        # POST retry without token — must be rejected.
        response = client.post(
            f"/api/chat/sessions/{session_id}/attempts/{uuid.uuid4()}/retry"
        )
        assert response.status_code == 400