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