""" Unit tests for cancel/interrupt functionality. Tests the integration between cancel_stream() and agent.interrupt(). """ import queue import threading from unittest.mock import Mock from api.streaming import cancel_stream from api.config import AGENT_INSTANCES, STREAMS, CANCEL_FLAGS, ACTIVE_RUNS, SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK class TestCancelInterrupt: """Test suite for cancel/interrupt functionality""" def setup_method(self): """Clean up before each test""" AGENT_INSTANCES.clear() STREAMS.clear() CANCEL_FLAGS.clear() ACTIVE_RUNS.clear() with SESSION_AGENT_CACHE_LOCK: SESSION_AGENT_CACHE.clear() def teardown_method(self): """Clean up after each test""" AGENT_INSTANCES.clear() STREAMS.clear() CANCEL_FLAGS.clear() ACTIVE_RUNS.clear() with SESSION_AGENT_CACHE_LOCK: SESSION_AGENT_CACHE.clear() def test_cancel_calls_agent_interrupt(self): """Verify that cancel_stream() calls agent.interrupt() when agent exists""" # Setup stream_id = "test_stream_123" mock_agent = Mock() mock_agent.interrupt = Mock() STREAMS[stream_id] = queue.Queue() CANCEL_FLAGS[stream_id] = threading.Event() AGENT_INSTANCES[stream_id] = mock_agent # Execute result = cancel_stream(stream_id) # Assert assert result is True mock_agent.interrupt.assert_called_once_with("Cancelled by user") # CANCEL_FLAGS is eagerly popped after cancel (#776 fix) so the flag # is no longer in the dict — verify the pop happened instead assert stream_id not in CANCEL_FLAGS, \ "cancel_stream() should eagerly pop CANCEL_FLAGS after signalling" def test_cancel_handles_interrupt_exception(self): """Verify that cancel_stream() handles interrupt() exceptions gracefully""" stream_id = "test_stream_456" mock_agent = Mock() mock_agent.interrupt = Mock(side_effect=RuntimeError("Agent error")) STREAMS[stream_id] = queue.Queue() CANCEL_FLAGS[stream_id] = threading.Event() AGENT_INSTANCES[stream_id] = mock_agent # Should not raise exception result = cancel_stream(stream_id) # Assert assert result is True mock_agent.interrupt.assert_called_once() assert stream_id not in CANCEL_FLAGS, \ "cancel_stream() should eagerly pop CANCEL_FLAGS even on interrupt exception" def test_cancel_before_agent_ready(self): """Test cancel when agent not yet stored in AGENT_INSTANCES (race condition)""" stream_id = "test_stream_789" STREAMS[stream_id] = queue.Queue() CANCEL_FLAGS[stream_id] = threading.Event() # Note: AGENT_INSTANCES[stream_id] not set (simulating race condition) # Should succeed even without agent result = cancel_stream(stream_id) # Assert assert result is True # CANCEL_FLAGS is eagerly popped; the agent thread checks the event # object it already has a reference to — pop doesn't clear the event assert stream_id not in CANCEL_FLAGS, \ "cancel_stream() should eagerly pop CANCEL_FLAGS even without an agent" # Agent will check this flag (it holds a reference to the event object) def test_cancel_nonexistent_stream(self): """Test cancel for a stream that doesn't exist""" result = cancel_stream("nonexistent_stream") assert result is False def test_cancel_falls_back_to_active_run_registry(self): """Cancel should still work when STREAMS is gone but the worker is alive.""" from unittest.mock import patch stream_id = "detached_stream_123" session_id = "sess_detached_123" mock_agent = Mock() mock_agent.interrupt = Mock() mock_agent.session_id = session_id mock_session = Mock() mock_session.session_id = session_id mock_session.active_stream_id = stream_id mock_session.pending_user_message = "hello" mock_session.pending_attachments = ["file.txt"] mock_session.pending_started_at = 1234567890.0 mock_session.messages = [] mock_session.save = Mock() ACTIVE_RUNS[stream_id] = { "session_id": session_id, "started_at": 1234567890.0, "phase": "running", } with SESSION_AGENT_CACHE_LOCK: SESSION_AGENT_CACHE[session_id] = (mock_agent, "sig") with patch("api.streaming.get_session", return_value=mock_session): result = cancel_stream(stream_id) assert result is True assert ACTIVE_RUNS[stream_id]["phase"] == "cancelling" mock_agent.interrupt.assert_called_once_with("Cancelled by user") assert mock_session.active_stream_id is None assert mock_session.pending_user_message is None assert mock_session.pending_attachments == [] assert mock_session.pending_started_at is None mock_session.save.assert_called_once() def test_cancel_sets_cancel_event(self): """Verify that cancel_stream() sets the cancel_event flag""" stream_id = "test_stream_event" STREAMS[stream_id] = queue.Queue() cancel_event = threading.Event() CANCEL_FLAGS[stream_id] = cancel_event result = cancel_stream(stream_id) assert result is True assert cancel_event.is_set() def test_cancel_puts_sentinel_in_queue(self): """Verify that cancel_stream() puts cancel sentinel in queue""" stream_id = "test_stream_queue" q = queue.Queue() STREAMS[stream_id] = q CANCEL_FLAGS[stream_id] = threading.Event() result = cancel_stream(stream_id) assert result is True # Check that cancel message was queued assert not q.empty() event_type, data = q.get_nowait() assert event_type == 'cancel' assert data['message'] == 'Cancelled by user' def test_cancel_preserves_partial_text_when_interrupt_pops_buffers(self): """Regression (Codex pre-release finding on #3475): the partial-text / reasoning / tool-call buffers must be snapshotted UNDER streams_lock BEFORE agent.interrupt() runs. Otherwise the worker's finally block can pop those live buffers (it does so under STREAMS_LOCK) the instant the interrupt wakes it, and the cancelled turn silently loses its already-streamed partial text. We simulate that race deterministically: the mock agent's interrupt() clears the live STREAM_* maps, mimicking the worker finally. The fix must still persist the partial text captured before the interrupt. """ from unittest.mock import patch from api.config import ( STREAM_PARTIAL_TEXT, STREAM_REASONING_TEXT, STREAM_LIVE_TOOL_CALLS, ) stream_id = "race_stream_partial" session_id = "sess_race_partial" q = queue.Queue() STREAMS[stream_id] = q CANCEL_FLAGS[stream_id] = threading.Event() STREAM_PARTIAL_TEXT[stream_id] = "partial answer so far" STREAM_REASONING_TEXT[stream_id] = "thinking..." STREAM_LIVE_TOOL_CALLS[stream_id] = [{"name": "search"}] mock_agent = Mock() mock_agent.session_id = session_id def _interrupt(_msg): # Mimic the worker's finally block popping the live buffers the # moment the interrupt wakes it (it runs under STREAMS_LOCK in prod; # here we just clear to model the worst-case post-interrupt state). STREAM_PARTIAL_TEXT.pop(stream_id, None) STREAM_REASONING_TEXT.pop(stream_id, None) STREAM_LIVE_TOOL_CALLS.pop(stream_id, None) mock_agent.interrupt = Mock(side_effect=_interrupt) AGENT_INSTANCES[stream_id] = mock_agent mock_session = Mock() mock_session.session_id = session_id mock_session.active_stream_id = stream_id mock_session.pending_user_message = "q" mock_session.pending_attachments = [] mock_session.pending_started_at = 1.0 mock_session.messages = [] mock_session.save = Mock() with patch("api.streaming.get_session", return_value=mock_session): result = cancel_stream(stream_id) assert result is True mock_agent.interrupt.assert_called_once_with("Cancelled by user") # The cancelled turn must carry the partial text that was live BEFORE the # interrupt popped the buffers. Find it in the appended messages. appended = [m for m in mock_session.messages if isinstance(m, dict)] joined = " ".join(str(m.get("content", "")) for m in appended) assert "partial answer so far" in joined, ( "cancelled turn lost its already-streamed partial text — the snapshot " "must be captured under streams_lock BEFORE agent.interrupt()" ) def test_cancel_preserves_partial_text_on_detached_active_run_path(self): """Regression (Codex 2nd pre-release finding on #3475): the under-lock snapshot must also cover the STREAMS-absent / ACTIVE_RUNS-present path. When the browser SSE has detached (no STREAMS entry) but the worker is still live in ACTIVE_RUNS, cancel resolves the agent from SESSION_AGENT_CACHE; the partial-text snapshot must still be taken before agent.interrupt() pops the live buffers. """ from unittest.mock import patch from api.config import ( STREAM_PARTIAL_TEXT, STREAM_REASONING_TEXT, STREAM_LIVE_TOOL_CALLS, ) stream_id = "detached_race_stream" session_id = "sess_detached_race" # NOTE: deliberately NO STREAMS entry — this is the detached path. STREAM_PARTIAL_TEXT[stream_id] = "detached partial text" STREAM_REASONING_TEXT[stream_id] = "detached reasoning" STREAM_LIVE_TOOL_CALLS[stream_id] = [{"name": "tool"}] mock_agent = Mock() mock_agent.session_id = session_id def _interrupt(_msg): STREAM_PARTIAL_TEXT.pop(stream_id, None) STREAM_REASONING_TEXT.pop(stream_id, None) STREAM_LIVE_TOOL_CALLS.pop(stream_id, None) mock_agent.interrupt = Mock(side_effect=_interrupt) ACTIVE_RUNS[stream_id] = { "session_id": session_id, "started_at": 1.0, "phase": "running", } with SESSION_AGENT_CACHE_LOCK: SESSION_AGENT_CACHE[session_id] = (mock_agent, "sig") mock_session = Mock() mock_session.session_id = session_id mock_session.active_stream_id = stream_id mock_session.pending_user_message = "q" mock_session.pending_attachments = [] mock_session.pending_started_at = 1.0 mock_session.messages = [] mock_session.save = Mock() with patch("api.streaming.get_session", return_value=mock_session), \ patch("api.streaming._cached_agent_matches_session", return_value=True): result = cancel_stream(stream_id) assert result is True mock_agent.interrupt.assert_called_once_with("Cancelled by user") appended = [m for m in mock_session.messages if isinstance(m, dict)] joined = " ".join(str(m.get("content", "")) for m in appended) assert "detached partial text" in joined, ( "detached-path cancel lost its partial text — the under-lock snapshot " "must cover the ACTIVE_RUNS-only path too, not just STREAMS-present" )