Merge PR #7748 (contributor head 861a10ec34)

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,245 @@
+"""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", "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()
+    created, publications, finalized = [], [], []
+    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"
+        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()
+        config.register_session_writeback_owner("original", "run")
+        stop_interrupt_seen.clear()
+
+    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", "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 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)
+                        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()
+                    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")
+        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 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
+        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
+    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", {})
+    monkeypatch.setattr(config, "SESSION_WRITEBACK_OWNERS", {})
+    config.register_session_writeback_owner("original", "run")
+
+    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_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

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

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