The persisted pin carried rendered session-context bytes but not the privacy.redact_pii they were rendered under. After a restart, the first internal event restored and reused them verbatim, so bytes rendered with redaction off reached the model after redaction was switched on. Base re-rendered under the current setting after a restart. The pin now records redact_pii (the in-memory ephemeral_pin becomes (change_key, text, redact_pii); the durable schema requires the flag). Bytes from a different policy are never reused, including by internal events in the same process, which had the same bypass until the next human turn. Simplified on the way: - sanitize_prompt_pin is the one validator; the snapshot builder and rehydrate no longer re-check the shape by hand, and use PROMPT_PIN_VERSION instead of a literal. - Rehydrate runs once, from the context path, and only when this process holds no pins for the key (a restart). Eviction clears only the context pin, so an evicted agent still re-renders instead of reviving the snapshot. The channel path no longer rehydrates or takes a session id. - _persist_prompt_pins drops its synchronous get_prompt_pin pre-check, which took the store lock on the event loop; set_prompt_pin already no-ops an unchanged snapshot. Its call-count test goes with it. - set_prompt_pin builds its candidate with dataclasses.replace like set_model_override, and no longer accepts None (no caller cleared). - The read-fence test folds into the restart/reset test.
270 lines
12 KiB
Python
270 lines
12 KiB
Python
"""Per-session gateway state in one container (replaces ~19 session_key-keyed dicts on
|
|
GatewayRunner that bred boundary drift and wholesale-reset races). Scopes follow where each dict
|
|
was CLEARED: ``turn`` at the end of every turn; ``conversation`` at conversation boundaries (/new,
|
|
/resume, auto-reset, expiry); ``persistent`` fields have their own lifecycles."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import MutableMapping
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Callable, Dict, Iterator, List, NamedTuple, Optional, Tuple
|
|
|
|
# /fast stores "priority" or None (explicit normal), so key PRESENCE decides, not truthiness.
|
|
_UNSET_TIER = object()
|
|
SERVICE_TIER_UNSET = _UNSET_TIER # public alias
|
|
|
|
|
|
@dataclass
|
|
class TurnState:
|
|
"""State scoped to one running gateway turn. ``lease_tokens`` is NOT touched by
|
|
``clear()``: ``_release_turn_lease`` owns it (release exactly once)."""
|
|
|
|
agent: Any = None # running AIAgent (or _AGENT_PENDING_SENTINEL); None = idle
|
|
# The MessageEvent that opened the running turn and the live TurnContext: a successful busy
|
|
# redirect re-anchors both to the redirecting message (#115001).
|
|
event: Any = None
|
|
ctx: Any = None
|
|
started_ts: float = 0.0 # 0.0 = not running
|
|
lease: Any = None # cross-process active-session slot lease
|
|
busy_ack_ts: float = 0.0 # debounce; 0.0 = never acked
|
|
# Held turn-lease tokens keyed by acquiring run generation: release/rebind resolve the
|
|
# token for their own generation, so a displaced turn's unwind frees only its own lease and
|
|
# never a successor's (an evicted turn and its replacement may both hold one briefly).
|
|
lease_tokens: Dict[int, Any] = field(default_factory=dict)
|
|
|
|
def clear(self) -> None:
|
|
"""Reset the per-turn slot. The caller pops ``lease`` first to release it."""
|
|
self.agent = self.lease = self.event = self.ctx = None
|
|
self.started_ts = self.busy_ack_ts = 0.0
|
|
|
|
|
|
@dataclass
|
|
class ConversationState:
|
|
"""State scoped to one conversation (survives turns, not boundaries)."""
|
|
|
|
model_override: Optional[Dict[str, Any]] = None # /model per-session override
|
|
one_turn_restore: Optional[Dict[str, Any]] = None # /model --once snapshot
|
|
reasoning_override: Optional[Dict[str, Any]] = None # /reasoning override
|
|
service_tier_override: Any = _UNSET_TIER # /fast: "priority" or None; _UNSET_TIER = absent
|
|
last_resolved_model: str = "" # last successfully-resolved non-empty model
|
|
queued_events: List[Any] = field(default_factory=list) # /queue overflow FIFO (head in adapter)
|
|
sidecar_notes: List[str] = field(default_factory=list) # one-shot must-deliver notes
|
|
ephemeral_pin: Optional[Tuple[Any, ...]] = None # pinned session-context (change_key, text, redact_pii)
|
|
# (channel_prompt, parent_chat_id) of the last non-internal turn; internal events reuse it
|
|
channel_pin: Optional[Tuple[Optional[str], Optional[str]]] = None
|
|
vc_last: Optional[str] = None # last voice-channel context delivered
|
|
|
|
def clear(self) -> None:
|
|
"""Reset every field to its default, so new fields are cleared automatically."""
|
|
self.__dict__.update(ConversationState().__dict__)
|
|
|
|
|
|
@dataclass
|
|
class PersistentState:
|
|
"""State with its own lifecycle — NOT cleared wholesale by turn or boundary resets
|
|
(approvals/update prompts ARE cleared, individually, by the boundary security funnel)."""
|
|
|
|
approvals: Optional[Dict[str, Any]] = None # {"command": ..., "pattern_key": ...}
|
|
update_prompt_pending: bool = False # /update prompt awaiting a reply
|
|
native_image_paths: List[str] = field(default_factory=list) # consumed one-shot
|
|
# Legacy runner-level pending text (flushed on shutdown); not the adapter-level one.
|
|
pending_command_text: Optional[str] = None
|
|
run_generation: int = 0 # monotonic; NEVER reset (stale-run detection depends on it)
|
|
# Consecutive hygiene compression failures (the in-agent ladder is unreachable: hygiene builds
|
|
# a FRESH AIAgent per run). Reset on success; process-local, mirrored to the DB by run.py.
|
|
# Monotonic run-generation counter (#28686). NEVER reset: clearing it would break stale-run detection.
|
|
# The in-agent compressor escalates repeat timeouts via ContextCompressor._consecutive_timeout_failures,
|
|
# but hygiene builds a FRESH AIAgent per run and bind_session_state() zeroes that counter, so the
|
|
# in-agent ladder is structurally unreachable from the gateway. Tracking the streak here — outside the
|
|
# per-run agent — lets hygiene escalate its cooldown instead of retrying on a flat interval forever.
|
|
# Reset on a successful compression, not by turn/boundary resets. PROCESS-LOCAL, deliberately:
|
|
# `PersistentState` means "survives turn and boundary resets", NOT "survives a restart" — this field has
|
|
# no disk flush (unlike `pending_command_text` above, #72680), so a gateway restart drops escalation
|
|
# back to rung 1 while the DB-backed deadline itself survives (#74136). Keying on `session_key` rather
|
|
# than `session_id` is what buys correctness across compaction ROTATION (the sid changes, the chat does
|
|
# not). gateway.run mirrors this value to the DB keyed by session_key so the same semantics also survive
|
|
# gateway restarts.
|
|
hygiene_failure_streak: int = 0
|
|
|
|
|
|
@dataclass
|
|
class SessionState:
|
|
"""All per-session gateway state, grouped by lifecycle scope."""
|
|
|
|
turn: TurnState = field(default_factory=TurnState)
|
|
conversation: ConversationState = field(default_factory=ConversationState)
|
|
persistent: PersistentState = field(default_factory=PersistentState)
|
|
|
|
|
|
# --- Legacy dict-view adapters: tests read/write the old dict attributes directly
|
|
# (``runner._running_agents = {}``); each view is a LIVE MutableMapping over one field.
|
|
|
|
|
|
class _FieldSpec(NamedTuple):
|
|
"""One legacy dict: scope attr, field name, default factory, presence test."""
|
|
scope: str
|
|
name: str
|
|
default: Callable[[], Any]
|
|
is_present: Callable[[Any], bool]
|
|
|
|
|
|
def _spec(scope: str, name: str, default: Any) -> _FieldSpec:
|
|
"""``default`` is a type (presence = truthiness) or a sentinel (presence = ``is not``)."""
|
|
if isinstance(default, type):
|
|
return _FieldSpec(scope, name, default, bool)
|
|
return _FieldSpec(scope, name, lambda: default, lambda v: v is not default)
|
|
|
|
|
|
class _RunnerView(MutableMapping):
|
|
"""Shared plumbing: live view over ``runner._sessions``, dict-comparable."""
|
|
|
|
__slots__ = ("_runner",)
|
|
|
|
def __init__(self, runner: Any) -> None:
|
|
self._runner = runner
|
|
|
|
def _sessions(self) -> Dict[str, SessionState]:
|
|
return self._runner.__dict__.get("_sessions") or {}
|
|
|
|
def __len__(self) -> int:
|
|
return sum(1 for _ in self)
|
|
|
|
def __eq__(self, other: object) -> bool: # Mapping has no __eq__; tests compare to dicts
|
|
if isinstance(other, (dict, MutableMapping)):
|
|
return dict(self.items()) == dict(other)
|
|
return NotImplemented
|
|
|
|
|
|
class SessionFieldView(_RunnerView):
|
|
"""Live dict-like view of one SessionState field across sessions."""
|
|
|
|
__slots__ = ("_spec",)
|
|
|
|
def __init__(self, runner: Any, spec: _FieldSpec) -> None:
|
|
super().__init__(runner)
|
|
self._spec = spec
|
|
|
|
def _value(self, state: SessionState) -> Any:
|
|
return getattr(getattr(state, self._spec.scope), self._spec.name)
|
|
|
|
def _set(self, state: SessionState, value: Any) -> None:
|
|
setattr(getattr(state, self._spec.scope), self._spec.name, value)
|
|
|
|
def _present(self, key: Any) -> Optional[SessionState]:
|
|
"""The session state for ``key`` if its field is present, else None."""
|
|
state = self._sessions().get(key)
|
|
return state if state is not None and self._spec.is_present(self._value(state)) else None
|
|
|
|
def _held(self, key: str) -> SessionState:
|
|
state = self._present(key)
|
|
if state is None:
|
|
raise KeyError(key)
|
|
return state
|
|
|
|
def __getitem__(self, key: str) -> Any:
|
|
return self._value(self._held(key))
|
|
|
|
def __setitem__(self, key: str, value: Any) -> None:
|
|
self._set(self._runner._session_state(key), value)
|
|
|
|
def __delitem__(self, key: str) -> None:
|
|
self._set(self._held(key), self._spec.default())
|
|
|
|
def __iter__(self) -> Iterator[str]:
|
|
return (k for k in list(self._sessions()) if self._present(k) is not None)
|
|
|
|
def __contains__(self, key: object) -> bool:
|
|
return self._present(key) is not None
|
|
|
|
def clear(self) -> None: # avoid MutableMapping's popitem loop
|
|
for state in list(self._sessions().values()):
|
|
self._set(state, self._spec.default())
|
|
|
|
def __repr__(self) -> str: # pragma: no cover - debug aid
|
|
return f"SessionFieldView({self._spec.scope}.{self._spec.name}, {dict(self.items())!r})"
|
|
|
|
|
|
class TurnLeaseTokenView(_RunnerView):
|
|
"""Legacy (session_key, generation)-keyed view of ``_turn_lease_tokens``; the lease registry
|
|
serializes acquisition per session, so one ``TurnState`` slot per key equals the old dict."""
|
|
|
|
__slots__ = ()
|
|
|
|
def _held(self, key: Any) -> TurnState:
|
|
"""TurnState for a currently-held (session_key, generation) or raise KeyError."""
|
|
if not isinstance(key, tuple) or len(key) != 2:
|
|
raise KeyError(key)
|
|
state = self._sessions().get(key[0])
|
|
if state is None or key[1] not in state.turn.lease_tokens:
|
|
raise KeyError(key)
|
|
return state.turn
|
|
|
|
def __getitem__(self, key: Any) -> Any:
|
|
return self._held(key).lease_tokens[key[1]]
|
|
|
|
def __setitem__(self, key: Any, value: Any) -> None:
|
|
if not isinstance(key, tuple) or len(key) != 2:
|
|
raise KeyError(key)
|
|
self._runner._session_state(key[0]).turn.lease_tokens[key[1]] = value
|
|
|
|
def __delitem__(self, key: Any) -> None:
|
|
del self._held(key).lease_tokens[key[1]]
|
|
|
|
def __iter__(self) -> Iterator[Tuple[str, Any]]:
|
|
return ((k, gen) for k, s in list(self._sessions().items()) for gen in list(s.turn.lease_tokens))
|
|
|
|
def clear(self) -> None: # avoid MutableMapping's popitem loop
|
|
for key in list(self):
|
|
del self[key]
|
|
|
|
|
|
# One spec per legacy dict attribute.
|
|
LEGACY_FIELD_SPECS: Dict[str, _FieldSpec] = {
|
|
"_running_agents": _spec("turn", "agent", None),
|
|
"_running_agents_ts": _spec("turn", "started_ts", float),
|
|
"_active_session_leases": _spec("turn", "lease", None),
|
|
"_busy_ack_ts": _spec("turn", "busy_ack_ts", float),
|
|
"_session_model_overrides": _spec("conversation", "model_override", None),
|
|
"_pending_one_turn_model_restores": _spec("conversation", "one_turn_restore", None),
|
|
"_session_reasoning_overrides": _spec("conversation", "reasoning_override", None),
|
|
"_session_service_tier_overrides": _spec("conversation", "service_tier_override", _UNSET_TIER),
|
|
"_last_resolved_model": _spec("conversation", "last_resolved_model", str),
|
|
"_queued_events": _spec("conversation", "queued_events", list),
|
|
"_pending_turn_sidecar_notes": _spec("conversation", "sidecar_notes", list),
|
|
"_session_ephemeral_pin": _spec("conversation", "ephemeral_pin", None),
|
|
"_session_vc_last": _spec("conversation", "vc_last", None),
|
|
"_pending_approvals": _spec("persistent", "approvals", None),
|
|
"_update_prompt_pending": _spec("persistent", "update_prompt_pending", bool),
|
|
"_pending_native_image_paths_by_session": _spec("persistent", "native_image_paths", list),
|
|
"_pending_messages": _spec("persistent", "pending_command_text", None),
|
|
"_session_run_generation": _spec("persistent", "run_generation", int),
|
|
}
|
|
|
|
|
|
def _legacy_property(make_view: Callable[[Any], MutableMapping], doc: str) -> property:
|
|
"""Dict-shaped @property over a live view; the setter takes a plain dict (test pattern
|
|
``runner._X = {...}``): reset the field on every session, then apply the entries."""
|
|
|
|
def fset(self: Any, mapping: Optional[Dict[Any, Any]]) -> None:
|
|
view = make_view(self)
|
|
view.clear()
|
|
view.update(mapping or {})
|
|
|
|
return property(make_view, fset, lambda self: make_view(self).clear(), doc=doc)
|
|
|
|
|
|
def legacy_dict_property(attr_name: str) -> property:
|
|
"""Legacy dict-shaped @property for one migrated attribute."""
|
|
spec = LEGACY_FIELD_SPECS[attr_name]
|
|
return _legacy_property(
|
|
lambda self: SessionFieldView(self, spec),
|
|
f"Legacy dict view over SessionState.{spec.scope}.{spec.name} (for pre-SessionState tests).",
|
|
)
|
|
|
|
|
|
def legacy_lease_token_property() -> property:
|
|
"""Legacy (session_key, generation)-keyed view of held turn-lease tokens."""
|
|
return _legacy_property(
|
|
TurnLeaseTokenView, "Legacy (session_key, generation)-keyed turn-lease token view."
|
|
)
|