fix: fence self-heal worker registration against cancellation

(cherry picked from commit 0fb38b7b5c9453781353464ca5c48c7365359cf1)

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,

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,

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