diff --git a/api/streaming.py b/api/streaming.py index a00337da..08913993 100644 --- a/api/streaming.py +++ b/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( diff --git a/docs/rfcs/webui-run-state-consistency-contract.md b/docs/rfcs/webui-run-state-consistency-contract.md index 82355ac4..e3da5bac 100644 --- a/docs/rfcs/webui-run-state-consistency-contract.md +++ b/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 diff --git a/tests/test_agent_max_turns_parity.py b/tests/test_agent_max_turns_parity.py index 4665f74f..c83d8777 100644 --- a/tests/test_agent_max_turns_parity.py +++ b/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):" ) diff --git a/tests/test_agent_publication_cancellation.py b/tests/test_agent_publication_cancellation.py new file mode 100644 index 00000000..c08acab8 --- /dev/null +++ b/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 diff --git a/tests/test_issue1897_profile_switch_agent_cache.py b/tests/test_issue1897_profile_switch_agent_cache.py index 684c5a01..2b1b0801 100644 --- a/tests/test_issue1897_profile_switch_agent_cache.py +++ b/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. diff --git a/tests/test_streaming_max_tokens_quota.py b/tests/test_streaming_max_tokens_quota.py index ed352290..e008c252 100644 --- a/tests/test_streaming_max_tokens_quota.py +++ b/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):" )