|
32 | 32 | from dataclasses import dataclass, field |
33 | 33 |
|
34 | 34 | from libtmux._internal.query_list import QueryList |
| 35 | +from libtmux.experimental.agents.state import AgentState |
35 | 36 | from libtmux.experimental.engines.base import TmuxEngine |
36 | 37 | from libtmux.experimental.ops import ( |
37 | 38 | ClearHistory, |
|
53 | 54 |
|
54 | 55 | from typing_extensions import Self |
55 | 56 |
|
| 57 | + from libtmux.experimental.agents.monitor import AgentMonitor |
| 58 | + from libtmux.experimental.agents.state import Agent |
56 | 59 | from libtmux.experimental.models.snapshots import PaneSnapshot |
57 | 60 | from libtmux.experimental.ops import Planner, PlanResult |
58 | 61 | from libtmux.experimental.ops._types import SlotRef, Target |
59 | 62 |
|
60 | 63 | #: A source of pane snapshots: an engine to read from, or pre-taken snapshots. |
61 | 64 | PaneSource = t.Union["TmuxEngine", "Sequence[PaneSnapshot]"] |
| 65 | +#: A source of agent records: a monitor to read its store, or pre-taken records. |
| 66 | +AgentSource = t.Union["AgentMonitor", "Sequence[Agent]"] |
62 | 67 | MappedT = t.TypeVar("MappedT") |
| 68 | +KeyT = t.TypeVar("KeyT") |
| 69 | + |
| 70 | +#: Default attention ladder for agent rollups (higher value = more urgent). The |
| 71 | +#: ordering is a *documented default* a caller can override per call: surveyed |
| 72 | +#: orchestrators disagree on the exact weighting, so it is policy, not a rule. |
| 73 | +ATTENTION: dict[AgentState, int] = { |
| 74 | + AgentState.AWAITING_INPUT: 5, |
| 75 | + AgentState.DONE: 4, |
| 76 | + AgentState.IDLE: 3, |
| 77 | + AgentState.RUNNING: 2, |
| 78 | + AgentState.UNKNOWN: 1, |
| 79 | + AgentState.EXITED: 0, |
| 80 | +} |
63 | 81 |
|
64 | 82 |
|
65 | 83 | def _snapshot_panes(source: PaneSource) -> tuple[PaneSnapshot, ...]: |
@@ -353,3 +371,152 @@ def panes() -> PaneQuery: |
353 | 371 | PaneQuery(lookups={'active': True}, order='pane_index', limit_count=1) |
354 | 372 | """ |
355 | 373 | return PaneQuery() |
| 374 | + |
| 375 | + |
| 376 | +def _query_agents(source: AgentSource) -> tuple[Agent, ...]: |
| 377 | + """Resolve *source* into agent records (read a monitor's store, or pass through). |
| 378 | +
|
| 379 | + A monitor is detected by its ``agents`` snapshot property (zero tmux calls -- |
| 380 | + the store is already populated by the monitor's own drain); any other value |
| 381 | + is taken as a pure sequence of :class:`~..agents.state.Agent` records. |
| 382 | + """ |
| 383 | + store_agents = getattr(source, "agents", None) |
| 384 | + if store_agents is not None: |
| 385 | + return tuple(store_agents) |
| 386 | + return tuple(t.cast("Sequence[Agent]", source)) |
| 387 | + |
| 388 | + |
| 389 | +@dataclass(frozen=True) |
| 390 | +class AgentQuery: |
| 391 | + """An immutable, chainable query over agents (the agent twin of PaneQuery). |
| 392 | +
|
| 393 | + Resolves against an :data:`AgentSource` -- an |
| 394 | + :class:`~..agents.monitor.AgentMonitor` (read straight from its in-process |
| 395 | + store, **zero tmux calls**) or a pure sequence of |
| 396 | + :class:`~..agents.state.Agent` records. Each method returns a new query; |
| 397 | + :meth:`all` / :meth:`first` resolve it. |
| 398 | +
|
| 399 | + Examples |
| 400 | + -------- |
| 401 | + >>> from libtmux.experimental.agents.state import Agent, AgentState |
| 402 | + >>> rows = [ |
| 403 | + ... Agent(pane_id="%1", key="%1", name="claude", |
| 404 | + ... state=AgentState.AWAITING_INPUT, since=0.0, source="option", |
| 405 | + ... pid=None, alive=True), |
| 406 | + ... Agent(pane_id="%2", key="%2", name="codex", |
| 407 | + ... state=AgentState.RUNNING, since=0.0, source="option", |
| 408 | + ... pid=None, alive=True), |
| 409 | + ... ] |
| 410 | + >>> agents().filter(state=AgentState.AWAITING_INPUT).map( |
| 411 | + ... lambda a: a.pane_id).all(rows) |
| 412 | + ('%1',) |
| 413 | + """ |
| 414 | + |
| 415 | + lookups: Mapping[str, t.Any] = field(default_factory=dict) |
| 416 | + order: str | None = None |
| 417 | + limit_count: int | None = None |
| 418 | + |
| 419 | + def filter(self, **lookups: t.Any) -> AgentQuery: |
| 420 | + """Narrow by QueryList lookups (e.g. ``state=AgentState.IDLE``, ``name=``).""" |
| 421 | + return dataclasses.replace(self, lookups={**self.lookups, **lookups}) |
| 422 | + |
| 423 | + def order_by(self, field_name: str) -> AgentQuery: |
| 424 | + """Sort the results by an Agent attribute (missing values last).""" |
| 425 | + return dataclasses.replace(self, order=field_name) |
| 426 | + |
| 427 | + def limit(self, count: int) -> AgentQuery: |
| 428 | + """Keep only the first *count* results.""" |
| 429 | + return dataclasses.replace(self, limit_count=count) |
| 430 | + |
| 431 | + def all(self, source: AgentSource) -> tuple[Agent, ...]: |
| 432 | + """Resolve the query against *source* and return the matched agents.""" |
| 433 | + rows: t.Any = QueryList(_query_agents(source)) |
| 434 | + if self.lookups: |
| 435 | + rows = rows.filter(**self.lookups) |
| 436 | + rows = list(rows) |
| 437 | + if self.order is not None: |
| 438 | + rows.sort(key=lambda agent: _order_key(agent, self.order)) |
| 439 | + if self.limit_count is not None: |
| 440 | + rows = rows[: self.limit_count] |
| 441 | + return tuple(rows) |
| 442 | + |
| 443 | + def first(self, source: AgentSource) -> Agent | None: |
| 444 | + """Return the first matched agent, or ``None`` when none match.""" |
| 445 | + rows = self.all(source) |
| 446 | + return rows[0] if rows else None |
| 447 | + |
| 448 | + def map(self, fn: Callable[[Agent], MappedT]) -> MappedAgentQuery[MappedT]: |
| 449 | + """Project each matched agent through *fn* (a pure read projection).""" |
| 450 | + return MappedAgentQuery(self, fn) |
| 451 | + |
| 452 | + def most_urgent( |
| 453 | + self, |
| 454 | + source: AgentSource, |
| 455 | + *, |
| 456 | + priority: Mapping[AgentState, int] = ATTENTION, |
| 457 | + ) -> Agent | None: |
| 458 | + """Return the matched agent whose state ranks highest in *priority*. |
| 459 | +
|
| 460 | + The "jump to the agent that needs me" primitive (ties keep input order); |
| 461 | + ``None`` when nothing matches. *priority* defaults to :data:`ATTENTION`. |
| 462 | + """ |
| 463 | + rows = self.all(source) |
| 464 | + if not rows: |
| 465 | + return None |
| 466 | + return max(rows, key=lambda agent: priority.get(agent.state, -1)) |
| 467 | + |
| 468 | + def rollup( |
| 469 | + self, |
| 470 | + source: AgentSource, |
| 471 | + *, |
| 472 | + key: Callable[[Agent], KeyT], |
| 473 | + priority: Mapping[AgentState, int] = ATTENTION, |
| 474 | + ) -> dict[KeyT, AgentState]: |
| 475 | + """Collapse each ``key(agent)`` group to its most-urgent state. |
| 476 | +
|
| 477 | + The fleet "who needs me" read model: group the matched agents by *key* |
| 478 | + (e.g. ``lambda a: a.name``) and report, per group, the state with the |
| 479 | + highest *priority*. *priority* defaults to :data:`ATTENTION` and is |
| 480 | + overridable -- the weighting is policy, not a fixed rule. |
| 481 | + """ |
| 482 | + best_rank: dict[KeyT, int] = {} |
| 483 | + out: dict[KeyT, AgentState] = {} |
| 484 | + for agent in self.all(source): |
| 485 | + group = key(agent) |
| 486 | + rank = priority.get(agent.state, -1) |
| 487 | + if group not in best_rank or rank > best_rank[group]: |
| 488 | + best_rank[group] = rank |
| 489 | + out[group] = agent.state |
| 490 | + return out |
| 491 | + |
| 492 | + |
| 493 | +@dataclass(frozen=True) |
| 494 | +class MappedAgentQuery(t.Generic[MappedT]): |
| 495 | + """An :class:`AgentQuery` whose rows are projected through a function.""" |
| 496 | + |
| 497 | + query: AgentQuery |
| 498 | + fn: Callable[[Agent], MappedT] |
| 499 | + |
| 500 | + def all(self, source: AgentSource) -> tuple[MappedT, ...]: |
| 501 | + """Resolve and project every matched agent.""" |
| 502 | + return tuple(self.fn(agent) for agent in self.query.all(source)) |
| 503 | + |
| 504 | + def first(self, source: AgentSource) -> MappedT | None: |
| 505 | + """Resolve and project the first matched agent, or ``None``.""" |
| 506 | + first = self.query.first(source) |
| 507 | + return self.fn(first) if first is not None else None |
| 508 | + |
| 509 | + |
| 510 | +def agents() -> AgentQuery: |
| 511 | + """Start a query over tracked coding agents (the agent twin of :func:`panes`). |
| 512 | +
|
| 513 | + Resolve it against an :class:`~..agents.monitor.AgentMonitor` (zero tmux |
| 514 | + calls -- the monitor's store is already live) or a pure sequence of |
| 515 | + :class:`~..agents.state.Agent` records. |
| 516 | +
|
| 517 | + Examples |
| 518 | + -------- |
| 519 | + >>> agents().filter(name="claude").limit(1) |
| 520 | + AgentQuery(lookups={'name': 'claude'}, order=None, limit_count=1) |
| 521 | + """ |
| 522 | + return AgentQuery() |
0 commit comments