Keep process-wakeup tags through cancel and recovery paths

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,

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

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

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
+        assert stream_id not in STREAMS