From 59c868fe8615a10583816e96231fbd8595dc22c1 Mon Sep 17 00:00:00 2001 From: Frank Song Date: Thu, 9 Jul 2026 09:06:19 +0800 Subject: [PATCH] fix(wakeup): preserve pause on gateway stop race --- api/gateway_chat.py | 20 ++++- tests/test_issue3929_process_wakeup_pause.py | 94 ++++++++++++++++++++ 2 files changed, 110 insertions(+), 4 deletions(-) diff --git a/api/gateway_chat.py b/api/gateway_chat.py index 96d14ecb..d05fb5c4 100644 --- a/api/gateway_chat.py +++ b/api/gateway_chat.py @@ -850,6 +850,7 @@ def _run_gateway_chat_streaming( assistant_msg["reasoning"] = saved_reasoning previous_messages = list(getattr(s, "messages", None) or []) previous_context = list(getattr(s, "context_messages", None) or getattr(s, "messages", None) or []) + previous_process_wakeup_pause = dict(getattr(s, "process_wakeup_pause", {}) or {}) # Stamp stable ids on the two new rows (shared with the display merge # below) so display and model-context copies share an id for the # fork/truncate aligner (#context-message-stable-id). @@ -908,16 +909,27 @@ def _run_gateway_chat_streaming( s.workspace = str(workspace) s.model = model s.model_provider = model_provider + + def _restore_cancelled_success_writeback(): + s.context_messages = previous_context + s.messages = previous_messages + s.process_wakeup_pause = dict(previous_process_wakeup_pause) + s.save() + put_gateway_event("cancel", {"message": "Cancelled by user"}) + # Recheck immediately before clearing the pause; Stop can arrive # while the success transcript is being assembled. if cancel_event.is_set(): - s.context_messages = previous_context - s.messages = previous_messages - s.save() - put_gateway_event("cancel", {"message": "Cancelled by user"}) + _restore_cancelled_success_writeback() return clear_process_wakeup_pause(s, reason="run_completed") + if cancel_event.is_set(): + _restore_cancelled_success_writeback() + return s.save() + if cancel_event.is_set(): + _restore_cancelled_success_writeback() + return try: from api.goals import evaluate_goal_after_turn, has_active_goal from api.profiles import get_hermes_home_for_profile diff --git a/tests/test_issue3929_process_wakeup_pause.py b/tests/test_issue3929_process_wakeup_pause.py index c07ba5b0..421c34cd 100644 --- a/tests/test_issue3929_process_wakeup_pause.py +++ b/tests/test_issue3929_process_wakeup_pause.py @@ -12,6 +12,7 @@ from unittest import mock import pytest import api.config as config +import api.gateway_chat as gateway_chat import api.models as models import api.profiles as profiles import api.providers as providers @@ -1159,6 +1160,99 @@ def test_process_wakeup_pause_empty_provider_lane_probes_after_fingerprint_chang assert saved.process_wakeup_pause == {} +def test_gateway_cancel_during_completion_save_restores_process_wakeup_pause(tmp_path, monkeypatch): + stream_id = "gateway-pause-save-race-stream" + session_id = "gateway_pause_save_race" + stream_queue = queue.Queue() + config.STREAMS[stream_id] = stream_queue + monkeypatch.setattr(gateway_chat, "RunJournalWriter", lambda *_args, **_kwargs: None) + monkeypatch.setattr(gateway_chat, "gateway_approval_unavailable_reason", lambda *_args, **_kwargs: None) + monkeypatch.setattr(config, "get_config", lambda: {"webui_gateway_base_url": "http://gateway.test"}) + + class _GatewayResponse: + def __enter__(self): + return self + + def __exit__(self, *_exc): + return False + + def __iter__(self): + payload = {"choices": [{"delta": {"content": "Gateway reply"}}]} + return iter([ + ("data: " + json.dumps(payload) + "\n").encode("utf-8"), + b"data: [DONE]\n", + ]) + + monkeypatch.setattr( + gateway_chat.urllib.request, + "urlopen", + lambda *_args, **_kwargs: _GatewayResponse(), + ) + + previous_pause = { + "paused": True, + "model": "claude-sonnet-test", + "provider": "test-provider", + "classification": "credential_pool_empty", + "first_paused_at": 1.0, + "last_visible_error_at": 1.0, + "visible_error_count": 1, + "suppressed_count": 0, + "credential_state_fingerprint": "fingerprint-before", + } + previous_messages = [{"role": "user", "content": "before", "timestamp": 1.0}] + session = Session( + session_id=session_id, + workspace=str(tmp_path), + model="claude-sonnet-test", + model_provider="test-provider", + messages=list(previous_messages), + context_messages=list(previous_messages), + active_stream_id=stream_id, + pending_user_message="wake up", + pending_user_source="process_wakeup", + process_wakeup_pause=dict(previous_pause), + ) + session.save() + models.SESSIONS[session_id] = session + + original_save = Session.save + save_calls = {"count": 0, "cleared_pause": 0} + + def _save_and_cancel_after_success_clear(self, *args, **kwargs): + save_calls["count"] += 1 + result = original_save(self, *args, **kwargs) + if ( + getattr(self, "session_id", None) == session_id + and not getattr(self, "process_wakeup_pause", None) + and save_calls["cleared_pause"] == 0 + ): + save_calls["cleared_pause"] += 1 + config.CANCEL_FLAGS[stream_id].set() + return result + + monkeypatch.setattr(Session, "save", _save_and_cancel_after_success_clear) + + gateway_chat._run_gateway_chat_streaming( + session_id, + "wake up", + "claude-sonnet-test", + str(tmp_path), + stream_id, + model_provider="test-provider", + ) + + assert save_calls["cleared_pause"] == 1 + saved = Session.load(session_id) + assert saved is not None + assert saved.process_wakeup_pause == previous_pause + assert saved.messages == previous_messages + assert saved.context_messages == previous_messages + queued_events = [item[0] for item in list(stream_queue.queue)] + assert "cancel" in queued_events + assert "done" not in queued_events + + def test_stale_credential_empty_process_wakeup_still_records_pause(tmp_path): session = Session( session_id="wakeup_pause_stale",