ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity

feat(task-agent): Turn-IR sub-harness + parent-tagged step events

Browse Source 
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.
This commit is contained in:
Patrick Buckley
2026-06-27 20:25:27 -07:00
parent 9837214414
commit 65eaacb341
6 changed files with 356 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_cancel.py
+37 -37
View File
@@ -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):
tests/test_model_registry.py
+13 -16
View File
@@ -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="<alias>") ----------
@@ -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")
# ---------------------------------------------------------------------------
tests/test_session.py
+63 -8
View File
@@ -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."""
tests/test_session_ui_base.py
+71
View File
@@ -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"
turnstone/core/session.py
+97 -53
View File
@@ -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)"
turnstone/core/session_ui_base.py
+75
View File
@@ -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_<hash>`` / ``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]]: