diff --git a/tests/tools/test_delegate_control_actions.py b/tests/tools/test_delegate_control_actions.py index 484760a2d2..aee8fd7350 100644 --- a/tests/tools/test_delegate_control_actions.py +++ b/tests/tools/test_delegate_control_actions.py @@ -14,9 +14,11 @@ import weakref from tools.delegate_tool import ( _handle_control_action, _is_descendant_of, + _owns_subagent_record, _register_subagent, _unregister_subagent, delegate_task, + get_subagent_attribution, ) @@ -291,10 +293,245 @@ def test_empty_tasks_array_with_goal_is_single_task_not_batch_error(): # --------------------------------------------------------------------------- -# Guardrail: control actions never consume the spawn cap +# Durable ownership: registry survives parent-agent object rebuilds +# (regression for deleg_88454b70 / sa-0-dc0100f4, 2026-08-17: CLI rebuilt its +# AIAgent mid-session; running child fell out of list/steer while completion +# delivery — which routes by durable session id — still worked) # --------------------------------------------------------------------------- +class _StubParentWithSession: + def __init__(self, session_id: str, session_db=None): + self.session_id = session_id + self._session_db = session_db + + +class _StubSessionDB: + """resolve_resume_session_id lineage: old_id -> tip mapping.""" + + def __init__(self, lineage=None): + self._lineage = dict(lineage or {}) + + def resolve_resume_session_id(self, session_id): + return self._lineage.get(session_id, session_id) + + +def test_list_finds_child_after_parent_agent_rebuild(): + """A REBUILT parent object (weakref chain broken) must still see its + running child via the durable owner_agent_session_id spine.""" + old_parent = _StubParentWithSession("sess-durable-1") + child = _StubChild(old_parent) + _register( + "sid-durable-list-1", child, owner_agent_session_id="sess-durable-1" + ) + # Simulate the CLI's `self.agent = None` + rebuild: a NEW object, same + # durable conversation/session id. The weakref chain no longer reaches it. + rebuilt_parent = _StubParentWithSession("sess-durable-1") + assert _is_descendant_of(child, rebuilt_parent) is False # identity broken + try: + out = json.loads(_handle_control_action("list", None, None, rebuilt_parent)) + assert out["count"] == 1 + assert out["subagents"][0]["subagent_id"] == "sid-durable-list-1" + finally: + _unregister_subagent("sid-durable-list-1") + + +def test_steer_resolves_after_parent_agent_rebuild(): + old_parent = _StubParentWithSession("sess-durable-2") + child = _StubChild(old_parent) + _register( + "sid-durable-steer-1", child, owner_agent_session_id="sess-durable-2" + ) + rebuilt_parent = _StubParentWithSession("sess-durable-2") + try: + out = json.loads( + _handle_control_action( + "steer", "sid-durable-steer-1", "keep going", rebuilt_parent + ) + ) + assert out["status"] == "queued" + assert child.steered == ["keep going"] + finally: + _unregister_subagent("sid-durable-steer-1") + + +def test_durable_ownership_does_not_leak_to_foreign_session(): + """A DIFFERENT conversation (different session id, no identity chain) + must still be refused — the durable spine widens recovery, not access.""" + owner = _StubParentWithSession("sess-durable-3") + child = _StubChild(owner) + _register( + "sid-durable-foreign-1", child, owner_agent_session_id="sess-durable-3" + ) + intruder = _StubParentWithSession("sess-other-99") + try: + out = _handle_control_action( + "steer", "sid-durable-foreign-1", "hijack", intruder + ) + assert "No live subagent" in out + assert child.steered == [] + out2 = json.loads(_handle_control_action("list", None, None, intruder)) + assert out2["count"] == 0 + finally: + _unregister_subagent("sid-durable-foreign-1") + + +def test_durable_ownership_resolves_compression_lineage(): + """Delegation registered under a pre-compression session id must match + the rotated parent whose SessionDB lineage maps old -> new.""" + db = _StubSessionDB({"sess-old-tip": "sess-new-tip"}) + child = _StubChild() # no identity chain at all + _register( + "sid-durable-lineage-1", child, owner_agent_session_id="sess-old-tip" + ) + rotated_parent = _StubParentWithSession("sess-new-tip", session_db=db) + try: + out = json.loads( + _handle_control_action("list", None, None, rotated_parent) + ) + assert out["count"] == 1 + finally: + _unregister_subagent("sid-durable-lineage-1") + + +def test_owns_subagent_record_requires_some_spine(): + """No identity chain AND no durable owner id -> not owned (fail closed).""" + child = _StubChild() + record = {"subagent_id": "x", "agent": child} + assert _owns_subagent_record(record, _StubParentWithSession("sess-a")) is False + # And a record with an owner id but a parent with no session_id fails too. + record2 = {"subagent_id": "y", "agent": child, "owner_agent_session_id": "s1"} + assert _owns_subagent_record(record2, _StubParent()) is False + + +# --------------------------------------------------------------------------- +# Process-notification attribution: get_subagent_attribution +# --------------------------------------------------------------------------- + + +def test_attribution_resolves_live_child(): + parent = _StubParentWithSession("sess-attr-1") + child = _StubChild(parent) + _register( + "sa-0-attr0001", + child, + delegation_id="deleg_attr_1", + owner_agent_session_id="sess-attr-1", + ) + try: + info = get_subagent_attribution("sa-0-attr0001") + assert info is not None + assert info["subagent_id"] == "sa-0-attr0001" + assert info["goal"] == "test goal" + assert info["delegation_id"] == "deleg_attr_1" + finally: + _unregister_subagent("sa-0-attr0001") + + +def test_attribution_survives_child_completion(): + """After the child finishes (unregistered), attribution must still + resolve — its background processes can outlive it.""" + parent = _StubParentWithSession("sess-attr-2") + child = _StubChild(parent) + _register( + "sa-0-attr0002", + child, + delegation_id="deleg_attr_2", + owner_agent_session_id="sess-attr-2", + ) + _unregister_subagent("sa-0-attr0002") + info = get_subagent_attribution("sa-0-attr0002") + assert info is not None + assert info["delegation_id"] == "deleg_attr_2" + assert info["goal"] == "test goal" + + +def test_attribution_unknown_task_id_is_none(): + assert get_subagent_attribution("proc-not-a-subagent") is None + assert get_subagent_attribution("") is None + assert get_subagent_attribution(None) is None + + +def test_completion_notification_carries_delegation_attribution(): + """format_process_notification on a child-started process completion must + name the subagent + delegation instead of an anonymous output wall.""" + from tools.process_registry import format_process_notification + + parent = _StubParentWithSession("sess-attr-3") + child = _StubChild(parent) + _register( + "sa-1-attr0003", + child, + delegation_id="deleg_attr_3", + goal="run the npm ci for the desktop app", + ) + try: + text = format_process_notification( + { + "type": "completion", + "session_id": "proc_deadbeef0001", + "task_id": "sa-1-attr0003", + "command": "npm ci", + "exit_code": 0, + "output": "added 1500 packages", + } + ) + assert text is not None + assert "Started by subagent sa-1-attr0003" in text + assert "deleg_attr_3" in text + assert "run the npm ci for the desktop app" in text + finally: + _unregister_subagent("sa-1-attr0003") + + +def test_completion_notification_trims_subagent_output_wall(): + from tools.process_registry import format_process_notification + + parent = _StubParentWithSession("sess-attr-4") + child = _StubChild(parent) + _register("sa-2-attr0004", child, delegation_id="deleg_attr_4") + try: + big_output = "npm noise line\n" * 500 + text = format_process_notification( + { + "type": "completion", + "session_id": "proc_deadbeef0002", + "task_id": "sa-2-attr0004", + "command": "npm ci", + "exit_code": 0, + "output": big_output, + } + ) + assert text is not None + assert "output trimmed — subagent-owned process" in text + assert len(text) < len(big_output) + finally: + _unregister_subagent("sa-2-attr0004") + + +def test_parent_owned_process_notification_unchanged(): + """Processes NOT started by a subagent keep the exact legacy shape.""" + from tools.process_registry import format_process_notification + + text = format_process_notification( + { + "type": "completion", + "session_id": "proc_parentowned", + "task_id": "20260817_154314_30d98f", # CLI session task_id + "command": "make build", + "exit_code": 0, + "output": "ok", + } + ) + assert text is not None + assert "Started by subagent" not in text + assert text.startswith("[IMPORTANT: Background process proc_parentowned") + assert "Command: make build\nOutput:\nok]" in text + + +# --------------------------------------------------------------------------- +# Guardrail: control actions never consume the spawn cap +# --------------------------------------------------------------------------- def test_spawn_count_zero_for_control_actions(): from agent.tool_guardrails import _subagent_spawn_count diff --git a/tools/delegate_tool.py b/tools/delegate_tool.py index e1ffb36c4a..8699edb984 100644 --- a/tools/delegate_tool.py +++ b/tools/delegate_tool.py @@ -152,6 +152,44 @@ _active_subagents_lock = threading.Lock() # 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. @@ -179,11 +217,26 @@ def _register_subagent(record: Dict[str, Any]) -> None: _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]: @@ -350,6 +403,66 @@ def _is_descendant_of(child_agent: Any, parent_agent: Any, max_hops: int = 8) -> _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], @@ -368,7 +481,7 @@ def _handle_control_action( entries = [] for r in records: agent = r.get("agent") - if not _is_descendant_of(agent, parent_agent): + if not _owns_subagent_record(r, parent_agent): continue started = r.get("started_at") entries.append( @@ -409,8 +522,7 @@ def _handle_control_action( ) with _active_subagents_lock: record = _active_subagents.get(sid) - target_agent = record.get("agent") if record else None - if record is None or not _is_descendant_of(target_agent, parent_agent): + 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 " @@ -2470,12 +2582,24 @@ def _run_single_child( _raw_depth = getattr(child, "_delegate_depth", 1) _tui_depth = max(0, _raw_depth - 1) if isinstance(_raw_depth, int) else 0 _parent_sid = getattr(child, "_parent_subagent_id", None) + # Durable ownership spine: the OWNING CONVERSATION's session id (the + # same lineage the delivery path routes completions by). Sourced from + # the child's _parent_session_id stamp so it stays correct even when + # parent_agent has been rebuilt between dispatch and this run. + _owner_agent_session_id = ( + str(getattr(child, "_parent_session_id", "") or "") + or str(getattr(parent_agent, "session_id", "") or "") + ) + _delegation_id = getattr(child, "_delegation_id", None) _register_subagent( { "subagent_id": _subagent_id, "parent_id": _parent_sid if isinstance(_parent_sid, str) else None, "depth": _tui_depth, "goal": goal, + "delegation_id": ( + _delegation_id if isinstance(_delegation_id, str) else None + ), "model": ( getattr(child, "model", None) if isinstance(getattr(child, "model", None), str) @@ -2485,6 +2609,11 @@ def _run_single_child( "status": "running", "tool_count": 0, "agent": child, + # Durable conversation lineage for the model-facing control + # plane (list/steer/stop). The weakref identity chain breaks + # when the CLI rebuilds its AIAgent mid-session; this id is + # the same spine completion delivery routes by. + "owner_agent_session_id": _owner_agent_session_id or None, # Immutable live gateway/TUI session that commissioned this # child. Empty outside those hosts; RPC authority fails closed. "owner_session_id": owner_session_id, @@ -3752,6 +3881,10 @@ def delegate_task( getattr(child, "tool_progress_callback", None), _writer ) child._live_transcript_path = str(_writer.path) + # Delegation identity for the live registry + process-notification + # attribution (child-started background processes report under it). + if live_deleg_id: + setattr(child, "_delegation_id", live_deleg_id) children.append((i, t, child)) def _execute_and_aggregate(*, honor_parent_interrupt: bool = True) -> dict: diff --git a/tools/process_registry.py b/tools/process_registry.py index e71cbb47e9..b5fdf7704d 100644 --- a/tools/process_registry.py +++ b/tools/process_registry.py @@ -629,6 +629,7 @@ class ProcessRegistry: notification = { "session_id": session.id, "session_key": session.session_key, + "task_id": session.task_id, "command": session.command, "type": "watch_match", "pattern": matched_pattern, @@ -1574,6 +1575,7 @@ class ProcessRegistry: "type": "completion", "session_id": session.id, "session_key": session.session_key, + "task_id": session.task_id, "command": session.command, "exit_code": session.exit_code, "completion_reason": session.completion_reason, @@ -2872,6 +2874,45 @@ def _format_async_delegation(evt: dict) -> str: return "\n".join(lines) +def _delegation_attribution_line(evt: dict) -> "str | None": + """One-line delegation attribution for a child-originated process event. + + Subagents run their terminal sessions under ``task_id == subagent_id`` + (delegate_tool._run_single_child). When a background process they started + completes, its notification is routed to the PARENT conversation by + design (children consume their own waits via process(wait); anything + that outlives the child must land where a durable consumer exists). + Without attribution the parent-facing user sees an anonymous raw output + wall mid-conversation with no hint it came from a delegation. Resolve + the task_id against the live + recently-finished subagent registry and + return a short provenance line, or None for parent-owned processes. + """ + task_id = str(evt.get("task_id") or "") + if not task_id.startswith("sa-"): + return None + try: + from tools.delegate_tool import get_subagent_attribution + + info = get_subagent_attribution(task_id) + except Exception: + info = None + if not info: + # The task_id shape says "subagent" even when the registry entry has + # aged out — still attribute generically rather than anonymously. + return f"Started by subagent {task_id} (delegate_task)." + goal = str(info.get("goal") or "").strip() + if len(goal) > 120: + goal = goal[:117] + "..." + deleg = info.get("delegation_id") + parts = [f"Started by subagent {task_id}"] + if deleg: + parts.append(f"of delegation {deleg}") + line = " ".join(parts) + "." + if goal: + line += f' Task: "{goal}"' + return line + + def format_process_notification(evt: dict) -> "str | None": """Format a process notification event into a [IMPORTANT: ...] message. @@ -2881,6 +2922,7 @@ def format_process_notification(evt: dict) -> "str | None": evt_type = evt.get("type", "completion") _sid = evt.get("session_id", "unknown") _cmd = evt.get("command", "unknown") + _attribution = _delegation_attribution_line(evt) if evt_type == "watch_disabled": return f"[IMPORTANT: {evt.get('message', '')}]" @@ -2898,6 +2940,10 @@ def format_process_notification(evt: dict) -> "str | None": text = ( f"[IMPORTANT: Background process {_sid} matched " f"watch pattern \"{_pat}\".\n" + ) + if _attribution: + text += f"{_attribution}\n" + text += ( f"Command: {_cmd}\n" f"Matched output:\n{_out}" ) @@ -2926,12 +2972,26 @@ def format_process_notification(evt: dict) -> "str | None": _status = "completed normally" else: _status = "exited" - return ( + text = ( f"[IMPORTANT: Background process {_sid} {_status} " f"(exit code {_exit}{_signal}).\n" + ) + if _attribution: + text += f"{_attribution}\n" + # A subagent-owned process's full output belongs in the child's + # transcript/summary, not as a raw wall in the parent conversation — + # trim the tail hard while keeping enough to recognise failures. + if isinstance(_out, str) and len(_out) > 600: + _out = ( + "...(output trimmed — subagent-owned process; see the " + "delegation's live transcript for full output)\n" + + _out[-600:] + ) + text += ( f"Command: {_cmd}\n" f"Output:\n{_out}]" ) + return text # ---------------------------------------------------------------------------