From 65eaacb341e9fe58a117b18dbd4b68ac62342f4f Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Sat, 27 Jun 2026 20:25:27 -0700 Subject: [PATCH] feat(task-agent): Turn-IR sub-harness + parent-tagged step events Rebuild the task_agent sub-harness on the canonical Turn trajectory (build list[Turn], lower via dicts_from_turns at the wire boundary) instead of hand-rolled OpenAI dicts; the cancel-ledger helpers read Turns. Tag each sub-tool's events with parent_call_id via a lock-guarded child registry stamped centrally in SessionUIBase._enqueue, so a later UI can nest a task agent's steps under its card. Getattr-guarded on the session side so CLI/eval/test UIs are unaffected. Behaviour-preserving (same wire shape, same cancellation semantics); the parent tag is wire-invisible and unconsumed until the frontend card lands. --- tests/test_cancel.py | 74 +++++++-------- tests/test_model_registry.py | 29 +++--- tests/test_session.py | 71 ++++++++++++-- tests/test_session_ui_base.py | 71 ++++++++++++++ turnstone/core/session.py | 150 +++++++++++++++++++----------- turnstone/core/session_ui_base.py | 75 +++++++++++++++ 6 files changed, 356 insertions(+), 114 deletions(-) diff --git a/tests/test_cancel.py b/tests/test_cancel.py index 3c99fcf4..2893789c 100644 --- a/tests/test_cancel.py +++ b/tests/test_cancel.py @@ -15,7 +15,14 @@ from turnstone.core.session import ( _CancelRef, _effect_status_meta, ) -from turnstone.core.trajectory import EffectStatus, Role, dicts_from_turns, turn_from_dict +from turnstone.core.trajectory import ( + EffectStatus, + Role, + ToolCall, + Turn, + dicts_from_turns, + turn_from_dict, +) class NullUI: @@ -1028,15 +1035,11 @@ class TestCancelledAgentDisposition: @staticmethod def _assistant(call_id, name): - return { - "role": "assistant", - "content": "", - "tool_calls": [{"id": call_id, "function": {"name": name}}], - } + return Turn.assistant("", tool_calls=(ToolCall(id=call_id, name=name, arguments=""),)) @staticmethod def _result(call_id, text="ok"): - return {"role": "tool", "tool_call_id": call_id, "content": text} + return Turn.tool(call_id, text) def test_status_none_when_no_actions(self): """Typed twin of the disposition: a task cancelled before any action is @@ -1107,14 +1110,13 @@ class TestCancelledAgentDisposition: # started" — inviting a re-run of the destructive bash. session = _make_session() msgs = [ - { - "role": "assistant", - "content": "", - "tool_calls": [ - {"id": "t1", "function": {"name": "bash"}}, - {"id": "t2", "function": {"name": "web_fetch"}}, - ], - } + Turn.assistant( + "", + tool_calls=( + ToolCall(id="t1", name="bash", arguments=""), + ToolCall(id="t2", name="web_fetch", arguments=""), + ), + ) ] # neither answered: bash raised mid-flight, web_fetch never ran out = session._cancelled_agent_disposition(msgs, "task") assert "In flight at cancel: bash" in out @@ -1127,26 +1129,24 @@ class TestCancelledAgentDisposition: # count summary, the first-gap boundary, and not-started. session = _make_session() msgs = [ - { - "role": "assistant", - "content": "", - "tool_calls": [ - {"id": "t1", "function": {"name": "bash"}}, - {"id": "t2", "function": {"name": "bash"}}, - {"id": "t3", "function": {"name": "read_file"}}, - ], - }, + Turn.assistant( + "", + tool_calls=( + ToolCall(id="t1", name="bash", arguments=""), + ToolCall(id="t2", name="bash", arguments=""), + ToolCall(id="t3", name="read_file", arguments=""), + ), + ), self._result("t1"), self._result("t2"), self._result("t3"), - { - "role": "assistant", - "content": "", - "tool_calls": [ - {"id": "t4", "function": {"name": "web_fetch"}}, - {"id": "t5", "function": {"name": "search"}}, - ], - }, + Turn.assistant( + "", + tool_calls=( + ToolCall(id="t4", name="web_fetch", arguments=""), + ToolCall(id="t5", name="search", arguments=""), + ), + ), ] out = session._cancelled_agent_disposition(msgs, "task") assert "Completed before cancel: bash×2, read_file." in out @@ -1155,13 +1155,13 @@ class TestCancelledAgentDisposition: def test_exec_task_routes_cancel_to_disposition(self, tmp_db): """_exec_task converts a GenerationCancelled from _run_agent into the - honest disposition, reading the in-place-mutated agent_messages.""" + honest disposition, reading the in-place-mutated agent_turns.""" session = _make_session() - def fake_run_agent(agent_messages, **kwargs): - agent_messages.append(self._assistant("t1", "bash")) - agent_messages.append(self._result("t1")) - agent_messages.append(self._assistant("t2", "web_fetch")) + def fake_run_agent(agent_turns, **kwargs): + agent_turns.append(self._assistant("t1", "bash")) + agent_turns.append(self._result("t1")) + agent_turns.append(self._assistant("t2", "web_fetch")) raise GenerationCancelled() with patch.object(session, "_run_agent", side_effect=fake_run_agent): diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index ced9d58b..6c7afb7c 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -15,6 +15,7 @@ from turnstone.core.model_registry import ( detect_model, load_model_registry, ) +from turnstone.core.trajectory import Turn # --------------------------------------------------------------------------- # ModelConfig @@ -1244,8 +1245,8 @@ class TestSessionAgentModel: agent_client.chat.completions.create = fake_create agent_msgs = [ - {"role": "developer", "content": "You are an agent."}, - {"role": "user", "content": "Do something."}, + Turn.system("You are an agent."), + Turn.user("Do something."), ] session._run_agent(agent_msgs) assert captured_model == "agent-model" @@ -1335,14 +1336,14 @@ class TestSessionAgentModel: reg = self._three_model_registry(agent_model="smart", task_model="fast") session = _make_session(registry=reg, model_alias="main") captured = self._capture(reg, "fast") - session._run_agent([{"role": "user", "content": "x"}], label="task") + session._run_agent([Turn.user("x")], label="task") assert captured["model"] == "fast-model" def test_plan_falls_back_to_agent_model(self) -> None: reg = self._three_model_registry(agent_model="fast") session = _make_session(registry=reg, model_alias="main") captured = self._capture(reg, "fast") - session._run_agent([{"role": "user", "content": "x"}], label="plan") + session._run_agent([Turn.user("x")], label="plan") assert captured["model"] == "fast-model" def test_plan_uses_session_model_when_no_overrides(self) -> None: @@ -1351,7 +1352,7 @@ class TestSessionAgentModel: reg = self._three_model_registry() session = _make_session(registry=reg, model_alias="main") captured = self._capture_on(session.client) - session._run_agent([{"role": "user", "content": "x"}], label="plan") + session._run_agent([Turn.user("x")], label="plan") assert captured["model"] == "test-model" def test_task_effort_inherits_session_when_unset(self) -> None: @@ -1362,7 +1363,7 @@ class TestSessionAgentModel: reg = self._three_model_registry() session = _make_session(registry=reg, model_alias="main", reasoning_effort="low") captured = self._capture_on(session.client) - session._run_agent([{"role": "user", "content": "x"}], label="task") + session._run_agent([Turn.user("x")], label="task") assert self._captured_effort(captured) == "low" def test_agent_model_routes_both_plan_and_task(self) -> None: @@ -1372,20 +1373,18 @@ class TestSessionAgentModel: session = _make_session(registry=reg, model_alias="main") plan_captured = self._capture(reg, "fast") - session._run_agent([{"role": "user", "content": "x"}], label="plan") + session._run_agent([Turn.user("x")], label="plan") assert plan_captured["model"] == "fast-model" task_captured = self._capture(reg, "fast") - session._run_agent([{"role": "user", "content": "y"}], label="task") + session._run_agent([Turn.user("y")], label="task") assert task_captured["model"] == "fast-model" def test_explicit_effort_wins_over_registry(self) -> None: reg = self._three_model_registry(task_effort="low") session = _make_session(registry=reg, model_alias="main") captured = self._capture_on(session.client) - session._run_agent( - [{"role": "user", "content": "x"}], label="task", reasoning_effort="minimal" - ) + session._run_agent([Turn.user("x")], label="task", reasoning_effort="minimal") assert self._captured_effort(captured) == "minimal" # -- per-call agent_alias override (LLM passes model="") ---------- @@ -1395,7 +1394,7 @@ class TestSessionAgentModel: reg = self._three_model_registry() session = _make_session(registry=reg, model_alias="main") captured = self._capture(reg, "fast") - session._run_agent([{"role": "user", "content": "x"}], label="task", agent_alias="fast") + session._run_agent([Turn.user("x")], label="task", agent_alias="fast") assert captured["model"] == "fast-model" def test_session_fallback_inherits_primary_alias_for_caps(self) -> None: @@ -1426,7 +1425,7 @@ class TestSessionAgentModel: session._resolve_capabilities = spy_resolve # type: ignore[method-assign] self._capture_on(session.client) # patch client.chat.completions.create - session._run_agent([{"role": "user", "content": "x"}], label="plan") + session._run_agent([Turn.user("x")], label="plan") assert captured_extra_alias and captured_extra_alias[-1] == "main", ( f"agent fallback path did not inherit primary alias for extra_params: " @@ -1443,9 +1442,7 @@ class TestSessionAgentModel: reg = self._three_model_registry() session = _make_session(registry=reg, model_alias="main") with pytest.raises(ValueError, match="Unknown agent_alias"): - session._run_agent( - [{"role": "user", "content": "x"}], label="plan", agent_alias="bogus" - ) + session._run_agent([Turn.user("x")], label="plan", agent_alias="bogus") # --------------------------------------------------------------------------- diff --git a/tests/test_session.py b/tests/test_session.py index 97384496..f328e3be 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -13,6 +13,7 @@ import pytest from turnstone.core.session import _IMAGE_EXTENSIONS, _IMAGE_SIZE_CAP, ChatSession from turnstone.core.trajectory import ( + Turn, dicts_from_turns, turn_from_dict, turn_to_dict, @@ -291,7 +292,7 @@ class TestTaskExec: with patch.object(session, "_run_agent", side_effect=fake_run_agent): session._exec_task(item) - return captured["messages"][0]["content"] + return captured["messages"][0].text def test_known_skill_renders_into_system_message(self, tmp_db) -> None: """Validated skill content (with template vars resolved) replaces @@ -1346,7 +1347,7 @@ class TestAgentOutputGuard: with patch.object(session, "_prepare_tool", side_effect=fake_prepare): session._run_agent( - [{"role": "user", "content": "test"}], + [Turn.user("test")], tools=[{"type": "function", "function": {"name": "read_file"}}], label="test", ) @@ -1409,7 +1410,7 @@ class TestAgentOutputGuard: with patch.object(session, "_prepare_tool", side_effect=fake_prepare): session._run_agent( - [{"role": "user", "content": "test"}], + [Turn.user("test")], tools=[{"type": "function", "function": {"name": "read_file"}}], label="test", ) @@ -1449,7 +1450,7 @@ class TestAgentOutputGuard: session.client.chat.completions.create = fake_create result = session._run_agent( - [{"role": "user", "content": "test"}], + [Turn.user("test")], tools=[{"type": "function", "function": {"name": "read_file"}}], label="plan", ) @@ -1487,7 +1488,7 @@ class TestAgentOutputGuard: session.client.chat.completions.create = fake_create result = session._run_agent( - [{"role": "user", "content": "test"}], + [Turn.user("test")], tools=[{"type": "function", "function": {"name": "read_file"}}], label="task", ) @@ -1522,8 +1523,8 @@ class TestAgentOutputGuard: session.client.chat.completions.create = fake_create result = session._run_agent( [ - {"role": "user", "content": "test"}, - {"role": "assistant", "content": prior}, + Turn.user("test"), + Turn.assistant(prior), ], tools=[{"type": "function", "function": {"name": "read_file"}}], label="plan", @@ -1587,7 +1588,7 @@ class TestAgentOutputGuard: with patch.object(session, "_prepare_tool", side_effect=fake_prepare): result = session._run_agent( - [{"role": "user", "content": "test"}], + [Turn.user("test")], tools=[{"type": "function", "function": {"name": "read_file"}}], label="task", ) @@ -1601,6 +1602,60 @@ class TestAgentOutputGuard: assert synth_args[2] == "task_agent_synthesis" +class TestAgentChildRegistration: + """_run_agent registers each sub-tool under the task's parent_call_id so the + UI can nest the step (the producer side of the SessionUIBase tagging).""" + + def test_sub_tool_registered_under_parent(self): + from turnstone.core.providers._openai_chat import OpenAIChatCompletionsProvider + + session = _make_session() + session._provider = OpenAIChatCompletionsProvider() + session.ui.note_agent_child = MagicMock() + + call_count = [0] + + def fake_create(**_kwargs): + call_count[0] += 1 + resp = MagicMock() + choice = MagicMock() + if call_count[0] == 1: + choice.finish_reason = "tool_calls" + tc = MagicMock() + tc.id = "call_1" + tc.function.name = "read_file" + tc.function.arguments = '{"path": "/tmp/x"}' + choice.message.tool_calls = [tc] + choice.message.content = None + else: + choice.finish_reason = "stop" + choice.message.tool_calls = None + choice.message.content = "done" + resp.choices = [choice] + resp.usage = MagicMock(prompt_tokens=10, completion_tokens=5) + return resp + + session.client.chat.completions.create = fake_create + + def fake_prepare(tc_dict, **_kwargs): + return { + "call_id": tc_dict["id"], + "func_name": "read_file", + "needs_approval": False, + "execute": lambda p: ("call_1", "contents"), + } + + with patch.object(session, "_prepare_tool", side_effect=fake_prepare): + session._run_agent( + [Turn.user("x")], + tools=[{"type": "function", "function": {"name": "read_file"}}], + label="task", + parent_call_id="task-1", + ) + + session.ui.note_agent_child.assert_called_once_with("call_1", "task-1") + + class TestEvaluateOutputLLMStage: """End-to-end coverage of _evaluate_output with the LLM judge stage.""" diff --git a/tests/test_session_ui_base.py b/tests/test_session_ui_base.py index 7b5a4edd..ed621114 100644 --- a/tests/test_session_ui_base.py +++ b/tests/test_session_ui_base.py @@ -2089,3 +2089,74 @@ def test_tool_pending_precedes_smart_approval_gate() -> None: assert approved is True assert captured and captured[0] == "tool_pending", captured + + +# --------------------------------------------------------------------------- +# Sub-agent step tagging (task_agent child events nest under the parent card) +# --------------------------------------------------------------------------- + + +class TestAgentChildTagging: + """``note_agent_child`` makes ``_enqueue`` stamp ``parent_call_id`` on a + sub-tool's events so the UI can nest a task agent's steps under its card. + Keyed on the immutable child call_id (correct under the parent's parallel + tool pool); cleared when the task agent finishes.""" + + def test_registered_child_event_is_stamped(self) -> None: + ui = _make_ui() + lq = ui._register_listener() + ui.note_agent_child("child-1", "task-A") + ui._enqueue({"type": "tool_result", "call_id": "child-1", "name": "bash", "output": "ok"}) + assert lq.get_nowait()["parent_call_id"] == "task-A" + + def test_unregistered_call_id_is_not_stamped(self) -> None: + ui = _make_ui() + lq = ui._register_listener() + ui.note_agent_child("child-1", "task-A") + ui._enqueue({"type": "tool_result", "call_id": "other", "name": "x", "output": "y"}) + assert "parent_call_id" not in lq.get_nowait() + + def test_no_registry_no_stamp(self) -> None: + """Empty registry short-circuits — events pass through untouched.""" + ui = _make_ui() + lq = ui._register_listener() + ui._enqueue({"type": "tool_result", "call_id": "child-1", "name": "x", "output": "y"}) + assert "parent_call_id" not in lq.get_nowait() + + def test_items_payload_is_stamped_per_entry(self) -> None: + """approve_request / tool_pending carry an ``items`` list; each child + entry is tagged independently, leaving non-child entries alone.""" + ui = _make_ui() + lq = ui._register_listener() + ui.note_agent_child("child-1", "task-A") + ui._enqueue( + { + "type": "tool_pending", + "items": [ + {"call_id": "child-1", "func_name": "bash"}, + {"call_id": "top-level", "func_name": "search"}, + ], + } + ) + items = lq.get_nowait()["items"] + assert items[0]["parent_call_id"] == "task-A" + assert "parent_call_id" not in items[1] + + def test_clear_agent_children_stops_stamping(self) -> None: + ui = _make_ui() + lq = ui._register_listener() + ui.note_agent_child("child-1", "task-A") + ui.clear_agent_children("task-A") + ui._enqueue({"type": "tool_result", "call_id": "child-1", "name": "x", "output": "y"}) + assert "parent_call_id" not in lq.get_nowait() + + def test_clear_is_scoped_to_one_parent(self) -> None: + """Two task agents in flight: clearing one leaves the other's children + tagged — the parallel-pool invariant.""" + ui = _make_ui() + lq = ui._register_listener() + ui.note_agent_child("child-A", "task-A") + ui.note_agent_child("child-B", "task-B") + ui.clear_agent_children("task-A") + ui._enqueue({"type": "tool_result", "call_id": "child-B", "name": "x", "output": "y"}) + assert lq.get_nowait()["parent_call_id"] == "task-B" diff --git a/turnstone/core/session.py b/turnstone/core/session.py index dfd4d078..0703cb61 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -143,6 +143,7 @@ from turnstone.core.tools import ( from turnstone.core.trajectory import ( EffectStatus, Role, + ToolCall, Turn, dicts_from_turns, turn_from_dict, @@ -12217,19 +12218,43 @@ class ChatSession: self._report_tool_result(call_id, "diff_file", desc) return call_id, output + def _note_agent_child(self, child_call_id: str, parent_call_id: str | None) -> None: + """Register a sub-agent tool call under its parent task_agent call so the + UI can nest the step (see ``SessionUIBase.note_agent_child``). No-op when + there's no parent (a top-level ``_run_agent``) or the UI doesn't support + it (CLI / eval / fixtures).""" + if not parent_call_id: + return + fn = getattr(self.ui, "note_agent_child", None) + if fn is not None: + fn(child_call_id, parent_call_id) + + def _clear_agent_children(self, parent_call_id: str | None) -> None: + """Drop a finished task agent's child registrations — getattr-guarded + twin of :meth:`_note_agent_child`.""" + if not parent_call_id: + return + fn = getattr(self.ui, "clear_agent_children", None) + if fn is not None: + fn(parent_call_id) + def _run_agent( self, - agent_messages: list[dict[str, Any]], + agent_turns: list[Turn], label: str = "agent", tools: list[dict[str, Any]] | None = None, auto_tools: set[str] | None = None, reasoning_effort: str | None = None, agent_alias: str | None = None, + parent_call_id: str | None = None, ) -> str: """Run an autonomous agent loop. Args: - agent_messages: Pre-built message list (system + developer + user). + agent_turns: Pre-built sub-harness trajectory (system + user) as + neutral ``Turn`` objects, lowered to wire dicts at the API + boundary. Mutated in place — every assistant turn and tool + result is appended as the loop runs. label: Display prefix for progress lines (e.g. "task"). tools: Tool definitions to send to the API. Defaults to the session's task tool set. @@ -12241,6 +12266,9 @@ class ChatSession: the registry's per-kind resolution when set. Caller is expected to have validated the alias against the registry; an unknown alias here raises ``ValueError``. + parent_call_id: The task_agent call_id this sub-agent runs under, + threaded so each sub-tool's events get tagged for UI nesting. + ``None`` for a top-level run (no nesting). Returns: Final content string from the agent. @@ -12290,7 +12318,7 @@ class ChatSession: ) def _api_call( - messages: list[dict[str, Any]], + turns: list[Turn], _tools: list[dict[str, Any]] | None = tools, ) -> CompletionResult: # NOTE: Phase 5 vLLM ``reasoning`` field replay is intentionally @@ -12300,13 +12328,17 @@ class ChatSession: # every turn anyway. Task agents are excluded from the # persistence/replay contract — their conversation history # is in-memory and rebuilt per ``_run_agent`` invocation. + # Lower the trajectory once, not once per retry attempt — ``turns`` + # is invariant across attempts (the retry path only sleeps and + # re-sends the same messages). + wire = dicts_from_turns(turns) last_err: Exception | None = None for attempt in range(self._MAX_RETRIES + 1): try: agent_result = agent_provider.create_completion( client=agent_client, model=agent_model, - messages=messages, + messages=wire, tools=_tools, max_tokens=self.max_tokens, temperature=self.temperature, @@ -12340,7 +12372,7 @@ class ChatSession: while max_tool_turns < 0 or turn < max_tool_turns: self._check_cancelled() try: - result = _api_call(agent_messages) + result = _api_call(agent_turns) except Exception as e: # Context-exceeded or other non-retryable API error. # Return what we have so far rather than crashing. @@ -12348,9 +12380,9 @@ class ChatSession: if "context" in err_str or "token" in err_str: self.ui.on_info(f"[{label}] context limit reached, stopping early") # Find the last assistant content we have - for msg in reversed(agent_messages): - if msg.get("role") == "assistant" and msg.get("content"): - return self._guard_subagent_synthesis(str(msg["content"]), label) + for t in reversed(agent_turns): + if t.role is Role.ASSISTANT and t.text: + return self._guard_subagent_synthesis(t.text, label) return f"({label} stopped: context limit exceeded)" raise @@ -12362,15 +12394,19 @@ class ChatSession: self.ui.on_info(f"[{label}] blocked by content filter") return "(content filter)" - # Build message dict for agent history - msg_dict: dict[str, Any] = { - "role": "assistant", - "content": result.content or "", - } + # Append the assistant turn to the sub-harness trajectory. + agent_tool_calls: tuple[ToolCall, ...] = () if result.tool_calls: self._ensure_tool_call_ids(result.tool_calls) - msg_dict["tool_calls"] = result.tool_calls - agent_messages.append(msg_dict) + agent_tool_calls = tuple( + ToolCall( + id=tc["id"], + name=tc.get("function", {}).get("name", ""), + arguments=tc.get("function", {}).get("arguments", ""), + ) + for tc in result.tool_calls + ) + agent_turns.append(Turn.assistant(result.content or "", tool_calls=agent_tool_calls)) if not result.tool_calls: content = result.content or "(no output)" @@ -12383,6 +12419,10 @@ class ChatSession: for tc_dict in result.tool_calls: self._check_cancelled() tool_name = tc_dict["function"]["name"].strip() + # Register every issued sub-tool under its parent task_agent — + # not only the execute path — so a guard-branch error's + # output_warning still nests under the task card. + self._note_agent_child(tc_dict["id"], parent_call_id) # Guard 1: block recursive agent calls. if tool_name == "task_agent": @@ -12435,28 +12475,29 @@ class ChatSession: if isinstance(output, str) and len(output) > 16000: output = output[:16000] + f"\n\n... (truncated from {len(output)} chars)" - agent_messages.append( - { - "role": "tool", - "tool_call_id": tc_dict["id"], - "content": output, - } - ) + # NOTE: for a vision tool result ``output`` is a list[dict] of + # inline content parts (read_file on an image). It lowers back + # to the same inline list on the wire (behaviour preserved), but + # it is NOT a valid ``TextBlock`` — ``Turn.text`` would raise on + # it. No consumer evaluates ``.text`` on a sub-agent tool turn + # today. The proper by-reference representation needs the + # attachment resolver wired into this sub-agent's + # ``create_completion`` (it currently isn't) plus content- + # addressed byte storage — deferred to the recall/persist work + # where that attachment path is already in scope. + agent_turns.append(Turn.tool(tc_dict["id"], output)) turn += 1 # Exhausted tool turns — force a final synthesis response. self.ui.on_info(f"[{label}] turn limit reached, requesting synthesis...") - agent_messages.append( - { - "role": "user", - "content": ( - "You have reached the tool call limit. " - "Provide your complete response now using " - "the information you have gathered so far." - ), - } + agent_turns.append( + Turn.user( + "You have reached the tool call limit. " + "Provide your complete response now using " + "the information you have gathered so far." + ) ) - result = _api_call(agent_messages, _tools=[]) + result = _api_call(agent_turns, _tools=[]) content = result.content or "(no output)" self.ui.on_info(f"[{label} done] {len(content)} chars") return self._guard_subagent_synthesis(content, label) @@ -12521,21 +12562,22 @@ class ChatSession: # history — it's an autonomous sub-agent. Merged to avoid # multi-system-message errors on models like Qwen. base = self._agent_system_messages[0]["content"] if self._agent_system_messages else "" - agent_messages = [ - {"role": "system", "content": base + "\n\n" + identity}, - {"role": "user", "content": prompt}, + agent_turns: list[Turn] = [ + Turn.system(base + "\n\n" + identity), + Turn.user(prompt), ] try: return call_id, self._run_agent( - agent_messages, + agent_turns, label="task", tools=self._task_tools, auto_tools=TASK_AUTO_TOOLS, agent_alias=item.get("model_override"), + parent_call_id=call_id, ) except GenerationCancelled: # Fold back an honest disposition built from the agent's own - # ledger. ``agent_messages`` is mutated in place by + # ledger. ``agent_turns`` is mutated in place by # ``_run_agent`` (it appends every assistant turn and tool # result), so at this catch point it holds the full record of # what the sub-agent did before cancel — including a partial @@ -12546,8 +12588,8 @@ class ChatSession: # See the cancellation appendix in HYPOTHESIS.md ("ρ may # fabricate the acknowledgment but must not fabricate the # outcome … unknown, never none"). - self._tool_status[call_id] = self._cancelled_agent_status(agent_messages) - return call_id, self._cancelled_agent_disposition(agent_messages, "task") + self._tool_status[call_id] = self._cancelled_agent_status(agent_turns) + return call_id, self._cancelled_agent_disposition(agent_turns, "task") except KeyboardInterrupt: # CLI Ctrl-C: keep the terse string and let the outer loop own # propagation (unchanged behavior). @@ -12555,10 +12597,12 @@ class ChatSession: except Exception as e: self.ui.on_info(f"[task error] {e}") return call_id, f"Task error: {e}" + finally: + self._clear_agent_children(call_id) @staticmethod def _cancel_ledger( - agent_messages: list[dict[str, Any]], + agent_turns: list[Turn], ) -> tuple[list[tuple[str, bool]], int | None]: """Read a cancelled sub-agent's ledger: every issued tool call as ``(name, was_answered)`` in order, plus the index of the first in-flight @@ -12573,32 +12617,32 @@ class ChatSession: its typed status so the two can't disagree. """ answered: set[str] = set() - for m in agent_messages: - if m.get("role") == "tool" and m.get("tool_call_id"): - answered.add(str(m["tool_call_id"])) + for t in agent_turns: + if t.role is Role.TOOL and t.tool_call_id: + answered.add(t.tool_call_id) issued: list[tuple[str, bool]] = [] - for m in agent_messages: - if m.get("role") != "assistant": + for t in agent_turns: + if t.role is not Role.ASSISTANT: continue - for tc in m.get("tool_calls") or []: - name = ((tc.get("function") or {}).get("name") or "tool").strip() - issued.append((name, str(tc.get("id") or "") in answered)) + for tc in t.tool_calls: + name = (tc.name or "tool").strip() + issued.append((name, tc.id in answered)) first_gap = next((i for i, (_n, ans) in enumerate(issued) if not ans), None) return issued, first_gap - def _cancelled_agent_status(self, agent_messages: list[dict[str, Any]]) -> EffectStatus: + def _cancelled_agent_status(self, agent_turns: list[Turn]) -> EffectStatus: """Typed twin of :meth:`_cancelled_agent_disposition`: ``none`` if the agent never acted, ``unknown`` if a tool was in flight when cancel landed (its effect unobserved), else ``partial`` — every issued call returned, but the agent was stopped before finishing.""" - issued, first_gap = self._cancel_ledger(agent_messages) + issued, first_gap = self._cancel_ledger(agent_turns) if not issued: return EffectStatus.NONE if first_gap is not None: return EffectStatus.UNKNOWN return EffectStatus.PARTIAL - def _cancelled_agent_disposition(self, agent_messages: list[dict[str, Any]], label: str) -> str: + def _cancelled_agent_disposition(self, agent_turns: list[Turn], label: str) -> str: """Build an honest, deterministic disposition for a cancelled sub-agent. A cancelled agent must not report a fabricated *outcome*. The bare @@ -12613,12 +12657,12 @@ class ChatSession: side effect may or may not have landed, its result never observed), and everything after it never ran. - Pure string assembly over the in-memory ``agent_messages`` — no + Pure string assembly over the in-memory ``agent_turns`` — no model call, because we are on the cancel path and the gate is closed. The owner (parent / coordinator) reads this to decide what, if anything, to compensate. """ - issued, first_gap = self._cancel_ledger(agent_messages) + issued, first_gap = self._cancel_ledger(agent_turns) if not issued: return f"({label} cancelled by user before any action — no side effects)" diff --git a/turnstone/core/session_ui_base.py b/turnstone/core/session_ui_base.py index 01d77de3..be0c80ed 100644 --- a/turnstone/core/session_ui_base.py +++ b/turnstone/core/session_ui_base.py @@ -251,6 +251,17 @@ class SessionUIBase: # in :meth:`_enqueue`); the snapshot helper for the in-progress # replay path captures it under ``_listeners_lock`` too. self._event_id: int = 0 + # Sub-agent step tagging: child tool call_id -> parent task_agent + # call_id. ``_enqueue`` reads it to stamp ``parent_call_id`` on every + # child event (tool_pending / approve_request / tool_result / + # tool_output_chunk / output_warning) so the UI can nest a task agent's + # steps under its card. Written by ``note_agent_child`` / + # ``clear_agent_children`` (the session brackets each ``_run_agent``); + # keyed on the immutable call_id so it stays correct under the parent's + # 4-wide parallel tool pool (several task agents in flight at once). Its + # own lock so the hot fan-out path never serializes on ``_listeners_lock``. + self._agent_children: dict[str, str] = {} + self._agent_children_lock = threading.Lock() # Approval blocking — the worker thread calls approve_tools # which waits on _approval_event; the /approve endpoint sets # it via resolve_approval. @@ -445,6 +456,11 @@ class SessionUIBase: """ if "ws_id" not in data: data = {**data, "ws_id": self.ws_id} + # Only events that can carry a child step — a top-level ``call_id`` or an + # ``items`` list — need the parent-tag lookup; skip the lock+scan for the + # high-frequency rest (content / reasoning / status / info / …). + if self._agent_children and ("call_id" in data or "items" in data): + data = self._stamp_agent_parent(data) with self._listeners_lock: self._event_id += 1 event_id = self._event_id @@ -462,6 +478,65 @@ class SessionUIBase: lq.put_nowait(data) return event_id + def _stamp_agent_parent(self, data: dict[str, Any]) -> dict[str, Any]: + """Stamp ``parent_call_id`` on a sub-agent's child event. + + A task agent's sub-tool events flow through the same emit path as + top-level tool events; the only thing marking them as a *child* is the + ``call_id`` registered in ``_agent_children`` by the session running + ``_run_agent``. Stamping at this one fan-out choke point (rather than at + each call site) also catches events emitted from a tool's own background + thread — a streaming bash's ``tool_output_chunk`` — which an ambient + context var set on the worker thread would miss. Returns a shallow copy + when it stamps; the input dict is never mutated.""" + with self._agent_children_lock: + if not self._agent_children: + return data + cid = data.get("call_id") + if isinstance(cid, str) and cid in self._agent_children: + data = {**data, "parent_call_id": self._agent_children[cid]} + items = data.get("items") + if isinstance(items, list) and any( + isinstance(it, dict) and it.get("call_id") in self._agent_children for it in items + ): + data = { + **data, + "items": [ + {**it, "parent_call_id": self._agent_children[it["call_id"]]} + if isinstance(it, dict) and it.get("call_id") in self._agent_children + else it + for it in items + ], + } + return data + + def note_agent_child(self, child_call_id: str, parent_call_id: str) -> None: + """Register a sub-agent's child tool call so ``_enqueue`` tags its events + with ``parent_call_id``. Called by the session for each sub-tool a + ``_run_agent`` issues, before the tool emits anything. + + KNOWN LIMITATION (to be closed with the chunk-3 nesting consumer): the + registry keys on ``child_call_id`` alone — unique for commercial + providers (``call_`` / ``toolu_…``) but NOT for local servers that + assign per-response sequential ids (``call_0``). Two task agents in the + parent's 4-wide pool can then collide and mis-nest steps. The fix + (namespacing child ids by parent, verified against each provider's + id-replay rules) lands with the frontend consumer that renders the + nesting — until then the tag is wire-invisible and unconsumed.""" + if not child_call_id or not parent_call_id: + return + with self._agent_children_lock: + self._agent_children[child_call_id] = parent_call_id + + def clear_agent_children(self, parent_call_id: str) -> None: + """Drop every child registered under ``parent_call_id`` (the task agent + finished). Bounds the registry to in-flight task agents. Deletes in + place rather than reallocating the whole dict, so one agent completing + doesn't churn other in-flight agents' entries.""" + with self._agent_children_lock: + for c in [c for c, p in self._agent_children.items() if p == parent_call_id]: + del self._agent_children[c] + def _register_listener( self, maxsize: int = _DEFAULT_LISTENER_QUEUE_MAX ) -> queue.Queue[dict[str, Any]]: