From 9b3951583d5ff113b47b8a27b2e7146e56b7b780 Mon Sep 17 00:00:00 2001 From: Frank Song Date: Tue, 22 Sep 2026 19:17:30 +0800 Subject: [PATCH 1/5] fix: fence self-heal worker registration against cancellation (cherry picked from commit 0fb38b7b5c9453781353464ca5c48c7365359cf1) --- api/streaming.py | 46 +++++---- .../webui-run-state-consistency-contract.md | 13 +++ tests/test_self_heal_worker_cancellation.py | 95 +++++++++++++++++++ 3 files changed, 136 insertions(+), 18 deletions(-) create mode 100644 tests/test_self_heal_worker_cancellation.py diff --git a/api/streaming.py b/api/streaming.py index b61490fb..92235574 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -9791,6 +9791,23 @@ def _run_agent_streaming( _streaming_skill_home_snapshot = None _restore_streaming_skill_home_modules = False _acquired_streaming_skill_home_patch_lock = False + def _register_agent_if_current(candidate): + # Every constructor, including both credential self-heal branches, must + # share the Stop publication edge. CANCEL_FLAGS may already be detached; + # the worker-retained event and stream membership remain authoritative. + with STREAMS_LOCK: + cancelled = cancel_event.is_set() or stream_id not in STREAMS + if not cancelled: + AGENT_INSTANCES[stream_id] = candidate + if not cancelled: + return True + # No Agent call, session persistence or SSE write under registry locks. + try: + candidate.interrupt("Cancelled before start") + except Exception: + logger.debug("Failed to interrupt cancelled candidate agent") + return False + # Initialised here (before any code that may raise) so the outer `finally` # block can safely check `if _checkpoint_stop is not None` even when an # exception fires before the checkpoint thread is created (Issue #765). @@ -11189,20 +11206,7 @@ def _run_agent_streaming( logger.debug('[webui] Evicted LRU agent from cache: %s', _evicted_sid) logger.debug('[webui] Created new agent for session %s', session_id) - # Stop may already have detached CANCEL_FLAGS while initialization - # was in flight. The worker-owned event and stream membership, not - # the removable registry entry, fence registration with cancellation. - with STREAMS_LOCK: - _cancelled_before_start = cancel_event.is_set() or stream_id not in STREAMS - if not _cancelled_before_start: - AGENT_INSTANCES[stream_id] = agent - if _cancelled_before_start: - # Interrupt, persistence and event writes must not hold the - # registry lock (the session lock can be held by another caller). - try: - agent.interrupt("Cancelled before start") - except Exception: - logger.debug("Failed to interrupt agent before start") + if not _register_agent_if_current(agent): with _agent_lock: _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) put('cancel', _cancel_event_payload('Cancelled by user')) @@ -11949,8 +11953,11 @@ def _run_agent_streaming( if 'credential_pool' in _agent_params: _agent_kwargs['credential_pool'] = _runtime_bundle['credential_pool'] agent = _AIAgent(**_agent_kwargs) - with STREAMS_LOCK: - AGENT_INSTANCES[stream_id] = agent + if not _register_agent_if_current(agent): + # Returned-error settlement already owns _agent_lock. + _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) + put('cancel', _cancel_event_payload('Cancelled by user')) + return _agent_sig = _compute_agent_cache_signature( resolved_model, resolved_api_key, @@ -13301,8 +13308,11 @@ def _run_agent_streaming( if 'credential_pool' in _agent_params: _heal_kwargs['credential_pool'] = _runtime_bundle['credential_pool'] _heal_agent = _AIAgent(**_heal_kwargs) - with STREAMS_LOCK: - AGENT_INSTANCES[stream_id] = _heal_agent + if not _register_agent_if_current(_heal_agent): + with _agent_lock: + _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) + put('cancel', _cancel_event_payload('Cancelled by user')) + return _agent_sig = _compute_agent_cache_signature( resolved_model, resolved_api_key, diff --git a/docs/rfcs/webui-run-state-consistency-contract.md b/docs/rfcs/webui-run-state-consistency-contract.md index d1362548..f014747b 100644 --- a/docs/rfcs/webui-run-state-consistency-contract.md +++ b/docs/rfcs/webui-run-state-consistency-contract.md @@ -30,6 +30,19 @@ This RFC defines a consistency contract for those layers. It complements the larger run adapter direction in #1925 by documenting what must remain coherent while WebUI still has multiple overlapping state stores. +## Agent registration after cancellation + +Initial and credential self-heal Agent construction use the same registration +boundary. Under `STREAMS_LOCK`, both the worker-retained cancellation event and +live stream membership must permit registration. A removed `CANCEL_FLAGS` entry +is not permission to restart. If Stop won during self-heal construction, the +replacement must not enter the reusable cache or call `run_conversation`. + +Interrupt and cancellation finalization occur outside the stream registry lock. +Session finalization still owns the session lock: the returned-error path already +holds it, while initial registration and the exception path acquire it. Do not +reacquire this non-reentrant lock from a branch that already owns it. + ## Inactive compression continuation recovery The Agent profile's SQLite compression lineage owns the canonical continuation, diff --git a/tests/test_self_heal_worker_cancellation.py b/tests/test_self_heal_worker_cancellation.py new file mode 100644 index 00000000..3d8f3327 --- /dev/null +++ b/tests/test_self_heal_worker_cancellation.py @@ -0,0 +1,95 @@ +"""Cancellation must fence both credential self-heal constructors. + +Reuse the real worker/Stop fixture; only the external Agent and credential +refresh are synthetic. No provider or tool operation is executed. +""" +import threading +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from api import config, streaming +from tests.test_steer_worker_boundaries import worker_scene as worker_scene + + +@pytest.mark.parametrize("failure_path", ["exception", "returned"]) +@pytest.mark.parametrize("cancellation", ["stop", "event-only", "detached-only", "none"]) +def test_self_heal_registration_obeys_retained_stop(worker_scene, monkeypatch, failure_path, cancellation): + scene = worker_scene + entered, release = threading.Event(), threading.Event() + created, interrupted, finalized = [], [], [] + + def on_init(): + candidate = scene.agent + created.append(candidate) + original_interrupt = candidate.interrupt + + def interrupt(reason): + interrupted.append(candidate) + return original_interrupt(reason) + + candidate.interrupt = interrupt + if len(created) == 2: + entered.set() + assert release.wait(8), "self-heal constructor barrier timed out" + + def on_run(): + if len(created) == 1: + if failure_path == "exception": + raise RuntimeError("401 unauthorized") + scene.result = {"error": "401 unauthorized", "messages": []} + else: + scene.result = None + + scene.on_init, scene.on_run = on_init, on_run + monkeypatch.setattr(streaming, "_attempt_credential_self_heal", lambda *args, **kwargs: { + "provider": "openai", "api_key": "synthetic-self-heal-fixture", + "base_url": "http://127.0.0.1:1/v1", + }) + finalize = streaming._finalize_cancelled_turn + + def observed_finalize(*args, **kwargs): + assert scene.lock.owner != threading.get_ident(), "finalize held STREAMS_LOCK" + finalized.append(True) + return finalize(*args, **kwargs) + + monkeypatch.setattr(streaming, "_finalize_cancelled_turn", observed_finalize) + with ThreadPoolExecutor(max_workers=2) as pool: + stop_future = None + future = pool.submit(scene.run) + try: + assert entered.wait(8), "real worker did not enter credential self-healing" + assert len(created) == 2 + retained = config.CANCEL_FLAGS["run"] + if cancellation == "stop": + # Stop can wait for the session lock after its atomic cancel + # edge. Do not block constructor release on that later write. + stop_future = pool.submit(streaming.cancel_stream, "run") + assert retained.wait(5), "Stop did not publish cancellation" + with config.STREAMS_LOCK: + assert "run" not in config.STREAMS + assert "run" not in config.CANCEL_FLAGS + elif cancellation == "event-only": + with config.STREAMS_LOCK: + retained.set() + config.CANCEL_FLAGS.pop("run") + elif cancellation == "detached-only": + with config.STREAMS_LOCK: + config.STREAMS.pop("run") + assert not retained.is_set() + finally: + release.set() + future.result(timeout=12) + if stop_future is not None: + assert stop_future.result(timeout=12) is True + + if cancellation == "none": + assert scene.calls.count("run") == 2, "valid credential refresh must still retry" + else: + assert scene.calls.count("run") == 1, "cancelled self-heal executed the replacement Agent" + cached = config.SESSION_AGENT_CACHE.get("original") + assert not cached or cached[0] is not created[-1], "cancelled replacement entered the reusable cache" + assert created[-1] in interrupted + assert finalized + assert "run" not in config.AGENT_INSTANCES + assert "run" not in config.ACTIVE_RUNS From 281e91aca36e797bc71458fded0c2e2b263700ec Mon Sep 17 00:00:00 2001 From: Frank Song Date: Tue, 22 Sep 2026 19:28:10 +0800 Subject: [PATCH 2/5] fix: publish newly constructed agents only after Stop admission (cherry picked from commit 84a12254fc36aaff628b7e8e43daf48dfd648e67) --- api/streaming.py | 117 +++++++++--------- .../webui-run-state-consistency-contract.md | 6 +- tests/test_steer_worker_boundaries.py | 9 ++ 3 files changed, 74 insertions(+), 58 deletions(-) diff --git a/api/streaming.py b/api/streaming.py index 92235574..a00337da 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -11030,6 +11030,7 @@ def _run_agent_streaming( # ── Agent cache: reuse across messages in the same session ── # Mirrors gateway _agent_cache. Keeps _user_turn_count alive so # injectionFrequency: "first-turn" actually suppresses after turn 1. + _cache_new_agent = False if ephemeral: agent = _AIAgent(**_agent_kwargs) logger.debug('[webui] Created ephemeral agent for session %s', session_id) @@ -11149,62 +11150,7 @@ def _run_agent_streaming( agent._interrupt_message = None else: agent = _AIAgent(**_agent_kwargs) - # Register the new agent with the memory lifecycle so - # its commit_memory_session() can be found later. - try: - from api.session_lifecycle import register_agent - register_agent(session_id, agent) - except Exception: - logger.debug("Lifecycle register_agent failed for new session %s", session_id, exc_info=True) - _evicted_items = [] - # Snapshot the set of session_ids with a LIVE agent worker - # BEFORE taking SESSION_AGENT_CACHE_LOCK, so LRU eviction never - # closes an agent mid-run AND we never nest ACTIVE_RUNS_LOCK - # inside SESSION_AGENT_CACHE_LOCK (avoids any lock-ordering - # deadlock). A cancel/reconnect can drop STREAMS while the - # worker is still unwinding or blocked in a provider call, so - # ACTIVE_RUNS (worker lifecycle) is the authoritative liveness - # signal, not STREAMS. (#3536 review round 2) - _active_sids = set() - try: - from api.config import ACTIVE_RUNS, ACTIVE_RUNS_LOCK - with ACTIVE_RUNS_LOCK: - for _entry in (ACTIVE_RUNS or {}).values(): - _sid = (_entry or {}).get("session_id") - if _sid: - _active_sids.add(_sid) - except Exception: - _active_sids = set() - with SESSION_AGENT_CACHE_LOCK: - SESSION_AGENT_CACHE[session_id] = (agent, _agent_sig) - SESSION_AGENT_CACHE.move_to_end(session_id) # LRU: mark as recently used - from api.config import SESSION_AGENT_CACHE_MAX - # Evict the oldest INACTIVE entries first. Walk LRU order - # (front = oldest); skip any session with a live run. If - # every over-cap entry is active, leave the cache - # temporarily above cap rather than close a live worker's - # agent — a later insertion/finalization trims it once the - # run ends. - while len(SESSION_AGENT_CACHE) > SESSION_AGENT_CACHE_MAX: - _evictable_sid = None - for _sid in list(SESSION_AGENT_CACHE.keys()): - if _sid not in _active_sids: - _evictable_sid = _sid - break - if _evictable_sid is None: - break # all over-cap entries are active; defer - evicted_entry = SESSION_AGENT_CACHE.pop(_evictable_sid) - _evicted_items.append((_evictable_sid, evicted_entry)) - # Commit and close evicted agents outside the cache lock so - # concurrent cache users are not blocked by provider I/O. - for _evicted_sid, _evicted_entry in _evicted_items: - try: - _evicted_agent = _evicted_entry[0] if isinstance(_evicted_entry, tuple) else None - _close_evicted_agent_at_session_boundary(_evicted_sid, _evicted_agent) - except Exception: - logger.debug("Failed to close evicted agent for session %s", _evicted_sid, exc_info=True) - logger.debug('[webui] Evicted LRU agent from cache: %s', _evicted_sid) - logger.debug('[webui] Created new agent for session %s', session_id) + _cache_new_agent = True if not _register_agent_if_current(agent): with _agent_lock: @@ -11212,6 +11158,65 @@ def _run_agent_streaming( put('cancel', _cancel_event_payload('Cancelled by user')) return + # Publish a newly built reusable agent only after Stop admission. + if _cache_new_agent: + # Register the new agent with the memory lifecycle so + # its commit_memory_session() can be found later. + try: + from api.session_lifecycle import register_agent + register_agent(session_id, agent) + except Exception: + logger.debug("Lifecycle register_agent failed for new session %s", session_id, exc_info=True) + _evicted_items = [] + # Snapshot the set of session_ids with a LIVE agent worker + # BEFORE taking SESSION_AGENT_CACHE_LOCK, so LRU eviction never + # closes an agent mid-run AND we never nest ACTIVE_RUNS_LOCK + # inside SESSION_AGENT_CACHE_LOCK (avoids any lock-ordering + # deadlock). A cancel/reconnect can drop STREAMS while the + # worker is still unwinding or blocked in a provider call, so + # ACTIVE_RUNS (worker lifecycle) is the authoritative liveness + # signal, not STREAMS. (#3536 review round 2) + _active_sids = set() + try: + from api.config import ACTIVE_RUNS, ACTIVE_RUNS_LOCK + with ACTIVE_RUNS_LOCK: + for _entry in (ACTIVE_RUNS or {}).values(): + _sid = (_entry or {}).get("session_id") + if _sid: + _active_sids.add(_sid) + except Exception: + _active_sids = set() + with SESSION_AGENT_CACHE_LOCK: + SESSION_AGENT_CACHE[session_id] = (agent, _agent_sig) + SESSION_AGENT_CACHE.move_to_end(session_id) # LRU: mark as recently used + from api.config import SESSION_AGENT_CACHE_MAX + # Evict the oldest INACTIVE entries first. Walk LRU order + # (front = oldest); skip any session with a live run. If + # every over-cap entry is active, leave the cache + # temporarily above cap rather than close a live worker's + # agent — a later insertion/finalization trims it once the + # run ends. + while len(SESSION_AGENT_CACHE) > SESSION_AGENT_CACHE_MAX: + _evictable_sid = None + for _sid in list(SESSION_AGENT_CACHE.keys()): + if _sid not in _active_sids: + _evictable_sid = _sid + break + if _evictable_sid is None: + break # all over-cap entries are active; defer + evicted_entry = SESSION_AGENT_CACHE.pop(_evictable_sid) + _evicted_items.append((_evictable_sid, evicted_entry)) + # Commit and close evicted agents outside the cache lock so + # concurrent cache users are not blocked by provider I/O. + for _evicted_sid, _evicted_entry in _evicted_items: + try: + _evicted_agent = _evicted_entry[0] if isinstance(_evicted_entry, tuple) else None + _close_evicted_agent_at_session_boundary(_evicted_sid, _evicted_agent) + except Exception: + logger.debug("Failed to close evicted agent for session %s", _evicted_sid, exc_info=True) + logger.debug('[webui] Evicted LRU agent from cache: %s', _evicted_sid) + logger.debug('[webui] Created new agent for session %s', session_id) + # Prepend workspace context so the agent always knows which directory # to use for file operations, regardless of session age or AGENTS.md defaults. workspace_ctx = _workspace_context_prefix(str(s.workspace)) diff --git a/docs/rfcs/webui-run-state-consistency-contract.md b/docs/rfcs/webui-run-state-consistency-contract.md index f014747b..82355ac4 100644 --- a/docs/rfcs/webui-run-state-consistency-contract.md +++ b/docs/rfcs/webui-run-state-consistency-contract.md @@ -35,8 +35,10 @@ while WebUI still has multiple overlapping state stores. Initial and credential self-heal Agent construction use the same registration boundary. Under `STREAMS_LOCK`, both the worker-retained cancellation event and live stream membership must permit registration. A removed `CANCEL_FLAGS` entry -is not permission to restart. If Stop won during self-heal construction, the -replacement must not enter the reusable cache or call `run_conversation`. +is not permission to restart. If Stop won during initial or self-heal +construction, the candidate must not enter the reusable cache or call +`run_conversation`. New-agent lifecycle/cache publication follows successful +registration; it must not precede the cancellation check. Interrupt and cancellation finalization occur outside the stream registry lock. Session finalization still owns the session lock: the returned-error path already diff --git a/tests/test_steer_worker_boundaries.py b/tests/test_steer_worker_boundaries.py index d55d3691..a97d450b 100644 --- a/tests/test_steer_worker_boundaries.py +++ b/tests/test_steer_worker_boundaries.py @@ -211,6 +211,8 @@ def test_stop_during_agent_creation_prevents_provider_run(worker_scene, monkeypa release.set() worker.result(timeout=10) assert "run" not in scene.calls, "accepted Stop must prevent the provider/tool run" + cached = config.SESSION_AGENT_CACHE.get("original") + assert not cached or cached[0] is not scene.agent, "cancelled initial candidate remained reusable" assert finalized, "cancelled turn was not finalized" assert "run" not in config.AGENT_INSTANCES assert "run" not in config.ACTIVE_RUNS @@ -276,3 +278,10 @@ def test_final_drain_fences_steer(worker_scene, monkeypatch, registered, rotated assert scene.agent.pending == [] assert finalizing_edges == [True], "final drain admission must close under the stream lock" assert "run" not in config.ACTIVE_RUNS + + +def test_successfully_registered_new_agent_remains_cached(worker_scene): + scene = worker_scene + scene.run() + assert scene.calls.count("run") == 1 + assert config.SESSION_AGENT_CACHE["original"][0] is scene.agent From acdcddf5c9ffc1c13d6052a52d2a719e8abfb715 Mon Sep 17 00:00:00 2001 From: Frank Song Date: Tue, 22 Sep 2026 20:03:32 +0800 Subject: [PATCH 3/5] fix: make Agent cache publication atomic with Stop (cherry picked from commit c40c43e0eee76a829b113d0119a476a651b60125) --- api/streaming.py | 108 +++++---- .../webui-run-state-consistency-contract.md | 14 +- tests/test_agent_max_turns_parity.py | 17 +- tests/test_agent_publication_cancellation.py | 206 ++++++++++++++++++ ...st_issue1897_profile_switch_agent_cache.py | 19 +- tests/test_streaming_max_tokens_quota.py | 17 +- 6 files changed, 313 insertions(+), 68 deletions(-) create mode 100644 tests/test_agent_publication_cancellation.py diff --git a/api/streaming.py b/api/streaming.py index a00337da..08913993 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -9791,7 +9791,7 @@ def _run_agent_streaming( _streaming_skill_home_snapshot = None _restore_streaming_skill_home_modules = False _acquired_streaming_skill_home_patch_lock = False - def _register_agent_if_current(candidate): + def _register_agent_if_current(candidate, cache_signature=None): # Every constructor, including both credential self-heal branches, must # share the Stop publication edge. CANCEL_FLAGS may already be detached; # the worker-retained event and stream membership remain authoritative. @@ -9799,6 +9799,20 @@ def _run_agent_streaming( cancelled = cancel_event.is_set() or stream_id not in STREAMS if not cancelled: AGENT_INSTANCES[stream_id] = candidate + # Check and publication share one admission, not a check followed + # by an unlocked cache/lifecycle write that could undo Stop. + if cache_signature is not None: + from api.config import SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK + with SESSION_AGENT_CACHE_LOCK: + SESSION_AGENT_CACHE[session_id] = (candidate, cache_signature) + SESSION_AGENT_CACHE.move_to_end(session_id) + if not ephemeral: + try: + # Pure in-memory bookkeeping; no Agent/provider call. + from api.session_lifecycle import register_agent + register_agent(session_id, candidate) + except Exception: + logger.debug("Lifecycle register_agent failed for session %s", session_id, exc_info=True) if not cancelled: return True # No Agent call, session persistence or SSE write under registry locks. @@ -9808,6 +9822,33 @@ def _run_agent_streaming( logger.debug("Failed to interrupt cancelled candidate agent") return False + def _agent_can_invoke(candidate): + # Prompt preparation and LRU cleanup can yield after registration. Admit + # each invocation again at its point of use; never hold registry locks + # across run_conversation or interrupt. A later Stop interrupts the + # already-admitted invocation through the registered Agent as usual. + from api.config import SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK + with STREAMS_LOCK: + current = (not cancel_event.is_set() and stream_id in STREAMS + and AGENT_INSTANCES.get(stream_id) is candidate) + if not current: + if AGENT_INSTANCES.get(stream_id) is candidate: + AGENT_INSTANCES.pop(stream_id, None) + with SESSION_AGENT_CACHE_LOCK: + entry = SESSION_AGENT_CACHE.get(session_id) + if entry and entry[0] is candidate: + SESSION_AGENT_CACHE.pop(session_id, None) + # Dirty memory segments retain their own owner; clear + # only this still-current reusable future-turn handle. + from api.session_lifecycle import unregister_agent + unregister_agent(session_id) + if not current: + try: + candidate.interrupt("Cancelled before invocation") + except Exception: + logger.debug("Failed to interrupt unstarted candidate agent") + return current + # Initialised here (before any code that may raise) so the outer `finally` # block can safely check `if _checkpoint_stop is not None` even when an # exception fires before the checkpoint thread is created (Issue #765). @@ -11073,14 +11114,6 @@ def _run_agent_streaming( session_id, _cached_agent_session_identity(_cached_agent), ) - if agent is not None: - # Reopened/cache-hit sessions must register the agent - # so later lifecycle commits can find it. - try: - from api.session_lifecycle import register_agent - register_agent(session_id, agent) - except Exception: - logger.debug("Lifecycle register_agent failed for cached session %s", session_id, exc_info=True) if _identity_mismatch_entry is not None: try: @@ -11152,21 +11185,15 @@ def _run_agent_streaming( agent = _AIAgent(**_agent_kwargs) _cache_new_agent = True - if not _register_agent_if_current(agent): + if not _register_agent_if_current(agent, _agent_sig if _cache_new_agent else None): with _agent_lock: _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) put('cancel', _cancel_event_payload('Cancelled by user')) return - # Publish a newly built reusable agent only after Stop admission. + # Cache publication is already fenced with Stop above. Keep LRU + # eviction and any memory-provider close outside the stream lock. if _cache_new_agent: - # Register the new agent with the memory lifecycle so - # its commit_memory_session() can be found later. - try: - from api.session_lifecycle import register_agent - register_agent(session_id, agent) - except Exception: - logger.debug("Lifecycle register_agent failed for new session %s", session_id, exc_info=True) _evicted_items = [] # Snapshot the set of session_ids with a LIVE agent worker # BEFORE taking SESSION_AGENT_CACHE_LOCK, so LRU eviction never @@ -11187,8 +11214,6 @@ def _run_agent_streaming( except Exception: _active_sids = set() with SESSION_AGENT_CACHE_LOCK: - SESSION_AGENT_CACHE[session_id] = (agent, _agent_sig) - SESSION_AGENT_CACHE.move_to_end(session_id) # LRU: mark as recently used from api.config import SESSION_AGENT_CACHE_MAX # Evict the oldest INACTIVE entries first. Walk LRU order # (front = oldest); skip any session with a live run. If @@ -11476,6 +11501,11 @@ def _run_agent_streaming( ) _run_conversation_kwargs["user_message"] = user_message _result_partial_pre_call_context = list(_previous_context_messages) + if not _agent_can_invoke(agent): + with _agent_lock: + _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) + put('cancel', _cancel_event_payload('Cancelled by user')) + return result = agent.run_conversation(**_run_conversation_kwargs) _remember_pending_steer_result(result) _active_turn_identity = _resolve_active_turn_authority( @@ -11958,11 +11988,6 @@ def _run_agent_streaming( if 'credential_pool' in _agent_params: _agent_kwargs['credential_pool'] = _runtime_bundle['credential_pool'] agent = _AIAgent(**_agent_kwargs) - if not _register_agent_if_current(agent): - # Returned-error settlement already owns _agent_lock. - _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) - put('cancel', _cancel_event_payload('Cancelled by user')) - return _agent_sig = _compute_agent_cache_signature( resolved_model, resolved_api_key, @@ -11979,10 +12004,11 @@ def _run_agent_streaming( profile_home=_profile_home, safe_profile_runtime_env=_safe_profile_runtime_env, ) - from api.config import SESSION_AGENT_CACHE as _SAC, SESSION_AGENT_CACHE_LOCK as _SAC_L - with _SAC_L: - _SAC[session_id] = (agent, _agent_sig) - _SAC.move_to_end(session_id) + if not _register_agent_if_current(agent, _agent_sig): + # Returned-error settlement already owns _agent_lock. + _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) + put('cancel', _cancel_event_payload('Cancelled by user')) + return # Retry the conversation once with fresh credentials _self_healed = True _token_sent = False @@ -12015,6 +12041,11 @@ def _run_agent_streaming( _result_partial_pre_call_context = list( _heal_context_messages ) + if not _agent_can_invoke(agent): + # Returned-error settlement already owns the lock. + _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) + put('cancel', _cancel_event_payload('Cancelled by user')) + return _heal_result = agent.run_conversation(**_heal_kwargs) _remember_pending_steer_result(_heal_result) _active_turn_identity = _resolve_active_turn_authority( @@ -13313,11 +13344,6 @@ def _run_agent_streaming( if 'credential_pool' in _agent_params: _heal_kwargs['credential_pool'] = _runtime_bundle['credential_pool'] _heal_agent = _AIAgent(**_heal_kwargs) - if not _register_agent_if_current(_heal_agent): - with _agent_lock: - _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) - put('cancel', _cancel_event_payload('Cancelled by user')) - return _agent_sig = _compute_agent_cache_signature( resolved_model, resolved_api_key, @@ -13334,10 +13360,11 @@ def _run_agent_streaming( profile_home=_profile_home, safe_profile_runtime_env=_safe_profile_runtime_env, ) - from api.config import SESSION_AGENT_CACHE as _SAC2, SESSION_AGENT_CACHE_LOCK as _SAC2_L - with _SAC2_L: - _SAC2[session_id] = (_heal_agent, _agent_sig) - _SAC2.move_to_end(session_id) + if not _register_agent_if_current(_heal_agent, _agent_sig): + with _agent_lock: + _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) + put('cancel', _cancel_event_payload('Cancelled by user')) + return # Retry the conversation _token_sent = False try: @@ -13369,6 +13396,11 @@ def _run_agent_streaming( _result_partial_pre_call_context = list( _heal_context_messages ) + if not _agent_can_invoke(_heal_agent): + with _agent_lock: + _finalize_cancelled_turn(s, ephemeral=ephemeral, message='Task cancelled before start.', stream_id=stream_id) + put('cancel', _cancel_event_payload('Cancelled by user')) + return _heal_result = _heal_agent.run_conversation(**_heal_kwargs2) _remember_pending_steer_result(_heal_result) _active_turn_identity = _resolve_active_turn_authority( diff --git a/docs/rfcs/webui-run-state-consistency-contract.md b/docs/rfcs/webui-run-state-consistency-contract.md index 82355ac4..e3da5bac 100644 --- a/docs/rfcs/webui-run-state-consistency-contract.md +++ b/docs/rfcs/webui-run-state-consistency-contract.md @@ -37,8 +37,18 @@ boundary. Under `STREAMS_LOCK`, both the worker-retained cancellation event and live stream membership must permit registration. A removed `CANCEL_FLAGS` entry is not permission to restart. If Stop won during initial or self-heal construction, the candidate must not enter the reusable cache or call -`run_conversation`. New-agent lifecycle/cache publication follows successful -registration; it must not precede the cancellation check. +`run_conversation`. Stream registration, reusable-cache publication and the +in-memory lifecycle handle share one atomic Stop admission, using lock order +`STREAMS_LOCK` then `SESSION_AGENT_CACHE_LOCK`. Merely moving a cache write +after a cancellation check does not protect that check-to-publication gap. + +After prompt preparation and immediately before each initial/self-heal invocation, +revalidate the retained cancel event, stream membership and exact registered +Agent. If Stop already won, remove only that Agent's matching reusable entry +and settle cancellation. Never clear a successor's cache or lifecycle handle. +Stop after invocation admission uses the existing Agent interrupt mechanism; +registry locks must not span provider or tool execution. LRU eviction/close stays +outside the stream lock and retains the existing active-worker policy. Interrupt and cancellation finalization occur outside the stream registry lock. Session finalization still owns the session lock: the returned-error path already diff --git a/tests/test_agent_max_turns_parity.py b/tests/test_agent_max_turns_parity.py index 4665f74f..c83d8777 100644 --- a/tests/test_agent_max_turns_parity.py +++ b/tests/test_agent_max_turns_parity.py @@ -79,20 +79,19 @@ RAISED_EXCEPTION_SELF_HEAL_MARKER = ( "retrying stream after credential refresh')" ) -# Each region ends at the cache write that stores the signature it just -# computed. Bounding the region at the write -- rather than at the next -# lifecycle anchor or at EOF -- is what proves the signature is recomputed -# *before* it is cached: a call moved below its own write lands outside the -# region and fails, instead of passing while the write consumes a stale -# `_agent_sig`. +# Each region ends at the atomic registration/publication call that consumes +# the signature it just computed. The shared helper now writes the cache under +# the Stop lock; the call boundary must still follow fresh signature computation. +# A signature call moved below its publication lands outside the region and +# fails, rather than silently reusing a stale `_agent_sig` on a retry. INITIAL_SEND_CACHE_WRITE_MARKER = ( - "SESSION_AGENT_CACHE[session_id] = (agent, _agent_sig)" + "if not _register_agent_if_current(agent, _agent_sig if _cache_new_agent else None):" ) RETURNED_ERROR_SELF_HEAL_CACHE_WRITE_MARKER = ( - "_SAC[session_id] = (agent, _agent_sig)" + "if not _register_agent_if_current(agent, _agent_sig):" ) RAISED_EXCEPTION_SELF_HEAL_CACHE_WRITE_MARKER = ( - "_SAC2[session_id] = (_heal_agent, _agent_sig)" + "if not _register_agent_if_current(_heal_agent, _agent_sig):" ) diff --git a/tests/test_agent_publication_cancellation.py b/tests/test_agent_publication_cancellation.py new file mode 100644 index 00000000..c08acab8 --- /dev/null +++ b/tests/test_agent_publication_cancellation.py @@ -0,0 +1,206 @@ +"""Stop must not lose to cache publication after an Agent was registered. + +Run the real worker and Stop primitive. Pause immediately after the registration +lock is released, not inside the constructor covered by the older tests. Only +external Agent/provider calls are synthetic; no live credentials are used. +""" +import threading +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from api import config, session_lifecycle, streaming +from tests.test_steer_worker_boundaries import worker_scene as worker_scene + + +@pytest.mark.parametrize("path", ["initial", "cached", "returned", "exception"]) +@pytest.mark.parametrize("cancellation", ["stop", "event-only", "detached-only", "successor", "none"]) +def test_stop_after_registration_before_cache_or_invocation(worker_scene, monkeypatch, path, cancellation): + scene = worker_scene + reached, release = threading.Event(), threading.Event() + created, publications, finalized = [], [], [] + worker_id = [None] + paused = [False] + successor = object() + monkeypatch.setattr(session_lifecycle, "_sessions", {}) + + def on_init(): + created.append(scene.agent) + + def on_run(): + assert scene.lock.owner != threading.get_ident(), "provider invocation under stream lock" + if path in ("returned", "exception") and len(created) == 1: + if path == "exception": + raise RuntimeError("401 unauthorized") + scene.result = {"error": "401 unauthorized", "messages": []} + else: + scene.result = None + + scene.on_init, scene.on_run = on_init, on_run + if path == "cached": + scene.run() + assert scene.calls.count("run") == 1 + assert config.SESSION_AGENT_CACHE["original"][0] is created[0] + scene.calls.clear() + config.STREAMS["run"] = scene.events + config.STREAM_SESSION_OWNERS["run"] = "original" + scene.session.active_stream_id = "run" + scene.session.pending_user_message = "Do the task." + scene.session.pending_started_at = 2.0 + scene.session.save() + + class ObservedCache(OrderedDict): + def __setitem__(self, key, value): + if key == "original": + publications.append((value[0], scene.lock.owner == threading.get_ident(), + "run" in config.STREAMS)) + super().__setitem__(key, value) + + cache = ObservedCache() + # Preserve an existing cache hit without counting fixture setup as publication. + for key, value in config.SESSION_AGENT_CACHE.items(): + OrderedDict.__setitem__(cache, key, value) + monkeypatch.setattr(config, "SESSION_AGENT_CACHE", cache) + if hasattr(streaming, "SESSION_AGENT_CACHE"): + monkeypatch.setattr(streaming, "SESSION_AGENT_CACHE", cache) + monkeypatch.setattr(streaming, "_attempt_credential_self_heal", lambda *a, **kw: { + "provider": "openai", "api_key": "synthetic-publication-fixture", + "base_url": "http://127.0.0.1:1/v1", + }) + + target_count = 2 if path in ("returned", "exception") else 1 + + def after_registration_unlock(): + if (worker_id[0] != threading.get_ident() or paused[0] + or len(created) != target_count + or config.AGENT_INSTANCES.get("run") is not created[-1]): + return + paused[0] = True + reached.set() + assert release.wait(8), "post-registration barrier timed out" + + scene.lock.after_release = after_registration_unlock + finalize = streaming._finalize_cancelled_turn + + def observed_finalize(*args, **kwargs): + assert scene.lock.owner != threading.get_ident(), "finalize under stream lock" + finalized.append(True) + return finalize(*args, **kwargs) + + monkeypatch.setattr(streaming, "_finalize_cancelled_turn", observed_finalize) + + def run(): + worker_id[0] = threading.get_ident() + scene.run() + + with ThreadPoolExecutor(max_workers=2) as pool: + stop_future = None + future = pool.submit(run) + try: + assert reached.wait(8), "worker never released target registration lock" + retained = config.CANCEL_FLAGS["run"] + if cancellation in ("stop", "successor"): + stop_future = pool.submit(streaming.cancel_stream, "run") + assert retained.wait(5), "Stop did not publish cancellation" + with config.STREAMS_LOCK: + assert "run" not in config.STREAMS + if cancellation == "successor": + with config.SESSION_AGENT_CACHE_LOCK: + OrderedDict.__setitem__(cache, "original", (successor, "new-signature")) + session_lifecycle.register_agent("original", successor) + config.AGENT_INSTANCES["successor"] = successor + elif cancellation == "event-only": + with config.STREAMS_LOCK: + retained.set() + config.CANCEL_FLAGS.pop("run") + elif cancellation == "detached-only": + with config.STREAMS_LOCK: + config.STREAMS.pop("run") + assert not retained.is_set() + finally: + release.set() + future.result(timeout=12) + if stop_future is not None: + assert stop_future.result(timeout=12) is True + + prior_runs = 1 if path in ("returned", "exception") else 0 + if cancellation == "none": + assert scene.calls.count("run") == prior_runs + 1 + assert config.SESSION_AGENT_CACHE["original"][0] is created[-1] + else: + assert scene.calls.count("run") == prior_runs, "cancelled candidate was invoked" + cached = config.SESSION_AGENT_CACHE.get("original") + assert not cached or cached[0] is not created[-1], "cancelled candidate stayed reusable" + assert finalized, "cancellation did not reach terminal settlement" + if cancellation == "successor": + assert config.SESSION_AGENT_CACHE["original"][0] is successor + assert session_lifecycle._sessions["original"]["agent"] is successor + assert config.AGENT_INSTANCES["successor"] is successor + assert all(held and live for _, held, live in publications), ( + "Agent cache publication escaped the stream/Stop admission boundary", publications) + assert "run" not in config.AGENT_INSTANCES + assert "run" not in config.ACTIVE_RUNS + + +def test_cache_publication_and_registration_exclude_competing_stop(worker_scene, monkeypatch): + """Block the actual cache write; Stop cannot observe a half-published Agent.""" + scene = worker_scene + entered, release = threading.Event(), threading.Event() + interrupted_state = [] + interrupt_observed = threading.Event() + monkeypatch.setattr(session_lifecycle, "_sessions", {}) + + class PublicationBarrier(OrderedDict): + def __setitem__(self, key, value): + if key == "original": + entered.set() + assert release.wait(8), "cache publication was not released" + super().__setitem__(key, value) + + cache = PublicationBarrier() + monkeypatch.setattr(config, "SESSION_AGENT_CACHE", cache) + if hasattr(streaming, "SESSION_AGENT_CACHE"): + monkeypatch.setattr(streaming, "SESSION_AGENT_CACHE", cache) + + def on_init(): + candidate = scene.agent + interrupt = candidate.interrupt + + def observed_interrupt(reason): + interrupted_state.append(( + cache.get("original", (None,))[0] is candidate, + session_lifecycle._sessions.get("original", {}).get("agent") is candidate, + )) + result = interrupt(reason) + interrupt_observed.set() + return result + + candidate.interrupt = observed_interrupt + + scene.on_init = on_init + with ThreadPoolExecutor(max_workers=2) as pool: + worker = pool.submit(scene.run) + stop = None + try: + assert entered.wait(8), "worker never reached cache publication" + retained = config.CANCEL_FLAGS["run"] + # Wait at the first publication unlock until Stop has won admission, + # so the provider cannot outrun the controlled schedule. + worker_ident = scene.lock.owner + def wait_for_stop(): + if threading.get_ident() == worker_ident: + scene.lock.after_release = None + assert interrupt_observed.wait(5), "Stop did not interrupt the published Agent" + scene.lock.after_release = wait_for_stop + stop = pool.submit(streaming.cancel_stream, "run") + assert scene.lock.contender.wait(5), "Stop bypassed publication's stream lock" + assert not retained.is_set(), "Stop observed half-published Agent state" + finally: + release.set() + worker.result(timeout=12) + if stop is not None: + assert stop.result(timeout=12) is True + assert interrupted_state[0] == (True, True), "Stop did not see coherent cache/lifecycle publication" + assert "run" not in scene.calls + assert "original" not in cache diff --git a/tests/test_issue1897_profile_switch_agent_cache.py b/tests/test_issue1897_profile_switch_agent_cache.py index 684c5a01..2b1b0801 100644 --- a/tests/test_issue1897_profile_switch_agent_cache.py +++ b/tests/test_issue1897_profile_switch_agent_cache.py @@ -79,20 +79,19 @@ RAISED_EXCEPTION_SELF_HEAL_MARKER = ( "retrying stream after credential refresh')" ) -# Each region ends at the cache write that stores the signature it just -# computed. Bounding the region at the write -- rather than at the next -# lifecycle anchor or at EOF -- is what proves the signature is recomputed -# *before* it is cached: a call moved below its own write lands outside the -# region and fails, instead of passing while the write consumes a stale -# `_agent_sig`. +# Each region ends at the atomic registration/publication call that consumes +# the signature it just computed. The shared helper now writes the cache under +# the Stop lock; the call boundary must still follow fresh signature computation. +# A signature call moved below its publication lands outside the region and +# fails, rather than silently reusing a stale `_agent_sig` on a retry. INITIAL_SEND_CACHE_WRITE_MARKER = ( - "SESSION_AGENT_CACHE[session_id] = (agent, _agent_sig)" + "if not _register_agent_if_current(agent, _agent_sig if _cache_new_agent else None):" ) RETURNED_ERROR_SELF_HEAL_CACHE_WRITE_MARKER = ( - "_SAC[session_id] = (agent, _agent_sig)" + "if not _register_agent_if_current(agent, _agent_sig):" ) RAISED_EXCEPTION_SELF_HEAL_CACHE_WRITE_MARKER = ( - "_SAC2[session_id] = (_heal_agent, _agent_sig)" + "if not _register_agent_if_current(_heal_agent, _agent_sig):" ) @@ -112,7 +111,7 @@ def _signature_calls_by_region( The regions are the initial send, the returned-error self-heal retry and the raised-exception self-heal retry; each runs from its lifecycle anchor - to the cache write that consumes `_agent_sig`. Requiring exactly one call + to the atomic publication call that consumes `_agent_sig`. Requiring one call inside those bounds means neither a dropped retry call (masked by its surviving siblings) nor a call recomputed after its own cache write (which would leave the write storing a stale signature) can read as correct. diff --git a/tests/test_streaming_max_tokens_quota.py b/tests/test_streaming_max_tokens_quota.py index ed352290..e008c252 100644 --- a/tests/test_streaming_max_tokens_quota.py +++ b/tests/test_streaming_max_tokens_quota.py @@ -84,20 +84,19 @@ RAISED_EXCEPTION_SELF_HEAL_MARKER = ( "retrying stream after credential refresh')" ) -# Each region ends at the cache write that stores the signature it just -# computed. Bounding the region at the write -- rather than at the next -# lifecycle anchor or at EOF -- is what proves the signature is recomputed -# *before* it is cached: a call moved below its own write lands outside the -# region and fails, instead of passing while the write consumes a stale -# `_agent_sig`. +# Each region ends at the atomic registration/publication call that consumes +# the signature it just computed. The shared helper now writes the cache under +# the Stop lock; the call boundary must still follow fresh signature computation. +# A signature call moved below its publication lands outside the region and +# fails, rather than silently reusing a stale `_agent_sig` on a retry. INITIAL_SEND_CACHE_WRITE_MARKER = ( - "SESSION_AGENT_CACHE[session_id] = (agent, _agent_sig)" + "if not _register_agent_if_current(agent, _agent_sig if _cache_new_agent else None):" ) RETURNED_ERROR_SELF_HEAL_CACHE_WRITE_MARKER = ( - "_SAC[session_id] = (agent, _agent_sig)" + "if not _register_agent_if_current(agent, _agent_sig):" ) RAISED_EXCEPTION_SELF_HEAL_CACHE_WRITE_MARKER = ( - "_SAC2[session_id] = (_heal_agent, _agent_sig)" + "if not _register_agent_if_current(_heal_agent, _agent_sig):" ) From 861a10ec34ad5a6a99603818b15548c8b1f10ac9 Mon Sep 17 00:00:00 2001 From: Frank Song Date: Tue, 22 Sep 2026 20:22:00 +0800 Subject: [PATCH 4/5] fix: preserve successor ownership of reused Agents (cherry picked from commit e4037f1d785aa86074e7cc84025942da5c0cc13a) --- api/streaming.py | 48 ++++++++++++----- .../webui-run-state-consistency-contract.md | 13 ++++- tests/test_agent_publication_cancellation.py | 53 ++++++++++++++++--- 3 files changed, 91 insertions(+), 23 deletions(-) diff --git a/api/streaming.py b/api/streaming.py index 08913993..4dfa5b63 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -9792,6 +9792,7 @@ def _run_agent_streaming( _restore_streaming_skill_home_modules = False _acquired_streaming_skill_home_patch_lock = False def _register_agent_if_current(candidate, cache_signature=None): + nonlocal agent # Every constructor, including both credential self-heal branches, must # share the Stop publication edge. CANCEL_FLAGS may already be detached; # the worker-retained event and stream membership remain authoritative. @@ -9815,6 +9816,15 @@ def _run_agent_streaming( logger.debug("Lifecycle register_agent failed for session %s", session_id, exc_info=True) if not cancelled: return True + # A cache hit is a borrowed reusable object, not this worker's private + # candidate. Stop may already have let a successor reuse it. Rejected + # cache-hit admission must not interrupt that other turn. + if cache_signature is None and not ephemeral: + if agent is candidate: + agent = None + return False + # Newly constructed candidates were never published on this failed + # admission; they cannot have been borrowed by a successor from us. # No Agent call, session persistence or SSE write under registry locks. try: candidate.interrupt("Cancelled before start") @@ -9823,30 +9833,40 @@ def _run_agent_streaming( return False def _agent_can_invoke(candidate): + nonlocal agent # Prompt preparation and LRU cleanup can yield after registration. Admit # each invocation again at its point of use; never hold registry locks # across run_conversation or interrupt. A later Stop interrupts the # already-admitted invocation through the registered Agent as usual. - from api.config import SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK + from api.config import ( + SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK, + SESSION_WRITEBACK_OWNERS, SESSION_WRITEBACK_OWNERS_LOCK, + ) with STREAMS_LOCK: current = (not cancel_event.is_set() and stream_id in STREAMS and AGENT_INSTANCES.get(stream_id) is candidate) if not current: if AGENT_INSTANCES.get(stream_id) is candidate: AGENT_INSTANCES.pop(stream_id, None) - with SESSION_AGENT_CACHE_LOCK: - entry = SESSION_AGENT_CACHE.get(session_id) - if entry and entry[0] is candidate: - SESSION_AGENT_CACHE.pop(session_id, None) - # Dirty memory segments retain their own owner; clear - # only this still-current reusable future-turn handle. - from api.session_lifecycle import unregister_agent - unregister_agent(session_id) - if not current: - try: - candidate.interrupt("Cancelled before invocation") - except Exception: - logger.debug("Failed to interrupt unstarted candidate agent") + # Agent identity is not turn identity: a successor can reuse the + # same cached object. Use the existing cancellation-surviving + # ownership record, holding its lock through cache retirement. + # Missing ownership also fails closed (successor may have ended). + with SESSION_WRITEBACK_OWNERS_LOCK: + if SESSION_WRITEBACK_OWNERS.get(session_id) == stream_id: + with SESSION_AGENT_CACHE_LOCK: + entry = SESSION_AGENT_CACHE.get(session_id) + if entry and entry[0] is candidate: + SESSION_AGENT_CACHE.pop(session_id, None) + # Dirty memory segments keep their original owner. + from api.session_lifecycle import unregister_agent + unregister_agent(session_id) + # This invocation never started. Stop owns any required interrupt of the + # running Agent; a late duplicate interrupt could hit its new borrower. + # Drop our local borrowed handle too: final Steer drain must not reach + # the successor through the old worker's fallback `agent` reference. + if not current and agent is candidate: + agent = None return current # Initialised here (before any code that may raise) so the outer `finally` diff --git a/docs/rfcs/webui-run-state-consistency-contract.md b/docs/rfcs/webui-run-state-consistency-contract.md index e3da5bac..4d1c4e19 100644 --- a/docs/rfcs/webui-run-state-consistency-contract.md +++ b/docs/rfcs/webui-run-state-consistency-contract.md @@ -44,8 +44,17 @@ after a cancellation check does not protect that check-to-publication gap. After prompt preparation and immediately before each initial/self-heal invocation, revalidate the retained cancel event, stream membership and exact registered -Agent. If Stop already won, remove only that Agent's matching reusable entry -and settle cancellation. Never clear a successor's cache or lifecycle handle. +Agent. If Stop already won, retire a matching reusable entry only while the +existing `SESSION_WRITEBACK_OWNERS` record still equals this exact stream. +Hold that ownership lock through cache/lifecycle retirement. Object identity is +insufficient because successors can reuse the same Agent; absent ownership is +not permission either, since a completed successor clears its record. Never +clear a successor's cache or lifecycle handle, and do not issue another Agent +interrupt for an invocation that never started. Rejected cache-hit registration +likewise must not interrupt the borrowed Agent; only a never-published newly +constructed candidate may receive construction-cancellation cleanup. Drop the +old worker's local borrowed handle as well, so final pending-Steer drain cannot +reach a successor through an object-reference fallback. Stop after invocation admission uses the existing Agent interrupt mechanism; registry locks must not span provider or tool execution. LRU eviction/close stays outside the stream lock and retains the existing active-worker policy. diff --git a/tests/test_agent_publication_cancellation.py b/tests/test_agent_publication_cancellation.py index c08acab8..5f854def 100644 --- a/tests/test_agent_publication_cancellation.py +++ b/tests/test_agent_publication_cancellation.py @@ -15,7 +15,7 @@ from tests.test_steer_worker_boundaries import worker_scene as worker_scene @pytest.mark.parametrize("path", ["initial", "cached", "returned", "exception"]) -@pytest.mark.parametrize("cancellation", ["stop", "event-only", "detached-only", "successor", "none"]) +@pytest.mark.parametrize("cancellation", ["stop", "event-only", "detached-only", "successor", "same-agent-successor", "same-agent-completed", "none"]) def test_stop_after_registration_before_cache_or_invocation(worker_scene, monkeypatch, path, cancellation): scene = worker_scene reached, release = threading.Event(), threading.Event() @@ -23,10 +23,20 @@ def test_stop_after_registration_before_cache_or_invocation(worker_scene, monkey worker_id = [None] paused = [False] successor = object() + stop_interrupt_seen = threading.Event() + successor_interrupt_count = None monkeypatch.setattr(session_lifecycle, "_sessions", {}) + monkeypatch.setattr(config, "SESSION_WRITEBACK_OWNERS", {}) + config.register_session_writeback_owner("original", "run") def on_init(): created.append(scene.agent) + interrupt = scene.agent.interrupt + def observed_interrupt(reason): + result = interrupt(reason) + stop_interrupt_seen.set() + return result + scene.agent.interrupt = observed_interrupt def on_run(): assert scene.lock.owner != threading.get_ident(), "provider invocation under stream lock" @@ -49,6 +59,8 @@ def test_stop_after_registration_before_cache_or_invocation(worker_scene, monkey scene.session.pending_user_message = "Do the task." scene.session.pending_started_at = 2.0 scene.session.save() + config.register_session_writeback_owner("original", "run") + stop_interrupt_seen.clear() class ObservedCache(OrderedDict): def __setitem__(self, key, value): @@ -100,16 +112,32 @@ def test_stop_after_registration_before_cache_or_invocation(worker_scene, monkey try: assert reached.wait(8), "worker never released target registration lock" retained = config.CANCEL_FLAGS["run"] - if cancellation in ("stop", "successor"): + if cancellation in ("stop", "successor", "same-agent-successor", "same-agent-completed"): stop_future = pool.submit(streaming.cancel_stream, "run") assert retained.wait(5), "Stop did not publish cancellation" + assert stop_interrupt_seen.wait(5), "Stop has not finished its own interrupt" with config.STREAMS_LOCK: assert "run" not in config.STREAMS - if cancellation == "successor": + if cancellation in ("successor", "same-agent-successor", "same-agent-completed"): + if cancellation.startswith("same-agent"): + successor = created[-1] + config.register_session_writeback_owner("original", "successor") with config.SESSION_AGENT_CACHE_LOCK: OrderedDict.__setitem__(cache, "original", (successor, "new-signature")) session_lifecycle.register_agent("original", successor) - config.AGENT_INSTANCES["successor"] = successor + scene.session.active_stream_id = "successor" + scene.session.pending_user_message = "Successor request" + if cancellation == "same-agent-completed": + config.clear_session_writeback_owner_if_owned("original", "successor") + scene.session.active_stream_id = None + scene.session.pending_user_message = None + else: + config.AGENT_INSTANCES["successor"] = successor + config.register_active_run("successor", session_id="original", backend="legacy") + if cancellation == "same-agent-successor": + with successor.pending_lock: + successor.pending.append("Successor-only guidance") + successor_interrupt_count = scene.calls.count("interrupt") elif cancellation == "event-only": with config.STREAMS_LOCK: retained.set() @@ -131,12 +159,21 @@ def test_stop_after_registration_before_cache_or_invocation(worker_scene, monkey else: assert scene.calls.count("run") == prior_runs, "cancelled candidate was invoked" cached = config.SESSION_AGENT_CACHE.get("original") - assert not cached or cached[0] is not created[-1], "cancelled candidate stayed reusable" + if not cancellation.startswith("same-agent"): + assert not cached or cached[0] is not created[-1], "cancelled candidate stayed reusable" assert finalized, "cancellation did not reach terminal settlement" - if cancellation == "successor": + if successor_interrupt_count is not None: + assert scene.calls.count("interrupt") == successor_interrupt_count, "stale worker interrupted its successor" assert config.SESSION_AGENT_CACHE["original"][0] is successor assert session_lifecycle._sessions["original"]["agent"] is successor - assert config.AGENT_INSTANCES["successor"] is successor + if cancellation == "same-agent-successor": + assert successor.pending == ["Successor-only guidance"], "stale terminal drain stole successor guidance" + if cancellation == "same-agent-completed": + assert scene.session.active_stream_id is None + else: + assert config.AGENT_INSTANCES["successor"] is successor + assert scene.session.active_stream_id == "successor" + assert scene.session.pending_user_message == "Successor request" assert all(held and live for _, held, live in publications), ( "Agent cache publication escaped the stream/Stop admission boundary", publications) assert "run" not in config.AGENT_INSTANCES @@ -150,6 +187,8 @@ def test_cache_publication_and_registration_exclude_competing_stop(worker_scene, interrupted_state = [] interrupt_observed = threading.Event() monkeypatch.setattr(session_lifecycle, "_sessions", {}) + monkeypatch.setattr(config, "SESSION_WRITEBACK_OWNERS", {}) + config.register_session_writeback_owner("original", "run") class PublicationBarrier(OrderedDict): def __setitem__(self, key, value): From 4437f5fb408d32beb2e272248ed933cc7bf418ba Mon Sep 17 00:00:00 2001 From: nesquena-hermes Date: Wed, 23 Sep 2026 06:23:42 +0000 Subject: [PATCH 5/5] docs(changelog): #7748 --- CHANGELOG.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5602de1a..ecbc5bad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,13 @@ ### Fixed +- **A stopped chat can no longer publish or reuse its agent after Stop.** A worker that was + cancelled could still publish its cached Agent or invoke it after Stop landed, and its late + cleanup could close or evict the same Agent object a successor turn had just picked up. The initial + path and both credential self-heal paths now take one Stop admission, the cancellation event, live + stream and exact Agent are rechecked immediately before invocation (no registry lock is held across + provider or tool execution), and cache/lifecycle handles are retired only while the surviving owner + still matches. (#7748 by @franksong2702) - **Reconnecting to a running session no longer redraws every tool card over and over.** After a reconnect restored the live activity scene from the run journal, the client also replayed its older cached in-flight tool list on top, so a turn with N tool cards redrew them N times. When the