fix: make Agent cache publication atomic with Stop

(cherry picked from commit c40c43e0eee76a829b113d0119a476a651b60125)

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(

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

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

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

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.

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