File size: 24,146 Bytes
d06fa57 8956871 d06fa57 8956871 d06fa57 8956871 d06fa57 8956871 d06fa57 8956871 d06fa57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 | """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:
"""Queued user guidance message waiting to be injected at the next safe agent boundary.
Attributes:
content: Non-empty user message text.
created_at: Unix timestamp when the message was queued.
"""
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:
"""Thread-safe controller managing execution state, pause/stop commands, and timeline events for one agent.
Args:
event_capacity: Maximum number of recent timeline events retained in memory.
"""
_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:
"""Initialize the interaction manager with an empty event ring buffer.
Args:
event_capacity: Maximum number of timeline events to keep in memory (defaults to ``1000``).
"""
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:
"""Initialize the symbolic-regression interaction manager.
Args:
event_capacity: Maximum number of timeline events retained in memory (defaults to ``1000``).
"""
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
|