"""Behaviour tests for the ``send_message`` stale-research reclaim path. The send_message reclaim loops in ``chat/routes.py`` run before the per-session and per-user concurrency caps are evaluated: 1. ResearchHistory rows scoped to this chat session whose worker thread is dead AND older than ``_STALE_RESEARCH_GRACE_SECONDS`` (default 30s) are flipped to FAILED, unblocking subsequent sends. 2. UserActiveResearch rows whose worker thread is dead AND older than the same grace cutoff are flipped, restoring per-user cap headroom. The grace window exists specifically so the sweep does not kill the freshly-spawned thread from a sibling concurrent send (which can be in flight between the DB commit and ``start_research_process``). Regressions in this glue silently produce one of: * permanently-409'd chats after any worker crash (sweep too narrow / wrong filter / missing query), * killed in-flight siblings (grace boundary flipped or shrunk), * inflated active-count tally (UserActiveResearch sweep skipped). These tests drive the production helper (``reclaim_stale_user_active_research``) against a real SQLite session and the ResearchHistory sweep contract via the chat archive guard (which shares the same query shape). """ import uuid from datetime import datetime, timedelta, UTC from unittest.mock import patch from src.local_deep_research.constants import ResearchStatus from src.local_deep_research.database.models import ( ChatSession, ChatSessionStatus, ResearchHistory, UserActiveResearch, ) from src.local_deep_research.web.routes.globals import ( is_research_thread_alive, reclaim_stale_user_active_research, ) # --------------------------------------------------------------------------- # Contract tests for is_research_thread_alive — the absent-thread → False # contract is what the reclaim sweep relies on to identify dead rows. # --------------------------------------------------------------------------- def test_is_research_thread_alive_returns_false_for_unknown_id(): """The sweep gates row reclamation on ``is_research_thread_alive``. If that helper grows a false-positive (returns True for an actually- dead thread), the sweep stops reclaiming and chats stay stuck at 409 after a crash. Lock the absent-thread → False contract. """ assert is_research_thread_alive("never-registered-id") is False # --------------------------------------------------------------------------- # UserActiveResearch sweep — driven through the shared # reclaim_stale_user_active_research helper. Both the chat send_message # and research_routes.start_research # now route through this helper; the grace_cutoff_dt argument is the only # behavioural difference between the two call sites. # --------------------------------------------------------------------------- def _seed_user_active_research(db, username, research_id, status, started_at): """Helper to insert a UserActiveResearch row at a specific timestamp.""" row = UserActiveResearch( username=username, research_id=research_id, status=status, started_at=started_at, ) db.add(row) db.commit() return row def test_reclaim_flips_dead_thread_row_to_failed( setup_database_for_all_tests, ): """A stale IN_PROGRESS row whose thread is not alive must be flipped to FAILED so the per-user cap regains headroom.""" SessionLocal = setup_database_for_all_tests username = f"alice_{uuid.uuid4().hex[:8]}" research_id = f"r-{uuid.uuid4().hex[:8]}" old_started_at = datetime.now(UTC) - timedelta(seconds=60) with SessionLocal() as db: _seed_user_active_research( db, username, research_id, ResearchStatus.IN_PROGRESS, old_started_at, ) # `is_research_thread_alive` returns False for any unregistered id # (verified by the contract test above), so the reclaim helper # treats this row as dead. with SessionLocal() as db: reclaimed = reclaim_stale_user_active_research( db, username, grace_cutoff_dt=datetime.now(UTC) ) # Helper does NOT commit (caller composes around it); commit here # so the assertion observes the persisted state. db.commit() assert reclaimed is True with SessionLocal() as db: row = ( db.query(UserActiveResearch) .filter_by(research_id=research_id) .one() ) assert row.status == ResearchStatus.FAILED def test_reclaim_skips_row_inside_grace_window(setup_database_for_all_tests): """A row whose ``started_at`` is at-or-after the grace cutoff must NOT be reclaimed — the grace window exists precisely to avoid killing a sibling request's just-spawned thread before its registration completes.""" SessionLocal = setup_database_for_all_tests username = f"bob_{uuid.uuid4().hex[:8]}" research_id = f"r-{uuid.uuid4().hex[:8]}" fresh_started_at = datetime.now(UTC) with SessionLocal() as db: _seed_user_active_research( db, username, research_id, ResearchStatus.IN_PROGRESS, fresh_started_at, ) # Cutoff is in the past, so the fresh row is filtered out entirely. cutoff = datetime.now(UTC) - timedelta(seconds=30) with SessionLocal() as db: reclaimed = reclaim_stale_user_active_research( db, username, grace_cutoff_dt=cutoff ) db.commit() assert reclaimed is False with SessionLocal() as db: row = ( db.query(UserActiveResearch) .filter_by(research_id=research_id) .one() ) assert row.status == ResearchStatus.IN_PROGRESS def test_reclaim_skips_live_thread(setup_database_for_all_tests): """When ``is_research_thread_alive`` reports a row's thread is live, the row stays IN_PROGRESS even if it is older than the grace cutoff. Mocks the alive check so the test does not depend on a real worker.""" SessionLocal = setup_database_for_all_tests username = f"carol_{uuid.uuid4().hex[:8]}" research_id = f"r-{uuid.uuid4().hex[:8]}" old_started_at = datetime.now(UTC) - timedelta(seconds=120) with SessionLocal() as db: _seed_user_active_research( db, username, research_id, ResearchStatus.IN_PROGRESS, old_started_at, ) with patch( "src.local_deep_research.web.routes.globals.is_research_thread_alive", return_value=True, ): with SessionLocal() as db: reclaimed = reclaim_stale_user_active_research( db, username, grace_cutoff_dt=datetime.now(UTC) ) db.commit() assert reclaimed is False with SessionLocal() as db: row = ( db.query(UserActiveResearch) .filter_by(research_id=research_id) .one() ) assert row.status == ResearchStatus.IN_PROGRESS def test_reclaim_without_grace_cutoff_immediate(setup_database_for_all_tests): """``research_routes.start_research`` does not pass a grace cutoff (it has always reclaimed dead-thread rows immediately). Verify the helper supports that mode.""" SessionLocal = setup_database_for_all_tests username = f"dave_{uuid.uuid4().hex[:8]}" research_id = f"r-{uuid.uuid4().hex[:8]}" # Even a brand-new row is reclaimed when there is no cutoff — # research_routes calls this without grace because it never spawns # a sibling that could race with the sweep. fresh_started_at = datetime.now(UTC) with SessionLocal() as db: _seed_user_active_research( db, username, research_id, ResearchStatus.IN_PROGRESS, fresh_started_at, ) with SessionLocal() as db: reclaimed = reclaim_stale_user_active_research(db, username) db.commit() assert reclaimed is True with SessionLocal() as db: row = ( db.query(UserActiveResearch) .filter_by(research_id=research_id) .one() ) assert row.status == ResearchStatus.FAILED def test_reclaim_scopes_by_username(setup_database_for_all_tests): """The sweep MUST be scoped by username — reclaiming another user's rows would corrupt their state and violate per-user isolation.""" SessionLocal = setup_database_for_all_tests user_a = f"alice_{uuid.uuid4().hex[:8]}" user_b = f"bob_{uuid.uuid4().hex[:8]}" research_b = f"r-{uuid.uuid4().hex[:8]}" old_started_at = datetime.now(UTC) - timedelta(seconds=60) with SessionLocal() as db: _seed_user_active_research( db, user_b, research_b, ResearchStatus.IN_PROGRESS, old_started_at, ) with SessionLocal() as db: reclaimed = reclaim_stale_user_active_research( db, user_a, grace_cutoff_dt=datetime.now(UTC) ) db.commit() assert reclaimed is False # user_a had nothing to reclaim with SessionLocal() as db: row = ( db.query(UserActiveResearch).filter_by(research_id=research_b).one() ) assert row.status == ResearchStatus.IN_PROGRESS # user_b untouched # --------------------------------------------------------------------------- # ResearchHistory sweep — still inline in chat/routes.py::send_message # because the chat-specific chat_session_id scoping cannot be expressed # by the shared helper. Test the underlying query shape end-to-end by # seeding a stale row and verifying the same filter combination the # inline sweep uses returns it. # --------------------------------------------------------------------------- def test_research_history_query_finds_stale_session_row( setup_database_for_all_tests, ): """The ResearchHistory sweep filters on: chat_session_id, status==IN_PROGRESS, created_at < grace_cutoff_iso. Build a row matching those bounds and verify the same query shape returns it — a future filter change that drops one of these predicates fails this test before any reclaim regression ships. """ SessionLocal = setup_database_for_all_tests session_id = str(uuid.uuid4()) research_id = str(uuid.uuid4()) grace_cutoff_iso = (datetime.now(UTC) - timedelta(seconds=30)).isoformat() with SessionLocal() as db: db.add( ChatSession( id=session_id, title="reclaim probe", status=ChatSessionStatus.ACTIVE.value, message_count=0, ) ) db.add( ResearchHistory( id=research_id, query="probe", mode="quick", status=ResearchStatus.IN_PROGRESS, # An ISO timestamp from well before the grace cutoff so # the < grace_cutoff_iso filter matches. created_at="2020-01-01T00:00:00+00:00", chat_session_id=session_id, ) ) db.commit() with SessionLocal() as db: rows = ( db.query(ResearchHistory) .filter( ResearchHistory.chat_session_id == session_id, ResearchHistory.status == ResearchStatus.IN_PROGRESS, ResearchHistory.created_at < grace_cutoff_iso, ) .all() ) assert [r.id for r in rows] == [research_id] def test_research_history_query_skips_fresh_row( setup_database_for_all_tests, ): """The same query MUST skip a row created after the cutoff — verifying the grace window prevents killing siblings.""" SessionLocal = setup_database_for_all_tests session_id = str(uuid.uuid4()) research_id = str(uuid.uuid4()) cutoff_iso = (datetime.now(UTC) - timedelta(seconds=30)).isoformat() fresh_iso = datetime.now(UTC).isoformat() with SessionLocal() as db: db.add( ChatSession( id=session_id, title="reclaim probe", status=ChatSessionStatus.ACTIVE.value, message_count=0, ) ) db.add( ResearchHistory( id=research_id, query="probe", mode="quick", status=ResearchStatus.IN_PROGRESS, created_at=fresh_iso, chat_session_id=session_id, ) ) db.commit() with SessionLocal() as db: rows = ( db.query(ResearchHistory) .filter( ResearchHistory.chat_session_id == session_id, ResearchHistory.status == ResearchStatus.IN_PROGRESS, ResearchHistory.created_at < cutoff_iso, ) .all() ) assert rows == []