| """Thread-safe interaction state and event streams for one interactive agent.""" |
|
|
| from __future__ import annotations |
|
|
| from itertools import islice |
| import threading |
| import time |
| import uuid |
| from collections import deque |
| from contextlib import contextmanager |
| from dataclasses import dataclass |
| from typing import Any, Callable, Iterator, Literal, Mapping, TypeAlias, cast |
|
|
| from ..core import json_value |
|
|
|
|
| InteractionState: TypeAlias = Literal[ |
| "idle", |
| "running", |
| "pausing", |
| "interrupting", |
| "paused", |
| ] |
| InteractionAction: TypeAlias = Literal["message", "pause", "force_pause"] |
| SRInteractionAction: TypeAlias = Literal[ |
| "message", |
| "pause", |
| "force_pause", |
| "next_c", |
| "next_r", |
| ] |
| InteractionEventKind: TypeAlias = Literal[ |
| "prompt_added", |
| "context", |
| "assistant_started", |
| "assistant_delta", |
| "assistant_completed", |
| "assistant_failed", |
| "tool_started", |
| "tool_completed", |
| "settings_applied", |
| "settings_failed", |
| "workspace_changed", |
| "execution_completed", |
| "execution_failed", |
| "command_received", |
| "state_changed", |
| "search_position_changed", |
| "topk_updated", |
| ] |
|
|
|
|
| @dataclass(frozen=True) |
| class PendingMessage: |
| """A user message waiting to be inserted at an agent boundary.""" |
|
|
| content: str |
| created_at: float |
|
|
|
|
| class _CancellationSignal: |
| """Read-only Event-like view of the manager's interruption state.""" |
|
|
| def __init__(self, manager: "InteractionManager") -> None: |
| self._manager = manager |
|
|
| def is_set(self) -> bool: |
| return self._manager.is_interrupting |
|
|
|
|
| class InteractionManager: |
| """Own the controls and observable event stream of exactly one agent. |
| |
| Args: |
| event_capacity: Maximum number of retained timeline events. |
| """ |
|
|
| _DELTA_KIND = "assistant_delta" |
| _TERMINAL_KINDS = {"assistant_completed", "assistant_failed"} |
| _EVENT_KINDS = { |
| "prompt_added", |
| "context", |
| "assistant_started", |
| "assistant_delta", |
| "assistant_completed", |
| "assistant_failed", |
| "tool_started", |
| "tool_completed", |
| "settings_applied", |
| "settings_failed", |
| "workspace_changed", |
| "execution_completed", |
| "execution_failed", |
| "command_received", |
| "state_changed", |
| "search_position_changed", |
| "topk_updated", |
| } |
| _STATE_FIELDS = { |
| "version", |
| "state", |
| "messages", |
| "events", |
| "next_sequence", |
| "discarded_through", |
| "event_capacity", |
| } |
|
|
| def __init__(self, *, event_capacity: int = 1000) -> None: |
| if event_capacity < 1: |
| raise ValueError("event_capacity must be positive") |
| self._condition = threading.Condition(threading.RLock()) |
| self._state: InteractionState = "idle" |
| self._messages: deque[PendingMessage] = deque() |
| self._events: deque[dict[str, Any]] = deque(maxlen=event_capacity) |
| self._next_sequence = 1 |
| self._discarded_through = 0 |
| self._cancel_current: Callable[[], str | None] | None = None |
| self._cancellation_signal = _CancellationSignal(self) |
|
|
| @property |
| def state(self) -> InteractionState: |
| """Return the authoritative execution state. |
| |
| Returns: |
| Current interaction state. |
| """ |
| with self._condition: |
| return self._state |
|
|
| @property |
| def is_interrupting(self) -> bool: |
| """Return whether the active operation should abort promptly. |
| |
| Returns: |
| Whether a force-pause request is interrupting active work. |
| """ |
| return self.state == "interrupting" |
|
|
| @property |
| def cancellation_signal(self) -> _CancellationSignal: |
| """Expose a read-only Event-like cancellation adapter to tools. |
| |
| Returns: |
| Object exposing ``is_set()`` for cancellation-aware tools. |
| """ |
| return self._cancellation_signal |
|
|
| def status(self) -> dict[str, Any]: |
| """Return a serializable control snapshot. |
| |
| Returns: |
| Current state, pending-input counts, event cursor, and server time. |
| """ |
| with self._condition: |
| return { |
| "interaction_state": self._state, |
| "paused": self._state in {"pausing", "interrupting", "paused"}, |
| "waiting_at_boundary": self._state == "paused", |
| "force_pause_requested": self._state == "interrupting", |
| "pending_messages": len(self._messages), |
| "last_event_seq": self._next_sequence - 1, |
| "server_time": time.time(), |
| } |
|
|
| def command(self, action: InteractionAction, message: str | None = None) -> dict[str, Any]: |
| """Apply a user control command to this agent. |
| |
| Args: |
| action: Message, pause, or force-pause command. |
| message: User guidance required by the ``message`` action. |
| |
| Returns: |
| Control snapshot after applying the command. |
| |
| Raises: |
| ValueError: If the action is invalid or a message is empty. |
| """ |
| if action not in {"message", "pause", "force_pause"}: |
| raise ValueError("action must be message, pause, or force_pause") |
| cancel: Callable[[], str | None] | None = None |
| with self._condition: |
| if action == "message": |
| content = (message or "").strip() |
| if not content: |
| raise ValueError("message command requires non-empty message") |
| self._messages.append(PendingMessage(content, time.time())) |
| elif action == "pause": |
| if self._state == "running": |
| self._set_state_locked("pausing") |
| elif self._state in {"running", "pausing"}: |
| self._set_state_locked("interrupting") |
| cancel = self._cancel_current |
| self._condition.notify_all() |
| cancellation_error = cancel() if cancel is not None else None |
| with self._condition: |
| if cancellation_error and self._state == "interrupting": |
| self._set_state_locked("pausing") |
| if action != "message": |
| payload = {"action": action} |
| if cancellation_error: |
| payload["error"] = cancellation_error |
| self._publish_event_locked("command_received", payload) |
| result = self.status() |
| if cancellation_error: |
| result["cancellation_error"] = cancellation_error |
| return result |
|
|
| def request_pause(self) -> None: |
| """Request a quiet pause, such as after a tool-free response.""" |
| with self._condition: |
| if self._state == "running": |
| self._set_state_locked("pausing") |
| self._condition.notify_all() |
|
|
| def start_agent_execution(self) -> list[PendingMessage]: |
| """Enter running state and consume messages that start this execution. |
| |
| Returns: |
| Messages queued before execution started. |
| |
| Raises: |
| RuntimeError: If the manager is not idle. |
| """ |
| with self._condition: |
| if self._state != "idle": |
| raise RuntimeError(f"cannot start agent execution while {self._state}") |
| messages = self._consume_messages_locked() |
| self._set_state_locked("running") |
| return messages |
|
|
| def finish_agent_execution(self) -> None: |
| """Mark a naturally completed agent execution as idle. |
| |
| Raises: |
| RuntimeError: If execution is paused at a boundary. |
| """ |
| with self._condition: |
| if self._state == "paused": |
| raise RuntimeError("cannot finish an execution while paused") |
| self._cancel_current = None |
| self._set_state_locked("idle") |
|
|
| @contextmanager |
| def wait(self) -> Iterator[list[PendingMessage]]: |
| """Wait at a safe boundary and yield queued messages. |
| |
| Yields: |
| Messages to insert before the next model request. |
| |
| Raises: |
| RuntimeError: If no agent execution is active. |
| """ |
| with self._condition: |
| if self._state == "running": |
| messages = self._consume_messages_locked() |
| paused = False |
| elif self._state in {"pausing", "interrupting", "paused"}: |
| if self._state != "paused": |
| self._set_state_locked("paused") |
| self._cancel_current = None |
| while not self._has_boundary_input_locked(): |
| self._condition.wait() |
| messages = self._consume_messages_locked() |
| paused = True |
| else: |
| raise RuntimeError(f"cannot wait at an agent boundary while {self._state}") |
| try: |
| yield messages |
| finally: |
| if paused: |
| with self._condition: |
| if self._state == "paused": |
| self._set_state_locked("running") |
|
|
| @contextmanager |
| def cancellable(self, cancel: Callable[[], str | None]) -> Iterator[None]: |
| """Register cancellation for the current blocking operation. |
| |
| Args: |
| cancel: Callback that promptly interrupts the active operation. |
| |
| Yields: |
| Control while the callback is registered. |
| |
| Raises: |
| RuntimeError: If another cancellable operation is already active. |
| """ |
| with self._condition: |
| if self._cancel_current is not None: |
| raise RuntimeError("another cancellable operation is already active") |
| self._cancel_current = cancel |
| cancel_immediately = self._state == "interrupting" |
| if cancel_immediately: |
| cancellation_error = cancel() |
| if cancellation_error: |
| with self._condition: |
| if self._state == "interrupting": |
| self._set_state_locked("pausing") |
| self._publish_event_locked( |
| "command_received", |
| {"action": "force_pause", "error": cancellation_error}, |
| ) |
| self._condition.notify_all() |
| try: |
| yield |
| finally: |
| with self._condition: |
| if self._cancel_current is cancel: |
| self._cancel_current = None |
|
|
| def publish_event( |
| self, kind: InteractionEventKind, payload: Mapping[str, Any] |
| ) -> dict[str, Any]: |
| """Append one observable event to this agent's timeline. |
| |
| Args: |
| kind: Supported event kind. |
| payload: JSON-serializable event data. |
| |
| Returns: |
| Stored event with sequence, identifier, and timestamp. |
| |
| Raises: |
| ValueError: If ``kind`` is unsupported. |
| """ |
| if kind not in self._EVENT_KINDS: |
| raise ValueError(f"unsupported interaction event kind: {kind}") |
| with self._condition: |
| return self._publish_event_locked(kind, payload) |
|
|
| def get_recent_events(self, after_sequence: int = 0, *, limit: int | None = None) -> dict[str, Any]: |
| """Return retained events newer than a consumer-owned sequence cursor. |
| |
| Args: |
| after_sequence: Last sequence already consumed by the caller. |
| limit: Maximum events per batch; omitted for an unlimited internal snapshot. |
| |
| Returns: |
| Event batch and cursor-reset or truncation metadata. |
| """ |
| with self._condition: |
| latest_sequence = self._next_sequence - 1 |
| cursor_reset = after_sequence > latest_sequence |
| effective_sequence = 0 if cursor_reset else after_sequence |
| return { |
| "events": [ |
| dict(event) for event in islice( |
| (item for item in self._events if item["seq"] > effective_sequence), |
| limit, |
| ) |
| ], |
| "truncated": effective_sequence < self._discarded_through, |
| "discarded_through": self._discarded_through, |
| "cursor_reset": cursor_reset, |
| "server_time": time.time(), |
| } |
|
|
| def export_state(self) -> dict[str, Any]: |
| """Return the durable portion of this manager's state. |
| |
| Returns: |
| Versioned interaction snapshot without runtime locks or callbacks. |
| """ |
| with self._condition: |
| return { |
| "version": 1, |
| "state": self._state, |
| "messages": [ |
| {"content": message.content, "created_at": message.created_at} |
| for message in self._messages |
| ], |
| "events": [dict(event) for event in self._events], |
| "next_sequence": self._next_sequence, |
| "discarded_through": self._discarded_through, |
| "event_capacity": self._events.maxlen, |
| } |
|
|
| def restore_state(self, snapshot: Mapping[str, Any]) -> bool: |
| """Restore durable state and interrupt work that died with the process. |
| |
| Runtime callbacks and locks are deliberately never restored. |
| |
| Args: |
| snapshot: State returned by :meth:`export_state`. |
| |
| Returns: |
| Whether in-flight work was converted into an interruption event. |
| """ |
| with self._condition: |
| if set(snapshot) != self._STATE_FIELDS: |
| raise ValueError("Persisted interaction state does not match the current schema") |
| if snapshot["version"] != 1: |
| raise ValueError("Unsupported InteractionManager snapshot version") |
| previous = snapshot["state"] |
| if previous not in {"idle", "running", "pausing", "interrupting", "paused"}: |
| raise ValueError(f"Unknown persisted interaction state: {previous!r}") |
| capacity = snapshot["event_capacity"] |
| if not isinstance(capacity, int) or isinstance(capacity, bool) or capacity < 1: |
| raise ValueError("Persisted event_capacity must be positive") |
| events = snapshot["events"] |
| if not isinstance(events, list) or any( |
| not isinstance(event, Mapping) for event in events |
| ): |
| raise TypeError("Persisted events must be a list of mappings") |
| required_event_fields = {"seq", "id", "kind", "timestamp", "payload"} |
| if any(set(event) != required_event_fields for event in events): |
| raise ValueError("Persisted events must use the current event schema") |
| if any(event["kind"] not in self._EVENT_KINDS for event in events): |
| raise ValueError("Persisted events contain an unsupported event kind") |
| if any( |
| not isinstance(event["seq"], int) |
| or isinstance(event["seq"], bool) |
| or event["seq"] < 1 |
| or not isinstance(event["id"], str) |
| or not event["id"] |
| or not isinstance(event["timestamp"], (int, float)) |
| or isinstance(event["timestamp"], bool) |
| or not isinstance(event["payload"], Mapping) |
| for event in events |
| ): |
| raise TypeError("Persisted events contain invalid field values") |
| if [event["seq"] for event in events] != sorted({event["seq"] for event in events}): |
| raise ValueError("Persisted event sequences must be unique and increasing") |
| self._events = deque( |
| (dict(event) for event in events), |
| maxlen=capacity, |
| ) |
| messages = snapshot["messages"] |
| if not isinstance(messages, list) or any( |
| not isinstance(item, Mapping) for item in messages |
| ): |
| raise TypeError("Persisted messages must be a list of mappings") |
| if any(set(item) != {"content", "created_at"} for item in messages): |
| raise ValueError("Persisted messages must use the current message schema") |
| if any( |
| not isinstance(item["content"], str) |
| or not item["content"] |
| or not isinstance(item["created_at"], (int, float)) |
| or isinstance(item["created_at"], bool) |
| for item in messages |
| ): |
| raise TypeError("Persisted messages contain invalid field values") |
| self._messages = deque( |
| PendingMessage(item["content"], item["created_at"]) for item in messages |
| ) |
| discarded_through = snapshot["discarded_through"] |
| if ( |
| not isinstance(discarded_through, int) |
| or isinstance(discarded_through, bool) |
| or discarded_through < 0 |
| ): |
| raise ValueError("Persisted discarded_through must be a non-negative integer") |
| self._discarded_through = discarded_through |
| highest_sequence = max( |
| (event["seq"] for event in self._events), |
| default=0, |
| ) |
| next_sequence = snapshot["next_sequence"] |
| if ( |
| not isinstance(next_sequence, int) |
| or isinstance(next_sequence, bool) |
| or next_sequence <= highest_sequence |
| or next_sequence <= discarded_through |
| ): |
| raise ValueError("Persisted next_sequence is inconsistent with retained events") |
| self._next_sequence = next_sequence |
| interrupted = previous in {"running", "pausing", "interrupting", "paused"} |
| self._state = "idle" |
| self._cancel_current = None |
| if interrupted: |
| self._publish_event_locked( |
| "command_received", |
| {"action": "force_pause", "restored_after_restart": True}, |
| ) |
| self._publish_event_locked( |
| "state_changed", |
| {"previous": previous, "current": "idle", "interrupted": True}, |
| ) |
| return interrupted |
|
|
| def _consume_messages_locked(self) -> list[PendingMessage]: |
| messages = list(self._messages) |
| self._messages.clear() |
| return messages |
|
|
| def _has_boundary_input_locked(self) -> bool: |
| return bool(self._messages) |
|
|
| def _set_state_locked(self, state: InteractionState) -> None: |
| previous = self._state |
| if previous == state: |
| return |
| self._state = state |
| self._publish_event_locked("state_changed", {"previous": previous, "current": state}) |
|
|
| def _publish_event_locked( |
| self, kind: InteractionEventKind, payload: Mapping[str, Any] |
| ) -> dict[str, Any]: |
| payload_dict = cast(dict[str, Any], json_value(dict(payload))) |
| sequence = self._next_sequence |
| self._next_sequence += 1 |
| event = { |
| "seq": sequence, |
| "id": uuid.uuid4().hex, |
| "kind": kind, |
| "timestamp": time.time(), |
| "payload": payload_dict, |
| } |
| response_id = payload_dict.get("response_id") |
| if response_id and kind == self._DELTA_KIND: |
| self._remove_stream_events_locked(response_id, {kind}) |
| elif response_id and kind in self._TERMINAL_KINDS: |
| self._remove_stream_events_locked(response_id, {self._DELTA_KIND}) |
| if len(self._events) == self._events.maxlen and self._events: |
| self._discarded_through = max(self._discarded_through, self._events[0]["seq"]) |
| self._events.append(event) |
| return dict(event) |
|
|
| def _remove_stream_events_locked(self, response_id: str, kinds: set[str]) -> None: |
| self._events = deque( |
| ( |
| event |
| for event in self._events |
| if not ( |
| event["kind"] in kinds |
| and event.get("payload", {}).get("response_id") == response_id |
| ) |
| ), |
| maxlen=self._events.maxlen, |
| ) |
|
|
|
|
| class SRInteractionManager(InteractionManager): |
| """Interaction manager with symbolic-regression branch commands.""" |
|
|
| _STATE_FIELDS = InteractionManager._STATE_FIELDS | {"search_transitions"} |
|
|
| def __init__(self, *, event_capacity: int = 1000) -> None: |
| super().__init__(event_capacity=event_capacity) |
| self._search_transitions: deque[Literal["next_c", "next_r"]] = deque() |
|
|
| def command(self, action: SRInteractionAction, message: str | None = None) -> dict[str, Any]: |
| """Apply a common command or queue an SR branch transition. |
| |
| Args: |
| action: Common interaction action or ``next_c``/``next_r``. |
| message: User guidance for a ``message`` action. |
| |
| Returns: |
| Control snapshot after applying the command. |
| """ |
| if action not in {"next_c", "next_r"}: |
| return super().command(cast(InteractionAction, action), message) |
| with self._condition: |
| self._search_transitions.clear() |
| self._search_transitions.append(action) |
| self._publish_event_locked("command_received", {"action": action}) |
| self._condition.notify_all() |
| return self.status() |
|
|
| def consume_search_transition(self) -> Literal["next_c", "next_r"] | None: |
| """Consume the newest queued branch transition. |
| |
| Returns: |
| Queued transition, or ``None`` when no transition is pending. |
| """ |
| with self._condition: |
| return self._search_transitions.popleft() if self._search_transitions else None |
|
|
| def _has_boundary_input_locked(self) -> bool: |
| return bool(self._messages or self._search_transitions) |
|
|
| def status(self) -> dict[str, Any]: |
| """Include the pending SR branch transition in the control snapshot. |
| |
| Returns: |
| Common control snapshot with ``pending_transition``. |
| """ |
| with self._condition: |
| status = super().status() |
| status["pending_transition"] = ( |
| self._search_transitions[-1] if self._search_transitions else None |
| ) |
| return status |
|
|
| def export_state(self) -> dict[str, Any]: |
| """Include pending search transitions in the durable snapshot. |
| |
| Returns: |
| Versioned interaction snapshot with symbolic-search transitions. |
| """ |
| with self._condition: |
| state = super().export_state() |
| state["search_transitions"] = list(self._search_transitions) |
| return state |
|
|
| def restore_state(self, snapshot: Mapping[str, Any]) -> bool: |
| """Restore common state plus queued symbolic-search transitions. |
| |
| Args: |
| snapshot: State returned by :meth:`export_state`. |
| |
| Returns: |
| Whether in-flight work was converted into an interruption event. |
| """ |
| interrupted = super().restore_state(snapshot) |
| with self._condition: |
| transitions = snapshot["search_transitions"] |
| if not isinstance(transitions, list) or any( |
| action not in {"next_c", "next_r"} for action in transitions |
| ): |
| raise ValueError("Persisted search_transitions contains an invalid action") |
| self._search_transitions = deque(transitions) |
| return interrupted |
|
|
