458 lines
17 KiB
Python
458 lines
17 KiB
Python
"""Live-subagent registry + model-facing control plane (list/steer/stop) for delegate_task.
|
|
|
|
Split out of ``tools/delegate_tool.py``; every moved name is re-imported there, so
|
|
``tools.delegate_tool.<name>`` keeps resolving (and monkeypatching) as before.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import json
|
|
import threading
|
|
import time
|
|
from typing import Any, Dict, List, Optional
|
|
from agent.interrupt_compat import request_hard_interrupt
|
|
from tools.registry import tool_error
|
|
|
|
# Log-record parity with the origin module.
|
|
logger = logging.getLogger("tools.delegate_tool")
|
|
|
|
_spawn_pause_lock = threading.Lock()
|
|
|
|
_spawn_paused: bool = False
|
|
|
|
_active_subagents_lock = threading.Lock()
|
|
|
|
# subagent_id -> mutable record tracking the live child agent. Stays only
|
|
# for the lifetime of the run; _run_single_child is the owner.
|
|
_active_subagents: Dict[str, Dict[str, Any]] = {}
|
|
|
|
# subagent_id -> {goal, delegation_id, parent_session_id} retained AFTER the
|
|
# child finishes (bounded FIFO). Child-started background processes routinely
|
|
# outlive the child itself (its npm ci with notify_on_complete=true finishes
|
|
# after the child's summary was delivered); their completion notifications
|
|
# reach the parent conversation via the shared completion_queue and need
|
|
# delegation attribution even though the live registry entry is gone.
|
|
_RECENT_SUBAGENTS_CAP = 200
|
|
|
|
_recent_subagents: Dict[str, Dict[str, Any]] = {}
|
|
|
|
|
|
def get_subagent_attribution(task_id: Optional[str]) -> Optional[Dict[str, Any]]:
|
|
"""Resolve a process task_id to its originating delegation, if any.
|
|
|
|
Children run their terminal sessions under ``task_id == subagent_id``
|
|
(see _run_single_child's child_task_id), so a background process spawned
|
|
by a subagent carries that id in ``ProcessSession.task_id``. Returns
|
|
``{subagent_id, goal, delegation_id}`` for live AND recently-finished
|
|
children, or None when the task_id is not a known subagent.
|
|
"""
|
|
if not task_id or not isinstance(task_id, str):
|
|
return None
|
|
with _active_subagents_lock:
|
|
record = _active_subagents.get(task_id)
|
|
if record is not None:
|
|
return {
|
|
"subagent_id": task_id,
|
|
"goal": record.get("goal"),
|
|
"delegation_id": record.get("delegation_id"),
|
|
}
|
|
retained = _recent_subagents.get(task_id)
|
|
if retained is not None:
|
|
return {
|
|
"subagent_id": task_id,
|
|
"goal": retained.get("goal"),
|
|
"delegation_id": retained.get("delegation_id"),
|
|
}
|
|
return None
|
|
|
|
|
|
def set_spawn_paused(paused: bool) -> bool:
|
|
"""Globally block/unblock new delegate_task spawns.
|
|
|
|
Active children keep running; only NEW calls to delegate_task fail fast
|
|
with a "spawning paused" error until unblocked. Returns the new state.
|
|
"""
|
|
global _spawn_paused
|
|
with _spawn_pause_lock:
|
|
_spawn_paused = bool(paused)
|
|
return _spawn_paused
|
|
|
|
|
|
def is_spawn_paused() -> bool:
|
|
with _spawn_pause_lock:
|
|
return _spawn_paused
|
|
|
|
|
|
def _register_subagent(record: Dict[str, Any]) -> None:
|
|
sid = record.get("subagent_id")
|
|
if not sid:
|
|
return
|
|
record.setdefault("accepting_steer", True)
|
|
with _active_subagents_lock:
|
|
_active_subagents[sid] = record
|
|
|
|
|
|
def _retain_recent_subagent(record: Dict[str, Any]) -> None:
|
|
"""Keep a bounded attribution stub after a child finishes (lock held)."""
|
|
sid = record.get("subagent_id")
|
|
if not sid:
|
|
return
|
|
_recent_subagents[sid] = {
|
|
"goal": record.get("goal"),
|
|
"delegation_id": record.get("delegation_id"),
|
|
"owner_agent_session_id": record.get("owner_agent_session_id"),
|
|
}
|
|
while len(_recent_subagents) > _RECENT_SUBAGENTS_CAP:
|
|
_recent_subagents.pop(next(iter(_recent_subagents)), None)
|
|
|
|
|
|
def _unregister_subagent(subagent_id: str, *, agent: Any = None) -> None:
|
|
with _active_subagents_lock:
|
|
record = _active_subagents.get(subagent_id)
|
|
if record is not None and (agent is None or record.get("agent") is agent):
|
|
_active_subagents.pop(subagent_id, None)
|
|
_retain_recent_subagent(record)
|
|
|
|
|
|
def _close_subagent_steering(subagent_id: str, agent: Any) -> Optional[str]:
|
|
"""Atomically close steer acceptance and drain its final durable artifact.
|
|
|
|
``steer_subagent`` holds the same registry lock through ``agent.steer``.
|
|
Therefore either acceptance wins and this drain sees its exact text, or
|
|
closure wins and the caller is rejected. Exact agent identity prevents a
|
|
finishing child with a recycled public id from closing its replacement.
|
|
"""
|
|
with _active_subagents_lock:
|
|
record = _active_subagents.get(subagent_id)
|
|
if record is None or record.get("agent") is not agent:
|
|
return None
|
|
record["accepting_steer"] = False
|
|
drain = getattr(agent, "_drain_pending_steer", None)
|
|
if not callable(drain):
|
|
return None
|
|
try:
|
|
pending = drain()
|
|
except Exception as exc:
|
|
logger.debug("final steer drain for %s failed: %s", subagent_id, exc)
|
|
return None
|
|
return pending if isinstance(pending, str) and pending.strip() else None
|
|
|
|
|
|
def interrupt_subagent(subagent_id: str) -> bool:
|
|
"""Request that a single running subagent stop at its next iteration boundary.
|
|
|
|
Does not hard-kill the worker thread (Python can't); sets the child's
|
|
interrupt flag which propagates to in-flight tools and recurses into
|
|
grandchildren via AIAgent.interrupt(). Returns True if a matching
|
|
subagent was found.
|
|
"""
|
|
with _active_subagents_lock:
|
|
record = _active_subagents.get(subagent_id)
|
|
if not record:
|
|
return False
|
|
agent = record.get("agent")
|
|
if agent is None:
|
|
return False
|
|
try:
|
|
if not request_hard_interrupt(agent, f"Interrupted via TUI ({subagent_id})"):
|
|
return False
|
|
except Exception as exc:
|
|
logger.debug("interrupt_subagent(%s) failed: %s", subagent_id, exc)
|
|
return False
|
|
return True
|
|
|
|
|
|
def steer_subagent(
|
|
subagent_id: str,
|
|
text: str,
|
|
*,
|
|
owner_session_id: Optional[str] = None,
|
|
owner_transport: Any = None,
|
|
owner_session_record: Any = None,
|
|
) -> bool:
|
|
"""Queue steering text into a single running subagent without stopping it.
|
|
|
|
The redirection-side mirror of interrupt_subagent(): resolves the live
|
|
child in the registry and calls AIAgent.steer(), which appends the text
|
|
to the child's last tool result at its next iteration boundary — the
|
|
current tool call is never cut. Returns True if a matching subagent
|
|
QUEUED the text while the child was still accepting work; False for an
|
|
unknown/closed id, an ownership mismatch, a record with no live agent, or
|
|
empty text. ``owner_session_id=None`` deliberately preserves the internal
|
|
in-process helper contract; gateway callers must pass exact authority.
|
|
|
|
Acceptance and completion are linearized by the registry lock. If
|
|
acceptance wins but no delivery boundary remains, ``_run_single_child``
|
|
drains the exact text into the completion entry as ``missed_steer``.
|
|
"""
|
|
if not text or not text.strip():
|
|
return False
|
|
with _active_subagents_lock:
|
|
record = _active_subagents.get(subagent_id)
|
|
if not record or not record.get("accepting_steer", False):
|
|
return False
|
|
if owner_session_id is not None and (
|
|
record.get("owner_session_id") != owner_session_id
|
|
or owner_transport is None
|
|
or record.get("owner_transport") is not owner_transport
|
|
or owner_session_record is None
|
|
or record.get("owner_session_record") is not owner_session_record
|
|
):
|
|
return False
|
|
agent = record.get("agent")
|
|
if agent is None:
|
|
return False
|
|
try:
|
|
return bool(agent.steer(text))
|
|
except Exception as exc:
|
|
logger.debug("steer_subagent(%s) failed: %s", subagent_id, exc)
|
|
return False
|
|
|
|
|
|
def _capture_gateway_steer_authority(
|
|
owner_session_id: Optional[str],
|
|
) -> tuple[Any, Any]:
|
|
"""Capture exact request transport + live session generation, if any.
|
|
|
|
This is intentionally an in-process bridge, not a serializable capability.
|
|
Non-gateway hosts (including the CLI helper path) receive ``(None, None)``.
|
|
"""
|
|
if not owner_session_id:
|
|
return None, None
|
|
try:
|
|
from tui_gateway.server import _current_session_steer_authority
|
|
|
|
return _current_session_steer_authority(owner_session_id)
|
|
except Exception:
|
|
return None, None
|
|
|
|
|
|
def list_active_subagents() -> List[Dict[str, Any]]:
|
|
"""Snapshot of the currently running subagent tree.
|
|
|
|
Each record: {subagent_id, parent_id, depth, goal, model, started_at,
|
|
tool_count, status}. Safe to call from any thread — returns a copy.
|
|
"""
|
|
with _active_subagents_lock:
|
|
return [
|
|
{
|
|
k: v
|
|
for k, v in r.items()
|
|
if k
|
|
not in {
|
|
"agent",
|
|
"owner_session_id",
|
|
"owner_transport",
|
|
"owner_session_record",
|
|
"accepting_steer",
|
|
}
|
|
}
|
|
for r in _active_subagents.values()
|
|
]
|
|
|
|
|
|
def _is_descendant_of(child_agent: Any, parent_agent: Any, max_hops: int = 8) -> bool:
|
|
"""True when *child_agent* sits below *parent_agent* in the spawn tree.
|
|
|
|
Walks the ``_delegate_parent_ref`` weakref chain stamped at build time.
|
|
Identity comparison only — a parent may steer/stop its own children and
|
|
grandchildren, never a sibling tree owned by another conversation.
|
|
"""
|
|
if child_agent is None or parent_agent is None:
|
|
return False
|
|
cur = child_agent
|
|
for _ in range(max_hops):
|
|
ref = getattr(cur, "_delegate_parent_ref", None)
|
|
ancestor = ref() if callable(ref) else None
|
|
if ancestor is None:
|
|
return False
|
|
if ancestor is parent_agent:
|
|
return True
|
|
cur = ancestor
|
|
return False
|
|
|
|
|
|
# Model-facing control actions accepted by delegate_task(action=...).
|
|
# "spawn" (or omitted) keeps the historical spawn semantics.
|
|
_CONTROL_ACTIONS = frozenset({"list", "steer", "stop"})
|
|
|
|
|
|
def _resolve_session_lineage(session_id: Optional[str], parent_agent: Any) -> str:
|
|
"""Resolve a session id to the tip of its compression lineage.
|
|
|
|
Best-effort: uses the parent's live SessionDB handle when present so a
|
|
delegation dispatched before a compression rotation still matches the
|
|
rotated parent. Returns the input unchanged when resolution fails.
|
|
"""
|
|
sid = str(session_id or "")
|
|
if not sid:
|
|
return ""
|
|
db = getattr(parent_agent, "_session_db", None)
|
|
if db is None:
|
|
return sid
|
|
try:
|
|
resolved = db.resolve_resume_session_id(sid)
|
|
return str(resolved) if resolved else sid
|
|
except Exception:
|
|
return sid
|
|
|
|
|
|
def _owns_subagent_record(record: Dict[str, Any], parent_agent: Any) -> bool:
|
|
"""True when *parent_agent*'s conversation owns this live-child record.
|
|
|
|
Two-tier check:
|
|
|
|
1. Object identity — the ``_delegate_parent_ref`` weakref chain stamped
|
|
at build time reaches *parent_agent*. Fast path for the common case
|
|
where the parent AIAgent object survives the whole run.
|
|
2. Durable conversation lineage — the child was registered with the
|
|
owning conversation's durable session id
|
|
(``owner_agent_session_id``); match it against the calling parent's
|
|
``session_id``, resolving compression-rotation lineage on both sides.
|
|
|
|
Tier 2 exists because the identity chain is BRITTLE across parent-agent
|
|
rebuilds: the CLI sets ``self.agent = None`` mid-session (route-signature
|
|
change, credential refresh, /model, MoA one-shots) and constructs a NEW
|
|
AIAgent for the next turn while the child keeps running with a weakref to
|
|
the old object. The delivery path always survived this (it routes by
|
|
durable session id); the control path must use the same durable spine or
|
|
running children go invisible/unsteerable (observed live: deleg_88454b70
|
|
/ sa-0-dc0100f4, 2026-08-17).
|
|
"""
|
|
agent = record.get("agent")
|
|
if _is_descendant_of(agent, parent_agent):
|
|
return True
|
|
owner_sid = str(record.get("owner_agent_session_id") or "")
|
|
if not owner_sid:
|
|
return False
|
|
parent_sid = str(getattr(parent_agent, "session_id", "") or "")
|
|
if not parent_sid:
|
|
return False
|
|
if owner_sid == parent_sid:
|
|
return True
|
|
# Compression rotation on either side: compare lineage tips.
|
|
return _resolve_session_lineage(owner_sid, parent_agent) in {
|
|
parent_sid,
|
|
_resolve_session_lineage(parent_sid, parent_agent),
|
|
}
|
|
|
|
|
|
def _handle_control_action(
|
|
action: str,
|
|
subagent_id: Optional[str],
|
|
message: Optional[str],
|
|
parent_agent: Any,
|
|
) -> str:
|
|
"""Synchronous control plane for delegate_task: list/steer/stop.
|
|
|
|
Runs in-turn (never backgrounded) and only over subagents descended from
|
|
*parent_agent* — the same registry the TUI overlay drives, but scoped so
|
|
a conversation can only control its own spawn tree.
|
|
"""
|
|
if action == "list":
|
|
with _active_subagents_lock:
|
|
records = list(_active_subagents.values())
|
|
entries = []
|
|
for r in records:
|
|
agent = r.get("agent")
|
|
if not _owns_subagent_record(r, parent_agent):
|
|
continue
|
|
started = r.get("started_at")
|
|
entries.append(
|
|
{
|
|
"subagent_id": r.get("subagent_id"),
|
|
"parent_id": r.get("parent_id"),
|
|
"goal": r.get("goal"),
|
|
"model": r.get("model"),
|
|
"status": r.get("status"),
|
|
"running_seconds": (
|
|
round(time.time() - started, 1)
|
|
if isinstance(started, (int, float))
|
|
else None
|
|
),
|
|
"accepting_steer": bool(r.get("accepting_steer", False)),
|
|
"live_transcript": getattr(agent, "_live_transcript_path", None),
|
|
}
|
|
)
|
|
payload: Dict[str, Any] = {
|
|
"action": "list",
|
|
"count": len(entries),
|
|
"subagents": entries,
|
|
}
|
|
if not entries:
|
|
payload["note"] = (
|
|
"No live subagents right now. Children that already finished "
|
|
"have delivered (or will deliver) their results as normal "
|
|
"completion messages — there is nothing to steer or stop."
|
|
)
|
|
return json.dumps(payload, ensure_ascii=False)
|
|
|
|
# steer / stop need a resolvable, owned target.
|
|
sid = (subagent_id or "").strip()
|
|
if not sid:
|
|
return tool_error(
|
|
f"action='{action}' requires subagent_id (from the spawn dispatch "
|
|
"response or action='list')."
|
|
)
|
|
with _active_subagents_lock:
|
|
record = _active_subagents.get(sid)
|
|
if record is None or not _owns_subagent_record(record, parent_agent):
|
|
return tool_error(
|
|
f"No live subagent '{sid}' in this conversation's spawn tree. It "
|
|
"may have already finished (its result arrives as a normal "
|
|
"completion message). Use action='list' to see live children."
|
|
)
|
|
|
|
if action == "stop":
|
|
if interrupt_subagent(sid):
|
|
return json.dumps(
|
|
{
|
|
"action": "stop",
|
|
"subagent_id": sid,
|
|
"status": "interrupt_requested",
|
|
"note": (
|
|
"The subagent stops at its next iteration boundary "
|
|
"(in-flight tool calls are asked to cancel). Its "
|
|
"partial result still re-enters the conversation as a "
|
|
"completion message — do not wait or poll."
|
|
),
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
return tool_error(
|
|
f"Could not interrupt '{sid}' — it likely finished in the last "
|
|
"moment. Its result arrives as a normal completion message."
|
|
)
|
|
|
|
if action == "steer":
|
|
text = (message or "").strip()
|
|
if not text:
|
|
return tool_error(
|
|
"action='steer' requires a non-empty 'message' describing the "
|
|
"course correction."
|
|
)
|
|
if steer_subagent(sid, text):
|
|
return json.dumps(
|
|
{
|
|
"action": "steer",
|
|
"subagent_id": sid,
|
|
"status": "queued",
|
|
"note": (
|
|
"Steering text queued. The subagent sees it appended "
|
|
"to its next tool result — the current tool call is "
|
|
"never cut. If the child finishes before a delivery "
|
|
"boundary remains, the text is reported back as "
|
|
"missed_steer in its completion entry."
|
|
),
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
return tool_error(
|
|
f"Subagent '{sid}' is no longer accepting steering (finishing or "
|
|
"already finished). Its result arrives as a normal completion "
|
|
"message; re-delegate a follow-up task if more work is needed."
|
|
)
|
|
|
|
return tool_error(f"Unknown action '{action}'. Use spawn, list, steer, or stop.")
|