diff --git a/api/streaming.py b/api/streaming.py index 08913993..4dfa5b63 100644 --- a/api/streaming.py +++ b/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` diff --git a/docs/rfcs/webui-run-state-consistency-contract.md b/docs/rfcs/webui-run-state-consistency-contract.md index e3da5bac..4d1c4e19 100644 --- a/docs/rfcs/webui-run-state-consistency-contract.md +++ b/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. diff --git a/tests/test_agent_publication_cancellation.py b/tests/test_agent_publication_cancellation.py index c08acab8..5f854def 100644 --- a/tests/test_agent_publication_cancellation.py +++ b/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):