fix: guard stale stream writebacks
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 ──")
|
||||
|
||||
@@ -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={}), \
|
||||
|
||||
89
tests/test_stale_stream_writeback.py
Normal file
89
tests/test_stale_stream_writeback.py
Normal 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
|
||||
Reference in New Issue
Block a user