diff --git a/gateway/run_agent_cache.py b/gateway/run_agent_cache.py index f2ddf9e64b..d503c82b8b 100644 --- a/gateway/run_agent_cache.py +++ b/gateway/run_agent_cache.py @@ -209,6 +209,16 @@ class GatewayAgentCacheMixin: override = self._session_model_override(session_key) return {"had_override": override is not None, "override": dict(override) if override is not None else None} + def _claim_one_turn_restore(self, session_key: str, snapshot: Optional[dict] = None) -> None: + """Arm the one-shot restore snapshot for ``/model --once`` / ``/moa``. A repeated one-shot + command before the turn runs keeps the EARLIEST snapshot: the later command's snapshot is + the first temporary model, not the user's standing override. Pass *snapshot* when the + caller captured the pre-switch state earlier (``/model --once`` applies its override before + arming); omit it to snapshot now.""" + conv = self._session_state(session_key).conversation + if not conv.one_turn_restore: + conv.one_turn_restore = dict(snapshot) if snapshot is not None else self._snapshot_session_model_override(session_key) + def _restore_session_model_override(self, session_key: str, snapshot: dict) -> None: """Restore the session override captured before a one-turn switch.""" if not session_key: @@ -256,14 +266,33 @@ class GatewayAgentCacheMixin: self._persist_active_agents() return True - def _drop_turn_slot(self, session_key: str) -> None: + def _drop_turn_slot(self, session_key: str, *, run_generation: Optional[int] = None) -> None: """Release the running-agent slot and evict the cached instance (/stop, eviction, reaper). ``_interrupt_requested`` is cleared only by the turn finalizer, so on a hung/still-draining run the flag would survive and silently kill the session's NEXT message (interrupted=True, api_calls=0, empty response); the next message rebuilds from history while the old agent - keeps its flag so a hung drain still dies (#44212).""" - self._release_running_agent_state(session_key) + keeps its flag so a hung drain still dies (#44212). With ``run_generation`` (the post-bump + value ``_interrupt_running_turn`` returns), the release is generation-guarded: an async + path awaits between bump and release, so a successor claiming the slot in that window must + not have its sentinel/lease wiped by the displaced path's tail. Then sweep lease tokens + from generations OLDER than the current one: a hung evicted turn's finalizer may never run, + and each such generation would otherwise pin its token (and its ``_SessionLease``) forever. + Identity-checked + idempotent, so a live successor's token is never affected.""" + self._release_running_agent_state(session_key, run_generation=run_generation) self._evict_cached_agent(session_key) + state = self._peek_session_state(session_key) + registry = getattr(self, "_turn_leases", None) + if state is None or registry is None: + return + current = int(state.persistent.run_generation or 0) + tokens = state.turn.lease_tokens + for gen in [g for g in tokens if int(g) < current]: + token = tokens.pop(gen) + try: + registry.release(token) + except Exception: + logger.debug("Failed to release displaced turn lease gen %s for %s", gen, session_key, + exc_info=True) def _held_turn_lease(self, session_key: str, run_generation: int): """Return ``(registry, lease_tokens)`` when ``session_key`` holds a lease token for @@ -436,7 +465,7 @@ class GatewayAgentCacheMixin: if not session_key: return state = self._peek_session_state(session_key) - self._interrupt_running_turn( + _generation_at_interrupt = self._interrupt_running_turn( session_key, interrupt_reason=interrupt_reason, invalidation_reason=invalidation_reason, ) adapter = self._adapter_for_source(source) @@ -452,7 +481,9 @@ class GatewayAgentCacheMixin: if state is not None: state.persistent.pending_command_text = None if release_running_state: - self._drop_turn_slot(session_key) + # Guarded release: a message that arrived during the awaits above may already run as + # the successor generation — the displaced /stop tail must not wipe its slot. + self._drop_turn_slot(session_key, run_generation=_generation_at_interrupt) async def _refresh_agent_cache_message_count(self, session_key: str, session_id: Optional[str]) -> None: """Re-baseline a cached agent's stored message_count after THIS turn — the coherence guard diff --git a/gateway/run_inbound.py b/gateway/run_inbound.py index 7dcd4cdf91..3c5028cb61 100644 --- a/gateway/run_inbound.py +++ b/gateway/run_inbound.py @@ -490,8 +490,9 @@ class GatewayInboundMixin: def _hm_evict_running_agent(self, _quick_key: str, reason: str) -> None: from gateway.run import _INTERRUPT_REASON_EVICTED - self._interrupt_running_turn(_quick_key, interrupt_reason=_INTERRUPT_REASON_EVICTED, invalidation_reason=reason) - self._drop_turn_slot(_quick_key) + _generation_at_interrupt = self._interrupt_running_turn( + _quick_key, interrupt_reason=_INTERRUPT_REASON_EVICTED, invalidation_reason=reason) + self._drop_turn_slot(_quick_key, run_generation=_generation_at_interrupt) def _hm_merge_pending_for_source( self, source: SessionSource, _quick_key: str, event: "MessageEvent", *, merge_text: bool = False @@ -903,8 +904,7 @@ class GatewayInboundMixin: event.text = moa_payload _moa_state = self._session_state(_quick_key) # Same one-shot snapshot `/model --once` uses, so eviction/stop/finalizer settle both alike. - if not _moa_state.conversation.one_turn_restore: - _moa_state.conversation.one_turn_restore = self._snapshot_session_model_override(_quick_key) + self._claim_one_turn_restore(_quick_key) _moa_state.conversation.model_override = { "provider": "moa", "model": moa_cfg["default_preset"], "base_url": "moa://local", "api_key": "moa-virtual-provider", "api_mode": "chat_completions", diff --git a/gateway/slash_commands_model.py b/gateway/slash_commands_model.py index 875707a874..c4ab97555a 100644 --- a/gateway/slash_commands_model.py +++ b/gateway/slash_commands_model.py @@ -267,13 +267,9 @@ class GatewayModelCommandsMixin: "capabilities": dict(result.runtime_capabilities or {}), } if one_turn: - if not hasattr(self, "_pending_one_turn_model_restores"): - self._pending_one_turn_model_restores = {} # A repeated --once before the turn runs must keep the EARLIEST snapshot: the later # command's snapshot is the first temporary model, not the user's standing override. - self._pending_one_turn_model_restores.setdefault( - ctx.session_key, dict(ctx.restore_snapshot or {"had_override": False, "override": None}), - ) + self._claim_one_turn_restore(ctx.session_key, ctx.restore_snapshot) elif not picker and hasattr(self, "_pending_one_turn_model_restores"): self._pending_one_turn_model_restores.pop(ctx.session_key, None) # Non-secret write-through so the override survives a restart (api_key/api_mode are