From 7befa11bf25b94a91edcdd11bb10161acbb03c1b Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:07:59 -0700 Subject: [PATCH] fix: retain subagent control after live session reattachment --- tests/tui_gateway/test_subagent_snapshot.py | 72 +++++++++++++++++++++ tools/delegate_tool_registry.py | 5 ++ tui_gateway/AGENTS.md | 4 +- tui_gateway/session_lifecycle.py | 14 +++- 4 files changed, 92 insertions(+), 3 deletions(-) diff --git a/tests/tui_gateway/test_subagent_snapshot.py b/tests/tui_gateway/test_subagent_snapshot.py index ca10d98cf1..2d38cd9b0a 100644 --- a/tests/tui_gateway/test_subagent_snapshot.py +++ b/tests/tui_gateway/test_subagent_snapshot.py @@ -148,3 +148,75 @@ def test_interrupt_requires_exact_live_owner_but_direct_helper_stays_legacy(runt assert not call("subagent.interrupt", subagent_id="child")["result"]["found"] finally: _unregister_subagent("child") + + + +def test_reattach_preserves_child_controls_including_late_registration(runtime, tmp_path): + from tools.delegate_tool_child_run import _register_child + + server, owner, old, call = runtime + new = type("Transport", (), {"write": lambda self, frame: True})() + transcript = tmp_path / "child.txt" + transcript.write_text("live child output") + steered, stopped = [], [] + + def register(sid): + child = SimpleNamespace(_subagent_id=sid, _delegate_depth=1, model="test", + _live_transcript_path=str(transcript), + steer=lambda text: steered.append(text) or True, + hard_interrupt=lambda text: stopped.append(text)) + _register_child(child, None, "owned", owner_session_id="ui-owner", + owner_transport=old, owner_session_record=owner) + + register("before") + owner["transport"] = server._detached_ws_transport + owner["history_lock"] = threading.Lock() + with server._session_resume_lock, owner["history_lock"]: + assert server._reattach_refusal(1, "ui-owner", owner) is None + server._rebind_live_transport("ui-owner", owner, new) + # A dispatch captured before reload may not construct its child until afterwards. + register("after") + assert {row["subagent_id"] for row in call("subagent.list", via=new)["result"]["subagents"]} == {"before", "after"} + # Closing a second authenticated viewer hands control back to the survivor. + popup = type(new)() + with server._session_resume_lock, owner["history_lock"]: + server._rebind_live_transport("ui-owner", owner, popup) + assert server._close_sessions_for_transport(popup) == (0, 0) + assert owner["transport"] is new + assert {row["subagent_id"] for row in call("subagent.list", via=new)["result"]["subagents"]} == {"before", "after"} + for sid in ("before", "after"): + assert call("subagent.tail", via=new, subagent_id=sid)["result"]["text"] == "live child output" + assert call("subagent.steer", via=new, subagent_id=sid, text=sid)["result"]["status"] == "queued" + assert call("subagent.interrupt", via=new, subagent_id=sid)["result"]["found"] + assert steered == ["before", "after"] and len(stopped) == 2 + for method in ("list", "tail", "steer", "interrupt"): + denied = call("subagent." + method, subagent_id="before", text="old") + assert "error" in denied or denied["result"].get("status") == "rejected" + assert steered == ["before", "after"] and len(stopped) == 2 + + +def test_reattach_does_not_adopt_foreign_or_retired_generations(runtime): + from tools.delegate_tool_child_run import _register_child + + server, owner, old, call = runtime + new = type("Transport", (), {"write": lambda self, frame: True})() + effects = [] + for sid, session_id, record in (("foreign", "other", owner), + ("retired", "ui-owner", {**owner})): + child = SimpleNamespace(_subagent_id=sid, _delegate_depth=1, model="test", + steer=lambda text: effects.append(text) or True, + hard_interrupt=lambda text: effects.append(text)) + _register_child(child, None, "private", owner_session_id=session_id, + owner_transport=old, owner_session_record=record) + with server._session_resume_lock: + assert server._reattach_refusal(1, "ui-owner", {**owner})["error"]["code"] == 4007 + owner["_client_gone_interrupt_requested"] = True + assert server._reattach_refusal(1, "ui-owner", owner)["error"]["code"] == 4009 + del owner["_client_gone_interrupt_requested"] + server._rebind_live_transport("ui-owner", owner, new) + assert call("subagent.list", via=new)["result"]["subagents"] == [] + for sid in ("foreign", "retired"): + assert not call("subagent.tail", via=new, subagent_id=sid)["result"]["available"] + assert call("subagent.steer", via=new, subagent_id=sid, text="deny")["result"]["status"] == "rejected" + assert not call("subagent.interrupt", via=new, subagent_id=sid)["result"]["found"] + assert effects == [] diff --git a/tools/delegate_tool_registry.py b/tools/delegate_tool_registry.py index 43d295d9ec..dd2a225e12 100644 --- a/tools/delegate_tool_registry.py +++ b/tools/delegate_tool_registry.py @@ -53,6 +53,11 @@ def _register_subagent(record: Dict[str, Any]) -> None: return record.setdefault("accepting_steer", True) with _active_subagents_lock: + owner = record.get("owner_session_record") + if owner is not None and record.get("owner_transport") is not None: + # Child construction can finish after its captured dispatch transport + # was replaced. The exact session object retains generation authority. + record["owner_transport"] = owner.get("transport") _active_subagents[sid] = record def _unregister_subagent(subagent_id: str, *, agent: Any = None) -> None: diff --git a/tui_gateway/AGENTS.md b/tui_gateway/AGENTS.md index bb571fd345..d4181645bd 100644 --- a/tui_gateway/AGENTS.md +++ b/tui_gateway/AGENTS.md @@ -43,7 +43,9 @@ existing topical sibling, registered in the table — no `if method == ...` chai `subagent.list({session_id})` returns `{subagents, delegations}` for the calling transport's live session. Live child records are pinned to the exact session -record and transport. `last_tool` is the last started tool, not an in-flight +record and transport. Authenticated live reattachment transfers that exact generation's +child authority to the new transport (also for late child registration and surviving +viewers); foreign or retired generations remain inaccessible. `last_tool` is the last started tool, not an in-flight indicator. Async completion units are not agents and lack exact generation authority; `delegations` remains an empty array for wire compatibility. No dispatch context, results, callbacks, or routing keys are sent. Clients hydrate from this snapshot diff --git a/tui_gateway/session_lifecycle.py b/tui_gateway/session_lifecycle.py index 8792565c85..166abfe247 100644 --- a/tui_gateway/session_lifecycle.py +++ b/tui_gateway/session_lifecycle.py @@ -466,7 +466,17 @@ def _reattach_refusal(rid, sid: str, session: dict) -> dict | None: def _rebind_live_transport(sid: str, session: dict, transport: Transport) -> None: """Attach a live peer without displacing existing subscribers (caller holds ``history_lock``).""" - _attach_session_transport(session, transport) + from tools.delegate_tool_registry import _active_subagents, _active_subagents_lock + + # Transfer only this exact live generation's capabilities at the authenticated + # attachment seam, including records spawned through an older dispatch context. + with _active_subagents_lock: + _attach_session_transport(session, transport) + for record in _active_subagents.values(): + if (record.get("owner_session_id") == sid + and record.get("owner_session_record") is session + and record.get("owner_transport") is not None): + record["owner_transport"] = transport # Every transport that showed this session (pop-outs resume the same sid); on disconnect the last # viewer becomes the transport instead of the drop sentinel. session.setdefault("viewers", {})[transport] = time.time() @@ -632,7 +642,7 @@ def _close_sessions_for_transport(transport, *, end_reason: str = "ws_disconnect viewers.pop(transport, None) live = [vt for vt, ts in sorted(viewers.items(), key=lambda kv: kv[1]) if not _transport_is_dead(vt)] if live: - current["transport"] = live[-1] + _rebind_live_transport(sid, current, live[-1]) else: current["transport"] = _detached_ws_transport current.pop("_client_gone_interrupt_requested", None)