fix: preserve successor ownership of reused Agents

(cherry picked from commit e4037f1d785aa86074e7cc84025942da5c0cc13a)

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`

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.

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):