diff --git a/api/models.py b/api/models.py index a6994908..02fa8d21 100644 --- a/api/models.py +++ b/api/models.py @@ -1896,6 +1896,7 @@ def _apply_core_sync_or_error_marker( session.pending_user_message = None session.pending_attachments = [] session.pending_started_at = None + session.pending_user_source = None session.save(touch_updated_at=touch_updated_at) logger.info( "Session %s: cleared stale pending state for completed stream %s without error marker", @@ -1922,6 +1923,7 @@ def _apply_core_sync_or_error_marker( session.pending_user_message = None session.pending_attachments = [] session.pending_started_at = None + session.pending_user_source = None session.messages.append( _build_recovery_marker_with_retry_hook( recovered_output=recovered_output, @@ -1976,6 +1978,7 @@ def _apply_core_sync_or_error_marker( session.pending_user_message = None session.pending_attachments = [] session.pending_started_at = None + session.pending_user_source = None if recovered_output: session.messages.append( _interrupted_recovery_marker( @@ -2023,6 +2026,7 @@ def _apply_core_sync_or_error_marker( session.pending_user_message = None session.pending_attachments = [] session.pending_started_at = None + session.pending_user_source = None session.messages.append( _build_recovery_marker_with_retry_hook( recovered_output=recovered_output, diff --git a/api/streaming.py b/api/streaming.py index f0dac4f9..02ad77ce 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -5008,6 +5008,9 @@ def _materialize_pending_user_turn_before_error(session) -> bool: 'timestamp': recovered_ts, '_recovered': True, } + pending_source = getattr(session, 'pending_user_source', None) + if pending_source and pending_source != 'webui': + recovered['_source'] = pending_source pending_attachments = getattr(session, 'pending_attachments', None) if pending_attachments: recovered['attachments'] = list(pending_attachments) @@ -8960,6 +8963,7 @@ def cancel_stream(stream_id: str) -> bool: # the cleanup. try: _pending_user = getattr(_cs, 'pending_user_message', None) + _pending_source = getattr(_cs, 'pending_user_source', None) _pending_atts_raw = getattr(_cs, 'pending_attachments', None) _pending_atts = list(_pending_atts_raw) if isinstance(_pending_atts_raw, (list, tuple)) else [] _pending_started = getattr(_cs, 'pending_started_at', None) or 0 @@ -8985,11 +8989,16 @@ def cancel_stream(stream_id: str) -> bool: if _pending_user == _last_content or _pending_user in _last_content: _already_persisted = True if not _already_persisted: + _recovered_ts = int(time.time()) + if isinstance(_pending_started, (int, float)) and _pending_started > 0: + _recovered_ts = int(_pending_started) _user_turn: dict = { 'role': 'user', 'content': _pending_user, - 'timestamp': int(time.time()), + 'timestamp': _recovered_ts, } + if _pending_source and _pending_source != 'webui': + _user_turn['_source'] = _pending_source if _pending_atts: _user_turn['attachments'] = _pending_atts _msgs_for_recovery.append(_user_turn) @@ -9002,6 +9011,7 @@ def cancel_stream(stream_id: str) -> bool: _cs.pending_user_message = None _cs.pending_attachments = [] _cs.pending_started_at = None + _cs.pending_user_source = None # Persist any partial assistant text that was streamed before cancel (#893). # Preserving partial content means the user sees what the agent had # produced rather than losing it entirely. The marker is _partial=True diff --git a/tests/test_process_wakeup_synthetic.py b/tests/test_process_wakeup_synthetic.py index e1ae731a..c5a0c12c 100644 --- a/tests/test_process_wakeup_synthetic.py +++ b/tests/test_process_wakeup_synthetic.py @@ -1,9 +1,11 @@ """Test coverage for source field threading on process-wakeup synthetic turns.""" import time -from api.models import Session, _append_recovered_pending_turn +from pathlib import Path + +from api.models import Session, _append_recovered_pending_turn, _apply_core_sync_or_error_marker from api.routes import _checkpoint_user_message_for_eager_session_save -from api.streaming import _merge_display_messages_after_agent_result +from api.streaming import _materialize_pending_user_turn_before_error, _merge_display_messages_after_agent_result def test_append_recovered_pending_turn_stamps_process_wakeup_source(): @@ -162,3 +164,29 @@ def test_merge_display_leaves_webui_user_turn_unmarked(): assert merged[0]["role"] == "user" assert "_source" not in merged[0] + + +def test_materialize_pending_user_turn_before_error_stamps_process_wakeup_source(): + s = Session( + session_id="test-session-8", + pending_user_message="[IMPORTANT: Wakeup prompt]", + pending_user_source="process_wakeup", + ) + s.messages = [] + + assert _materialize_pending_user_turn_before_error(s) is True + assert s.messages[0]["_source"] == "process_wakeup" + + +def test_apply_core_sync_or_error_marker_clears_pending_user_source(tmp_path): + s = Session( + session_id="test-session-9", + pending_user_message="[IMPORTANT: Wakeup prompt]", + pending_user_source="process_wakeup", + ) + s.messages = [{"role": "assistant", "content": "done"}] + s.pending_started_at = time.time() + s.save = lambda *args, **kwargs: None + + assert _apply_core_sync_or_error_marker(s, Path(tmp_path / "missing-core.json")) is True + assert s.pending_user_source is None diff --git a/tests/test_sprint51.py b/tests/test_sprint51.py index e1e6074d..377efe9b 100644 --- a/tests/test_sprint51.py +++ b/tests/test_sprint51.py @@ -86,6 +86,8 @@ class TestCancelStreamEagerRelease: mock_session.pending_user_message = "hello" mock_session.pending_attachments = ["file.txt"] mock_session.pending_started_at = 1234567890.0 + mock_session.pending_user_source = "process_wakeup" + mock_session.messages = [] STREAMS[stream_id] = queue.Queue() CANCEL_FLAGS[stream_id] = threading.Event() @@ -102,6 +104,14 @@ class TestCancelStreamEagerRelease: "cancel_stream() should clear session.pending_attachments" assert mock_session.pending_started_at is None, \ "cancel_stream() should clear session.pending_started_at" + assert mock_session.pending_user_source is None, \ + "cancel_stream() should clear session.pending_user_source" + assert any( + isinstance(msg, dict) + and msg.get("role") == "user" + and msg.get("_source") == "process_wakeup" + for msg in mock_session.messages + ), "cancel_stream() should preserve process_wakeup source on recovered user turns" mock_session.save.assert_called_once() def test_cancel_without_agent_still_pops_streams(self): @@ -168,4 +178,4 @@ class TestCancelStreamEagerRelease: result = cancel_stream(stream_id) assert result is True - assert stream_id not in STREAMS \ No newline at end of file + assert stream_id not in STREAMS