diff --git a/api/streaming.py b/api/streaming.py index fd772414..b3963101 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -1705,6 +1705,16 @@ def _drop_checkpointed_current_user_from_context(messages, msg_text): return history +def _stream_writeback_is_current(session, stream_id): + """Return True only while a worker still owns the session writeback. + + cancel_stream() intentionally clears ``active_stream_id`` early so the UI can + accept a follow-up turn while the old worker is unwinding. That old worker + must not later persist its stale result over the newer transcript. + """ + return bool(stream_id) and getattr(session, 'active_stream_id', None) == stream_id + + def _merge_display_messages_after_agent_result(previous_display, previous_context, result_messages, msg_text): """Keep UI transcript durable while allowing model context to compact. @@ -3121,6 +3131,14 @@ def _run_agent_streaming( if _ckpt_thread is not None: _ckpt_thread.join(timeout=15) with _agent_lock: + if not ephemeral and not _stream_writeback_is_current(s, stream_id): + logger.info( + "Skipping stale stream writeback for session %s stream %s; active_stream_id=%s", + getattr(s, 'session_id', session_id), + stream_id, + getattr(s, 'active_stream_id', None), + ) + return _result_messages = result.get('messages') or _previous_context_messages _next_context_messages = _restore_reasoning_metadata( _previous_context_messages, @@ -4239,6 +4257,14 @@ def cancel_stream(stream_id: str) -> bool: with _get_session_agent_lock(_cancel_session_id): try: _cs = get_session(_cancel_session_id) + if not _stream_writeback_is_current(_cs, stream_id): + logger.info( + "Skipping stale cancel writeback for session %s stream %s; active_stream_id=%s", + _cancel_session_id, + stream_id, + getattr(_cs, 'active_stream_id', None), + ) + return True # ── Preserve the user's typed message before clearing pending state (#1298) ── # The agent's internal messages list (where the user message was appended at # the start of run_conversation()) may not have been merged back into diff --git a/tests/test_issue1361_cancel_data_loss.py b/tests/test_issue1361_cancel_data_loss.py index 09fc77f0..7fd082c4 100644 --- a/tests/test_issue1361_cancel_data_loss.py +++ b/tests/test_issue1361_cancel_data_loss.py @@ -95,6 +95,7 @@ def _make_session(session_id="cancel_sid_1361", def _setup_cancel_state(session_id, stream_id="stream_1361"): """Wire up STREAMS/CANCEL_FLAGS/AGENT_INSTANCES for cancel_stream().""" + models.SESSIONS[session_id].active_stream_id = stream_id config.STREAMS[stream_id] = queue.Queue() config.CANCEL_FLAGS[stream_id] = threading.Event() mock_agent = Mock() diff --git a/tests/test_issue1857_usage_overwrite.py b/tests/test_issue1857_usage_overwrite.py index 948fd332..affbd72b 100644 --- a/tests/test_issue1857_usage_overwrite.py +++ b/tests/test_issue1857_usage_overwrite.py @@ -110,6 +110,7 @@ def test_stream_completion_overwrites_session_usage_with_latest_turn(cleanup_tes fake_session = FakeSession() fake_stream_id = "stream_issue1857_usage_overwrite" + fake_session.active_stream_id = fake_stream_id fake_queue = queue.Queue() fake_runtime_module = types.ModuleType("hermes_cli.runtime_provider") fake_runtime_module.resolve_runtime_provider = mock.Mock( diff --git a/tests/test_issue1897_profile_switch_agent_cache.py b/tests/test_issue1897_profile_switch_agent_cache.py index 0b82a3a6..97190f66 100644 --- a/tests/test_issue1897_profile_switch_agent_cache.py +++ b/tests/test_issue1897_profile_switch_agent_cache.py @@ -188,6 +188,7 @@ def test_same_session_profile_switch_rebuilds_agent_under_new_soul_home(tmp_path def run_turn(profile_name: str, stream_id: str, text: str): fake_session.profile = profile_name + fake_session.active_stream_id = stream_id streaming.STREAMS[stream_id] = queue.Queue() streaming._run_agent_streaming( session_id=fake_session.session_id, diff --git a/tests/test_issue765_streaming_persistence.py b/tests/test_issue765_streaming_persistence.py index f55bb188..cc693429 100644 --- a/tests/test_issue765_streaming_persistence.py +++ b/tests/test_issue765_streaming_persistence.py @@ -333,7 +333,10 @@ class TestIssue765FollowupHardening: ) stop_idx = src.find("if _checkpoint_stop is not None:\n _checkpoint_stop.set()") join_idx = src.find("if _ckpt_thread is not None:\n _ckpt_thread.join(timeout=15)") - lock_idx = src.find("with _agent_lock:\n _result_messages =") + lock_idx = src.find( + "with _agent_lock:\n" + " if not ephemeral and not _stream_writeback_is_current(s, stream_id):" + ) save_idx = src.find("s.context_messages = _next_context_messages") assert stop_idx != -1, "Success path must stop the checkpoint thread" @@ -353,7 +356,10 @@ class TestIssue765FollowupHardening: src = (Path(__file__).parent.parent / "api" / "streaming.py").read_text( encoding="utf-8" ) - outer_lock_idx = src.find("with _agent_lock:\n _result_messages =") + outer_lock_idx = src.find( + "with _agent_lock:\n" + " if not ephemeral and not _stream_writeback_is_current(s, stream_id):" + ) silent_failure_idx = src.find("if not _assistant_added and not _token_sent:") inner_lock_idx = src.find("with _agent_lock:", outer_lock_idx + 1) compression_idx = src.find("# ── Handle context compression side effects ──") diff --git a/tests/test_sprint42.py b/tests/test_sprint42.py index 79a3314b..995a6614 100644 --- a/tests/test_sprint42.py +++ b/tests/test_sprint42.py @@ -213,6 +213,7 @@ class TestRuntimeRouteInjection(unittest.TestCase): fake_session = FakeSession() fake_stream_id = "stream-runtime-route" + fake_session.active_stream_id = fake_stream_id fake_queue = queue.Queue() fake_runtime_module = types.ModuleType("hermes_cli.runtime_provider") fake_runtime_module.resolve_runtime_provider = resolve_runtime_provider @@ -362,7 +363,10 @@ class TestRuntimeRouteInjection(unittest.TestCase): fake_hermes_state = types.ModuleType("hermes_state") fake_hermes_state.SessionDB = mock.Mock(return_value=object()) - with mock.patch.object(streaming, "get_session", return_value=FakeSession()), \ + fake_session = FakeSession() + fake_session.active_stream_id = fake_stream_id + + with mock.patch.object(streaming, "get_session", return_value=fake_session), \ mock.patch.object(streaming, "_get_ai_agent", return_value=CapturingAgent), \ mock.patch.object(streaming, "resolve_model_provider", return_value=("gpt-4o", "openai-codex", None)), \ mock.patch("api.config.get_config", return_value={}), \ @@ -506,7 +510,10 @@ class TestRuntimeRouteInjection(unittest.TestCase): fake_hermes_state = types.ModuleType("hermes_state") fake_hermes_state.SessionDB = mock.Mock(return_value=object()) - with mock.patch.object(streaming, "get_session", return_value=FakeSession()), \ + fake_session = FakeSession() + fake_session.active_stream_id = fake_stream_id + + with mock.patch.object(streaming, "get_session", return_value=fake_session), \ mock.patch.object(streaming, "_get_ai_agent", return_value=CapturingAgent), \ mock.patch.object(streaming, "resolve_model_provider", return_value=("gpt-5.4", "openai-codex", None)), \ mock.patch.object(streaming, "get_config", return_value={"clarify": {"timeout": 300}}), \ @@ -841,7 +848,10 @@ class TestCredentialPoolBackwardCompat(unittest.TestCase): fake_hermes_state = types.ModuleType("hermes_state") fake_hermes_state.SessionDB = mock.Mock(return_value=None) - with mock.patch.object(streaming, "get_session", return_value=FakeSession()), \ + fake_session = FakeSession() + fake_session.active_stream_id = fake_stream_id + + with mock.patch.object(streaming, "get_session", return_value=fake_session), \ mock.patch.object(streaming, "_get_ai_agent", return_value=OlderAgent), \ mock.patch.object(streaming, "resolve_model_provider", return_value=("gpt-4o", "openai", None)), \ mock.patch("api.config.get_config", return_value={}), \ diff --git a/tests/test_stale_stream_writeback.py b/tests/test_stale_stream_writeback.py new file mode 100644 index 00000000..03aa55da --- /dev/null +++ b/tests/test_stale_stream_writeback.py @@ -0,0 +1,89 @@ +import queue +import threading +from pathlib import Path +from unittest.mock import Mock + +import pytest + +import api.config as config +import api.models as models +import api.streaming as streaming +from api.models import Session + + +@pytest.fixture(autouse=True) +def _isolate_sessions(tmp_path, monkeypatch): + session_dir = tmp_path / "sessions" + session_dir.mkdir() + index_file = session_dir / "_index.json" + monkeypatch.setattr(models, "SESSION_DIR", session_dir) + monkeypatch.setattr(models, "SESSION_INDEX_FILE", index_file) + monkeypatch.setattr(streaming, "SESSION_DIR", session_dir) + monkeypatch.setattr(config, "SESSION_INDEX_FILE", index_file, raising=False) + models.SESSIONS.clear() + config.STREAMS.clear() + config.CANCEL_FLAGS.clear() + config.AGENT_INSTANCES.clear() + config.SESSION_AGENT_LOCKS.clear() + yield + models.SESSIONS.clear() + config.STREAMS.clear() + config.CANCEL_FLAGS.clear() + config.AGENT_INSTANCES.clear() + config.SESSION_AGENT_LOCKS.clear() + + +def test_stream_writeback_requires_active_stream_ownership(): + s = Session(session_id="ownership", messages=[]) + s.active_stream_id = "current-stream" + + assert streaming._stream_writeback_is_current(s, "current-stream") is True + + s.active_stream_id = None + assert streaming._stream_writeback_is_current(s, "current-stream") is False + + s.active_stream_id = "newer-stream" + assert streaming._stream_writeback_is_current(s, "current-stream") is False + + +def test_cancel_stream_does_not_append_marker_after_stream_ownership_rotated(): + sid = "rotated_cancel_sid" + old_stream = "old-stream" + s = Session( + session_id=sid, + title="Rotated stream", + messages=[{"role": "user", "content": "newer prompt"}], + ) + s.active_stream_id = "newer-stream" + s.pending_user_message = "newer prompt" + s.pending_started_at = 456.0 + s.save() + models.SESSIONS[sid] = s + + config.STREAMS[old_stream] = queue.Queue() + config.CANCEL_FLAGS[old_stream] = threading.Event() + mock_agent = Mock() + mock_agent.session_id = sid + mock_agent.interrupt = Mock() + config.AGENT_INSTANCES[old_stream] = mock_agent + + assert streaming.cancel_stream(old_stream) is True + + assert s.active_stream_id == "newer-stream" + assert s.pending_user_message == "newer prompt" + assert [m["content"] for m in s.messages] == ["newer prompt"] + assert all(m.get("content") != "*Task cancelled.*" for m in s.messages) + + +def test_success_path_checks_stream_ownership_before_persisting_result(): + src = Path("api/streaming.py").read_text(encoding="utf-8") + guard = "if not ephemeral and not _stream_writeback_is_current(s, stream_id):" + guard_pos = src.find(guard) + result_merge_pos = src.find("_result_messages = result.get('messages') or _previous_context_messages") + compression_pos = src.find("Handle context compression side effects") + + assert guard_pos != -1 + assert result_merge_pos != -1 + assert compression_pos != -1 + assert guard_pos < result_merge_pos + assert guard_pos < compression_pos