diff --git a/api/streaming.py b/api/streaming.py index b61490fb..4dfa5b63 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -9791,6 +9791,84 @@ 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, 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. + with STREAMS_LOCK: + 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 + # 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") + except Exception: + logger.debug("Failed to interrupt cancelled candidate agent") + 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, + 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) + # 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` # block can safely check `if _checkpoint_stop is not None` even when an # exception fires before the checkpoint thread is created (Issue #765). @@ -11013,6 +11091,7 @@ def _run_agent_streaming( # ── Agent cache: reuse across messages in the same session ── # Mirrors gateway _agent_cache. Keeps _user_turn_count alive so # injectionFrequency: "first-turn" actually suppresses after turn 1. + _cache_new_agent = False if ephemeral: agent = _AIAgent(**_agent_kwargs) logger.debug('[webui] Created ephemeral agent for session %s', session_id) @@ -11055,14 +11134,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: @@ -11132,82 +11203,65 @@ def _run_agent_streaming( agent._interrupt_message = None else: agent = _AIAgent(**_agent_kwargs) - # 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 - # closes an agent mid-run AND we never nest ACTIVE_RUNS_LOCK - # inside SESSION_AGENT_CACHE_LOCK (avoids any lock-ordering - # deadlock). A cancel/reconnect can drop STREAMS while the - # worker is still unwinding or blocked in a provider call, so - # ACTIVE_RUNS (worker lifecycle) is the authoritative liveness - # signal, not STREAMS. (#3536 review round 2) - _active_sids = set() - try: - from api.config import ACTIVE_RUNS, ACTIVE_RUNS_LOCK - with ACTIVE_RUNS_LOCK: - for _entry in (ACTIVE_RUNS or {}).values(): - _sid = (_entry or {}).get("session_id") - if _sid: - _active_sids.add(_sid) - 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 - # every over-cap entry is active, leave the cache - # temporarily above cap rather than close a live worker's - # agent — a later insertion/finalization trims it once the - # run ends. - while len(SESSION_AGENT_CACHE) > SESSION_AGENT_CACHE_MAX: - _evictable_sid = None - for _sid in list(SESSION_AGENT_CACHE.keys()): - if _sid not in _active_sids: - _evictable_sid = _sid - break - if _evictable_sid is None: - break # all over-cap entries are active; defer - evicted_entry = SESSION_AGENT_CACHE.pop(_evictable_sid) - _evicted_items.append((_evictable_sid, evicted_entry)) - # Commit and close evicted agents outside the cache lock so - # concurrent cache users are not blocked by provider I/O. - for _evicted_sid, _evicted_entry in _evicted_items: - try: - _evicted_agent = _evicted_entry[0] if isinstance(_evicted_entry, tuple) else None - _close_evicted_agent_at_session_boundary(_evicted_sid, _evicted_agent) - except Exception: - logger.debug("Failed to close evicted agent for session %s", _evicted_sid, exc_info=True) - logger.debug('[webui] Evicted LRU agent from cache: %s', _evicted_sid) - logger.debug('[webui] Created new agent for session %s', session_id) + _cache_new_agent = True - # Stop may already have detached CANCEL_FLAGS while initialization - # was in flight. The worker-owned event and stream membership, not - # the removable registry entry, fence registration with cancellation. - with STREAMS_LOCK: - _cancelled_before_start = cancel_event.is_set() or stream_id not in STREAMS - if not _cancelled_before_start: - AGENT_INSTANCES[stream_id] = agent - if _cancelled_before_start: - # Interrupt, persistence and event writes must not hold the - # registry lock (the session lock can be held by another caller). - try: - agent.interrupt("Cancelled before start") - except Exception: - logger.debug("Failed to interrupt agent before start") + 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 + # Cache publication is already fenced with Stop above. Keep LRU + # eviction and any memory-provider close outside the stream lock. + if _cache_new_agent: + _evicted_items = [] + # Snapshot the set of session_ids with a LIVE agent worker + # BEFORE taking SESSION_AGENT_CACHE_LOCK, so LRU eviction never + # closes an agent mid-run AND we never nest ACTIVE_RUNS_LOCK + # inside SESSION_AGENT_CACHE_LOCK (avoids any lock-ordering + # deadlock). A cancel/reconnect can drop STREAMS while the + # worker is still unwinding or blocked in a provider call, so + # ACTIVE_RUNS (worker lifecycle) is the authoritative liveness + # signal, not STREAMS. (#3536 review round 2) + _active_sids = set() + try: + from api.config import ACTIVE_RUNS, ACTIVE_RUNS_LOCK + with ACTIVE_RUNS_LOCK: + for _entry in (ACTIVE_RUNS or {}).values(): + _sid = (_entry or {}).get("session_id") + if _sid: + _active_sids.add(_sid) + except Exception: + _active_sids = set() + with SESSION_AGENT_CACHE_LOCK: + 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 + # every over-cap entry is active, leave the cache + # temporarily above cap rather than close a live worker's + # agent — a later insertion/finalization trims it once the + # run ends. + while len(SESSION_AGENT_CACHE) > SESSION_AGENT_CACHE_MAX: + _evictable_sid = None + for _sid in list(SESSION_AGENT_CACHE.keys()): + if _sid not in _active_sids: + _evictable_sid = _sid + break + if _evictable_sid is None: + break # all over-cap entries are active; defer + evicted_entry = SESSION_AGENT_CACHE.pop(_evictable_sid) + _evicted_items.append((_evictable_sid, evicted_entry)) + # Commit and close evicted agents outside the cache lock so + # concurrent cache users are not blocked by provider I/O. + for _evicted_sid, _evicted_entry in _evicted_items: + try: + _evicted_agent = _evicted_entry[0] if isinstance(_evicted_entry, tuple) else None + _close_evicted_agent_at_session_boundary(_evicted_sid, _evicted_agent) + except Exception: + logger.debug("Failed to close evicted agent for session %s", _evicted_sid, exc_info=True) + logger.debug('[webui] Evicted LRU agent from cache: %s', _evicted_sid) + logger.debug('[webui] Created new agent for session %s', session_id) + # Prepend workspace context so the agent always knows which directory # to use for file operations, regardless of session age or AGENTS.md defaults. workspace_ctx = _workspace_context_prefix(str(s.workspace)) @@ -11467,6 +11521,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( @@ -11949,8 +12008,6 @@ def _run_agent_streaming( if 'credential_pool' in _agent_params: _agent_kwargs['credential_pool'] = _runtime_bundle['credential_pool'] agent = _AIAgent(**_agent_kwargs) - with STREAMS_LOCK: - AGENT_INSTANCES[stream_id] = agent _agent_sig = _compute_agent_cache_signature( resolved_model, resolved_api_key, @@ -11967,10 +12024,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 @@ -12003,6 +12061,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( @@ -13301,8 +13364,6 @@ def _run_agent_streaming( if 'credential_pool' in _agent_params: _heal_kwargs['credential_pool'] = _runtime_bundle['credential_pool'] _heal_agent = _AIAgent(**_heal_kwargs) - with STREAMS_LOCK: - AGENT_INSTANCES[stream_id] = _heal_agent _agent_sig = _compute_agent_cache_signature( resolved_model, resolved_api_key, @@ -13319,10 +13380,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: @@ -13354,6 +13416,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 d1362548..4d1c4e19 100644 --- a/docs/rfcs/webui-run-state-consistency-contract.md +++ b/docs/rfcs/webui-run-state-consistency-contract.md @@ -30,6 +30,40 @@ This RFC defines a consistency contract for those layers. It complements the larger run adapter direction in #1925 by documenting what must remain coherent while WebUI still has multiple overlapping state stores. +## Agent registration after cancellation + +Initial and credential self-heal Agent construction use the same registration +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`. 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, 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. + +Interrupt and cancellation finalization occur outside the stream registry lock. +Session finalization still owns the session lock: the returned-error path already +holds it, while initial registration and the exception path acquire it. Do not +reacquire this non-reentrant lock from a branch that already owns it. + ## Inactive compression continuation recovery The Agent profile's SQLite compression lineage owns the canonical continuation, 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..5f854def --- /dev/null +++ b/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 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_self_heal_worker_cancellation.py b/tests/test_self_heal_worker_cancellation.py new file mode 100644 index 00000000..3d8f3327 --- /dev/null +++ b/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 diff --git a/tests/test_steer_worker_boundaries.py b/tests/test_steer_worker_boundaries.py index d55d3691..a97d450b 100644 --- a/tests/test_steer_worker_boundaries.py +++ b/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 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):" )