fix: guard stale stream writebacks

This commit is contained in:
Lumen Yang
2026-05-12 13:55:31 +02:00
parent 15d620392f
commit 4b57b202a0
7 changed files with 139 additions and 5 deletions

View File

@@ -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

View File

@@ -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()

View File

@@ -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(

View File

@@ -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,

View File

@@ -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 ──")

View File

@@ -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={}), \

View File

@@ -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