mirror of
https://github.com/usestrix/strix.git
synced 2026-08-18 17:52:32 +02:00
537 lines
21 KiB
Python
537 lines
21 KiB
Python
"""SDK-native state for Strix's addressable agent graph."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import tempfile
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Literal, cast
|
|
|
|
from strix.core.sessions import session_write_lock
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
from agents.items import TResponseInputItem
|
|
from agents.memory import Session
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed", "budget_paused"]
|
|
|
|
# Why an agent parked. The user can message any agent, so this - not the agent's
|
|
# position in the tree - decides whether waiting is bounded: only an agent waiting
|
|
# on other agents is re-checked on a timer.
|
|
WaitKind = Literal["user", "agents", "stalled"]
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class AgentRuntime:
|
|
session: Session | None = None
|
|
task: asyncio.Task[Any] | None = None
|
|
stream: Any | None = None
|
|
interrupt_on_message: bool = False
|
|
wake: asyncio.Event = field(default_factory=asyncio.Event)
|
|
mailbox: list[dict[str, Any]] = field(default_factory=list)
|
|
user_wake_required: bool = False
|
|
|
|
|
|
class AgentCoordinator:
|
|
"""Single owner for graph state, SDK runtimes, messages, and resume snapshots."""
|
|
|
|
def __init__(self) -> None:
|
|
self.statuses: dict[str, Status] = {}
|
|
self.parent_of: dict[str, str | None] = {}
|
|
self.names: dict[str, str] = {}
|
|
self.metadata: dict[str, dict[str, Any]] = {}
|
|
self.pending_counts: dict[str, int] = {}
|
|
self.errors: dict[str, str] = {}
|
|
self.recovery_counts: dict[str, int] = {}
|
|
self.idle_resume_counts: dict[str, int] = {}
|
|
self.wait_kinds: dict[str, WaitKind] = {}
|
|
self.runtimes: dict[str, AgentRuntime] = {}
|
|
self._parent_notified: set[str] = set()
|
|
self._lock = asyncio.Lock()
|
|
self._snapshot_path: Path | None = None
|
|
self.is_shutting_down = False
|
|
self._budget_stopped = False
|
|
self._reserve_stopped = False
|
|
self._budget_paused = False
|
|
self._extend_budget: Callable[[], None] | None = None
|
|
|
|
def set_snapshot_path(self, path: Path) -> None:
|
|
self._snapshot_path = path
|
|
|
|
def mark_shutting_down(self) -> None:
|
|
self.is_shutting_down = True
|
|
|
|
@property
|
|
def budget_stopped(self) -> bool:
|
|
return self._budget_stopped
|
|
|
|
async def trigger_budget_stop(self) -> None:
|
|
"""Signal a scan-wide budget stop and wake every parked agent so it exits."""
|
|
async with self._lock:
|
|
self._budget_stopped = True
|
|
for runtime in self.runtimes.values():
|
|
runtime.wake.set()
|
|
|
|
@property
|
|
def reserve_stopped(self) -> bool:
|
|
return self._reserve_stopped
|
|
|
|
@property
|
|
def budget_paused(self) -> bool:
|
|
return self._budget_paused
|
|
|
|
def set_budget_extender(self, extend: Callable[[], None]) -> None:
|
|
self._extend_budget = extend
|
|
|
|
async def pause_for_budget(self, agent_id: str) -> None:
|
|
async with self._lock:
|
|
self._budget_paused = True
|
|
await self.set_status(agent_id, "budget_paused")
|
|
|
|
async def resume_from_budget_pause(self, *, exclude: str | None = None) -> None:
|
|
async with self._lock:
|
|
if not self._budget_paused:
|
|
return
|
|
self._budget_paused = False
|
|
paused = [aid for aid, status in self.statuses.items() if status == "budget_paused"]
|
|
if self._extend_budget is not None:
|
|
self._extend_budget()
|
|
for aid in paused:
|
|
await self.set_status(aid, "waiting")
|
|
if aid != exclude:
|
|
await self.send(
|
|
aid,
|
|
{
|
|
"from": "system",
|
|
"type": "budget_extended",
|
|
"content": (
|
|
"[Budget] The user extended the scan budget \u2014 continue your "
|
|
"current task."
|
|
),
|
|
},
|
|
)
|
|
|
|
async def reset_budget_stops(
|
|
self,
|
|
*,
|
|
budget_stopped: bool,
|
|
reserve_stopped: bool,
|
|
budget_paused: bool = False,
|
|
) -> None:
|
|
async with self._lock:
|
|
self._budget_stopped = budget_stopped
|
|
self._reserve_stopped = reserve_stopped
|
|
if not budget_paused:
|
|
self._budget_paused = False
|
|
for aid, status in self.statuses.items():
|
|
if status == "budget_paused":
|
|
self.statuses[aid] = "waiting"
|
|
await self._maybe_snapshot()
|
|
|
|
async def claim_reserve_notification(self) -> str | None:
|
|
async with self._lock:
|
|
if self._reserve_stopped:
|
|
return None
|
|
self._reserve_stopped = True
|
|
for runtime in self.runtimes.values():
|
|
runtime.wake.set()
|
|
return next((aid for aid, parent in self.parent_of.items() if parent is None), None)
|
|
|
|
async def register(
|
|
self,
|
|
agent_id: str,
|
|
name: str,
|
|
parent_id: str | None,
|
|
*,
|
|
task: str | None = None,
|
|
skills: list[str] | None = None,
|
|
) -> None:
|
|
async with self._lock:
|
|
self.statuses[agent_id] = "running"
|
|
self.parent_of[agent_id] = parent_id
|
|
self.names[agent_id] = name
|
|
self.pending_counts.setdefault(agent_id, 0)
|
|
self.metadata[agent_id] = {
|
|
"task": task or "",
|
|
"skills": list(skills or []),
|
|
}
|
|
self.runtimes.setdefault(agent_id, AgentRuntime())
|
|
logger.info("agent.register %s (%s) parent=%s", agent_id, name, parent_id or "-")
|
|
await self._maybe_snapshot()
|
|
|
|
async def attach_runtime(
|
|
self,
|
|
agent_id: str,
|
|
*,
|
|
session: Session | None = None,
|
|
task: asyncio.Task[Any] | None = None,
|
|
interrupt_on_message: bool | None = None,
|
|
) -> None:
|
|
async with self._lock:
|
|
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
|
if session is not None:
|
|
runtime.session = session
|
|
if task is not None:
|
|
runtime.task = task
|
|
if interrupt_on_message is not None:
|
|
runtime.interrupt_on_message = interrupt_on_message
|
|
|
|
async def mark_running(self, agent_id: str) -> None:
|
|
async with self._lock:
|
|
if agent_id in self.statuses:
|
|
self.statuses[agent_id] = "running"
|
|
self.errors.pop(agent_id, None)
|
|
self.wait_kinds.pop(agent_id, None)
|
|
self.runtimes.setdefault(agent_id, AgentRuntime()).user_wake_required = False
|
|
self._parent_notified.discard(agent_id)
|
|
await self._maybe_snapshot()
|
|
|
|
async def park_waiting(self, agent_id: str, *, wait_kind: WaitKind) -> None:
|
|
"""Park an agent, recording what it is waiting on so the driver can time it."""
|
|
async with self._lock:
|
|
if agent_id in self.statuses:
|
|
self.wait_kinds[agent_id] = wait_kind
|
|
await self.set_status(agent_id, "waiting")
|
|
|
|
async def wait_kind_of(self, agent_id: str) -> WaitKind | None:
|
|
async with self._lock:
|
|
return self.wait_kinds.get(agent_id)
|
|
|
|
async def record_recovery(self, agent_id: str) -> int:
|
|
"""Count a turn that ended without a lifecycle tool call; return the new total.
|
|
|
|
Persisted so a resumed agent cannot earn a fresh nudge budget on every
|
|
auto-resume and loop forever.
|
|
"""
|
|
async with self._lock:
|
|
count = self.recovery_counts.get(agent_id, 0) + 1
|
|
self.recovery_counts[agent_id] = count
|
|
await self._maybe_snapshot()
|
|
return count
|
|
|
|
async def reset_recovery(self, agent_id: str) -> None:
|
|
"""Clear the nudge budget after real progress (new message or a lifecycle tool)."""
|
|
async with self._lock:
|
|
if self.recovery_counts.pop(agent_id, None) is None:
|
|
return
|
|
await self._maybe_snapshot()
|
|
|
|
async def record_idle_resume(self, agent_id: str) -> int:
|
|
"""Count an auto-resume that no message triggered; return the new total.
|
|
|
|
An agent that parks again after every auto-resume would otherwise burn a
|
|
model turn per timeout for the rest of the scan.
|
|
"""
|
|
async with self._lock:
|
|
count = self.idle_resume_counts.get(agent_id, 0) + 1
|
|
self.idle_resume_counts[agent_id] = count
|
|
await self._maybe_snapshot()
|
|
return count
|
|
|
|
async def reset_idle_resumes(self, agent_id: str) -> None:
|
|
async with self._lock:
|
|
if self.idle_resume_counts.pop(agent_id, None) is None:
|
|
return
|
|
await self._maybe_snapshot()
|
|
|
|
async def set_status(
|
|
self, agent_id: str, status: Status | str, *, error: str | None = None
|
|
) -> None:
|
|
async with self._lock:
|
|
if agent_id not in self.statuses:
|
|
return
|
|
self.statuses[agent_id] = status # type: ignore[assignment]
|
|
if error is not None:
|
|
self.errors[agent_id] = error
|
|
elif status == "running":
|
|
self.errors.pop(agent_id, None)
|
|
if status == "running":
|
|
# Running again means a fresh stint that owes its parent its own notice.
|
|
self._parent_notified.discard(agent_id)
|
|
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
|
runtime.user_wake_required = status in {"failed", "crashed"}
|
|
runtime.wake.set()
|
|
logger.info("agent.status %s=%s", agent_id, status)
|
|
await self._maybe_snapshot()
|
|
|
|
async def claim_parent_notice(self, agent_id: str) -> bool:
|
|
"""Reserve the one notice a child owes its parent when it stops running.
|
|
|
|
A completion report and a terminal notice carry the same information, so
|
|
whichever comes first claims the slot and the other is skipped.
|
|
"""
|
|
async with self._lock:
|
|
if agent_id in self._parent_notified:
|
|
return False
|
|
self._parent_notified.add(agent_id)
|
|
return True
|
|
|
|
async def send(
|
|
self, target_agent_id: str, message: dict[str, Any], *, interrupt: bool = True
|
|
) -> bool:
|
|
"""Queue a user/peer message in the target's mailbox and wake it."""
|
|
from_user = message.get("from") == "user"
|
|
if from_user and self._budget_paused:
|
|
await self.resume_from_budget_pause(exclude=target_agent_id)
|
|
async with self._lock:
|
|
if target_agent_id not in self.statuses:
|
|
logger.debug("agent.send dropped unknown target=%s", target_agent_id)
|
|
return False
|
|
runtime = self.runtimes.setdefault(target_agent_id, AgentRuntime())
|
|
runtime.mailbox.append(dict(message))
|
|
self.pending_counts[target_agent_id] = self.pending_counts.get(target_agent_id, 0) + 1
|
|
if from_user:
|
|
runtime.user_wake_required = False
|
|
runtime.wake.set()
|
|
stream = runtime.stream
|
|
interrupt_on_message = runtime.interrupt_on_message
|
|
if stream is not None and interrupt and interrupt_on_message:
|
|
stream.cancel(mode="immediate")
|
|
await self._maybe_snapshot()
|
|
return True
|
|
|
|
async def wait_for_message(self, agent_id: str, *, timeout: float | None = None) -> bool:
|
|
"""Wait until a message is ready for ``agent_id``; False on ``timeout``."""
|
|
while True:
|
|
async with self._lock:
|
|
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
|
reserve_exit = self._reserve_stopped and self.parent_of.get(agent_id) is not None
|
|
pending_ready = (
|
|
self.pending_counts.get(agent_id, 0) > 0 and not runtime.user_wake_required
|
|
)
|
|
if self._budget_stopped or reserve_exit or pending_ready:
|
|
return True
|
|
wake = runtime.wake
|
|
wake.clear()
|
|
if timeout is None:
|
|
await wake.wait()
|
|
else:
|
|
try:
|
|
await asyncio.wait_for(wake.wait(), timeout)
|
|
except TimeoutError:
|
|
return False
|
|
|
|
async def consume_pending(
|
|
self,
|
|
agent_id: str,
|
|
*,
|
|
include_items: bool = False,
|
|
) -> tuple[int, list[Any]]:
|
|
"""Drain the agent's mailbox into its own SDK session."""
|
|
async with self._lock:
|
|
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
|
queued = list(runtime.mailbox)
|
|
runtime.mailbox.clear()
|
|
count = max(self.pending_counts.get(agent_id, 0), len(queued))
|
|
self.pending_counts[agent_id] = 0
|
|
session = runtime.session
|
|
if count <= 0:
|
|
return 0, []
|
|
items = [self._message_to_session_item(m) for m in queued]
|
|
if items:
|
|
if session is None:
|
|
logger.warning(
|
|
"agent %s has no SDK session attached; %d queued messages were not persisted",
|
|
agent_id,
|
|
len(items),
|
|
)
|
|
else:
|
|
try:
|
|
async with session_write_lock(session):
|
|
await session.add_items(items)
|
|
except Exception:
|
|
logger.exception(
|
|
"failed to append %d queued messages to the session of %s",
|
|
len(items),
|
|
agent_id,
|
|
)
|
|
await self._maybe_snapshot()
|
|
if not include_items:
|
|
return count, []
|
|
return count, items
|
|
|
|
async def request_stop(self, agent_id: str) -> None:
|
|
async with self._lock:
|
|
if agent_id not in self.statuses:
|
|
return
|
|
self.statuses[agent_id] = "stopped"
|
|
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
|
runtime.wake.set()
|
|
stream = runtime.stream
|
|
if stream is not None:
|
|
stream.cancel(mode="after_turn")
|
|
await self._maybe_snapshot()
|
|
|
|
async def cancel_descendants(self, agent_id: str) -> None:
|
|
tasks = []
|
|
async with self._lock:
|
|
for aid in reversed(self._subtree_order_locked(agent_id)):
|
|
task = self.runtimes.get(aid, AgentRuntime()).task
|
|
if task is not None and not task.done():
|
|
tasks.append(task)
|
|
for task in tasks:
|
|
task.cancel()
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
async def cancel_descendants_graceful(self, agent_id: str) -> list[str]:
|
|
"""Stop a subtree leaves-first and report which agents were stopped."""
|
|
async with self._lock:
|
|
order = self._subtree_order_locked(agent_id)
|
|
stopped = list(reversed(order))
|
|
for aid in stopped:
|
|
await self.request_stop(aid)
|
|
await self._maybe_snapshot()
|
|
return stopped
|
|
|
|
async def attach_stream(
|
|
self,
|
|
agent_id: str,
|
|
stream: Any,
|
|
) -> None:
|
|
async with self._lock:
|
|
self.runtimes.setdefault(agent_id, AgentRuntime()).stream = stream
|
|
|
|
async def detach_stream(
|
|
self,
|
|
agent_id: str,
|
|
stream: Any,
|
|
) -> None:
|
|
async with self._lock:
|
|
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
|
if runtime.stream is stream:
|
|
runtime.stream = None
|
|
|
|
async def active_agents_except(self, agent_id: str) -> list[dict[str, Any]]:
|
|
async with self._lock:
|
|
return [
|
|
{
|
|
"agent_id": aid,
|
|
"name": self.names.get(aid, aid),
|
|
"status": status,
|
|
"parent_id": self.parent_of.get(aid),
|
|
}
|
|
for aid, status in self.statuses.items()
|
|
if aid != agent_id and status in {"running", "waiting"}
|
|
]
|
|
|
|
async def graph_snapshot(
|
|
self,
|
|
) -> tuple[dict[str, str | None], dict[str, Status], dict[str, str], dict[str, str]]:
|
|
async with self._lock:
|
|
return (
|
|
dict(self.parent_of),
|
|
dict(self.statuses),
|
|
dict(self.names),
|
|
dict(self.errors),
|
|
)
|
|
|
|
def _message_to_session_item(self, message: dict[str, Any]) -> TResponseInputItem:
|
|
sender = str(message.get("from", "unknown"))
|
|
content = str(message.get("content", ""))
|
|
if sender == "user":
|
|
return cast("TResponseInputItem", {"role": "user", "content": content})
|
|
sender_name = self.names.get(sender, sender)
|
|
msg_type = message.get("type", "information")
|
|
priority = message.get("priority", "normal")
|
|
return cast(
|
|
"TResponseInputItem",
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"[Message from {sender_name} ({sender}) | type={msg_type} "
|
|
f"| priority={priority}]\n{content}"
|
|
),
|
|
},
|
|
)
|
|
|
|
def _subtree_order_locked(self, agent_id: str) -> list[str]:
|
|
queue = [agent_id]
|
|
order: list[str] = []
|
|
while queue:
|
|
aid = queue.pop()
|
|
order.append(aid)
|
|
queue.extend(child for child, parent in self.parent_of.items() if parent == aid)
|
|
return order
|
|
|
|
async def snapshot(self) -> dict[str, Any]:
|
|
async with self._lock:
|
|
return {
|
|
"statuses": dict(self.statuses),
|
|
"parent_of": dict(self.parent_of),
|
|
"names": dict(self.names),
|
|
"metadata": {aid: dict(md) for aid, md in self.metadata.items()},
|
|
"pending_counts": dict(self.pending_counts),
|
|
"recovery_counts": dict(self.recovery_counts),
|
|
"idle_resume_counts": dict(self.idle_resume_counts),
|
|
"wait_kinds": dict(self.wait_kinds),
|
|
"mailboxes": {
|
|
aid: [dict(m) for m in runtime.mailbox]
|
|
for aid, runtime in self.runtimes.items()
|
|
if runtime.mailbox
|
|
},
|
|
"errors": dict(self.errors),
|
|
"budget_stopped": self._budget_stopped,
|
|
"reserve_stopped": self._reserve_stopped,
|
|
"budget_paused": self._budget_paused,
|
|
}
|
|
|
|
async def restore(self, snap: dict[str, Any]) -> None:
|
|
async with self._lock:
|
|
self.statuses = dict(snap.get("statuses", {}))
|
|
self.parent_of = dict(snap.get("parent_of", {}))
|
|
self.names = dict(snap.get("names", {}))
|
|
self.metadata = {aid: dict(md) for aid, md in snap.get("metadata", {}).items()}
|
|
self.pending_counts = dict(snap.get("pending_counts", {}))
|
|
self.errors = dict(snap.get("errors", {}))
|
|
self.recovery_counts = dict(snap.get("recovery_counts", {}))
|
|
self.idle_resume_counts = dict(snap.get("idle_resume_counts", {}))
|
|
self.wait_kinds = dict(snap.get("wait_kinds", {}))
|
|
mailboxes = snap.get("mailboxes", {})
|
|
if isinstance(mailboxes, dict):
|
|
for aid, msgs in mailboxes.items():
|
|
if isinstance(msgs, list):
|
|
runtime = self.runtimes.setdefault(aid, AgentRuntime())
|
|
runtime.mailbox = [dict(m) for m in msgs if isinstance(m, dict)]
|
|
self._budget_stopped = bool(snap.get("budget_stopped", False))
|
|
self._reserve_stopped = bool(snap.get("reserve_stopped", False))
|
|
self._budget_paused = bool(snap.get("budget_paused", False))
|
|
for aid in self.statuses:
|
|
self.runtimes.setdefault(aid, AgentRuntime())
|
|
|
|
async def _maybe_snapshot(self) -> None:
|
|
path = self._snapshot_path
|
|
if path is None:
|
|
return
|
|
try:
|
|
data = await self.snapshot()
|
|
payload = json.dumps(data, ensure_ascii=False, default=str)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w",
|
|
encoding="utf-8",
|
|
dir=str(path.parent),
|
|
prefix=f".{path.name}.",
|
|
suffix=".tmp",
|
|
delete=False,
|
|
) as tmp:
|
|
tmp.write(payload)
|
|
tmp_path = Path(tmp.name)
|
|
tmp_path.replace(path)
|
|
except Exception:
|
|
logger.exception("coordinator snapshot to %s failed", path)
|
|
|
|
|
|
def coordinator_from_context(ctx: dict[str, Any]) -> AgentCoordinator | None:
|
|
coordinator = ctx.get("coordinator")
|
|
return coordinator if isinstance(coordinator, AgentCoordinator) else None
|