Merge remote-tracking branch 'origin/main' into fix/async-batch-task-failure-notice
# Conflicts: # tools/delegate_tool_dispatch.py
This commit is contained in:
+119
-64
@@ -1,6 +1,7 @@
|
||||
"""Batch execution + background dispatch for delegate_task: ``delegate_task`` builds a ``_Batch``
|
||||
(children + origin identity) and hands it to ``_run_batch``, which runs it synchronously or as ONE
|
||||
detached async unit."""
|
||||
(children + origin identity) and hands it to ``_run_batch``, which runs it synchronously or as
|
||||
detached async units — one per task ``group`` (ungrouped tasks are a unit each), so a finished
|
||||
review does not wait for the slowest sibling."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -9,9 +10,10 @@ import json
|
||||
import logging
|
||||
import time
|
||||
from concurrent.futures import FIRST_COMPLETED, wait as _cf_wait
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from tools.async_delegation import _new_delegation_id, record_unit_child
|
||||
from tools.delegate_tool_child_run import _detach_child, _fabricated_entry, _signal_child_stop
|
||||
from tools.delegate_tool_progress import (
|
||||
SUBAGENT_FAILURE_STATUSES, _clean_error_text, _print_completion_line, _quiet, format_batch_tag,
|
||||
@@ -42,6 +44,9 @@ class _Batch:
|
||||
origin_owner_transport: Any
|
||||
origin_owner_session_record: Any
|
||||
overall_start: float
|
||||
# Set on per-group units carved out by ``_dispatch_background``; None for the whole batch / ungrouped units.
|
||||
group: Optional[str] = None
|
||||
unit_id: Optional[str] = None # the async registry id this unit runs under (``<call_id>-k`` for split calls)
|
||||
|
||||
def owner_kwargs(self) -> Dict[str, Any]:
|
||||
"""Steer/stop authority of the originating session, passed to every child run."""
|
||||
@@ -58,7 +63,7 @@ class _Batch:
|
||||
def _announce_batch(parent_agent, n_tasks: int, live_deleg_id: Optional[str]) -> None:
|
||||
"""Announce the batch tag once so interleaved ``[tag n/N]`` lines are attributable."""
|
||||
if n_tasks > 1 and live_deleg_id:
|
||||
_hdr = f" 🔀 [{format_batch_tag(live_deleg_id)}] delegating {n_tasks} tasks"
|
||||
_hdr = f" 🔀 [{format_batch_tag(live_deleg_id, parent_agent)}] delegating {n_tasks} tasks"
|
||||
_print_completion_line(parent_agent, getattr(parent_agent, "_delegate_spinner", None), _hdr, console_line=_hdr)
|
||||
|
||||
def _capture_origin() -> tuple[str, str, Any, Any]:
|
||||
@@ -100,9 +105,10 @@ def _run_children_parallel(batch: _Batch, results: list, *, honor_parent_interru
|
||||
parent_agent, n_tasks = batch.parent_agent, len(batch.task_list)
|
||||
task_labels = [t["goal"][:40] for t in batch.task_list]
|
||||
spinner_ref = getattr(parent_agent, "_delegate_spinner", None)
|
||||
_tag = format_batch_tag(batch.live_deleg_id)
|
||||
_tag = format_batch_tag(batch.live_deleg_id, parent_agent)
|
||||
# Fabricated entries for still-pending / raised futures carry the correct _delegate_role.
|
||||
_child_by_index = {i: child for (i, _, child) in batch.children}
|
||||
n_here = len(batch.children) # a per-group unit runs a subset; ``n_tasks`` keeps the call-wide ``i/N`` slot
|
||||
|
||||
def _entry_of(future, idx):
|
||||
if not future.done():
|
||||
@@ -125,25 +131,30 @@ def _run_children_parallel(batch: _Batch, results: list, *, honor_parent_interru
|
||||
for future in done:
|
||||
entry = _entry_of(future, futures[future])
|
||||
results.append(entry)
|
||||
_report_child_done(parent_agent, spinner_ref, entry, _tag, task_labels, n_tasks, n_tasks - len(results))
|
||||
if detached and entry.get("status") in SUBAGENT_FAILURE_STATUSES and len(results) < n_tasks:
|
||||
# The parent is not watching a spinner: tell it NOW, not when the last sibling finishes.
|
||||
if not honor_parent_interrupt and batch.unit_id:
|
||||
# Detached unit: a crash before the join must not lose children that already finished.
|
||||
record_unit_child(batch.unit_id, entry)
|
||||
_report_child_done(parent_agent, spinner_ref, entry, _tag, task_labels, n_tasks, n_here - len(results))
|
||||
if (not honor_parent_interrupt and batch.unit_id and entry.get("status") in SUBAGENT_FAILURE_STATUSES
|
||||
and len(results) < n_here):
|
||||
# Detached unit, a sibling is still running: tell the parent NOW, not when the last one finishes.
|
||||
# Non-durable and separate from the unit's final result (which is still delivered once).
|
||||
with _quiet("task failure notice failed", exc_info=True):
|
||||
from tools.async_delegation import push_task_failure_notice
|
||||
_i = entry.get("task_index", -1)
|
||||
_live = batch.live_paths[_i] if isinstance(_i, int) and 0 <= _i < len(batch.live_paths) else None
|
||||
push_task_failure_notice(
|
||||
batch.live_deleg_id, {**entry, **({"live_transcript": _live} if _live else {})}, n_tasks=n_tasks)
|
||||
batch.unit_id, {**entry, **({"live_transcript": _live} if _live else {})}, n_tasks=n_tasks)
|
||||
results.sort(key=lambda r: r["task_index"]) # match input order
|
||||
|
||||
def _execute_and_aggregate(batch: _Batch, *, honor_parent_interrupt: bool = True, detached: bool = False) -> dict:
|
||||
"""Run all built children, join, finalize (hooks + cost rollup), return the combined dict. Shared by the sync path
|
||||
and the background runner: even in the background the batch JOINS on itself here so ONE consolidated results
|
||||
block re-enters the conversation. Live transcripts are finalized but retained as the full-fidelity record
|
||||
(retention pruning happens on future dispatches)."""
|
||||
def _execute_and_aggregate(batch: _Batch, *, honor_parent_interrupt: bool = True) -> dict:
|
||||
"""Run the batch's built children, join, finalize (hooks + cost rollup), return the combined dict. Shared by the
|
||||
sync path and the background runner; a background runner receives a per-group unit (a subset of the call's
|
||||
children) so each group JOINS only on itself. Live transcripts are finalized but retained as the full-fidelity
|
||||
record (retention pruning happens on future dispatches)."""
|
||||
from tools.delegation_live_log import update_manifest_statuses
|
||||
results: list = []
|
||||
if len(batch.task_list) == 1:
|
||||
if len(batch.children) == 1:
|
||||
results.append(batch.run_child(*batch.children[0]))
|
||||
else:
|
||||
_run_children_parallel(batch, results, honor_parent_interrupt=honor_parent_interrupt, detached=detached)
|
||||
@@ -160,8 +171,11 @@ def _execute_and_aggregate(batch: _Batch, *, honor_parent_interrupt: bool = True
|
||||
update_manifest_statuses(batch.live_deleg_id, results)
|
||||
|
||||
combined: Dict[str, Any] = {"results": results, "total_duration_seconds": total_duration}
|
||||
if batch.live_paths:
|
||||
combined["live_transcripts"] = list(batch.live_paths)
|
||||
unit_paths = [batch.live_paths[i] for (i, _, _) in batch.children if i < len(batch.live_paths)]
|
||||
if unit_paths:
|
||||
combined["live_transcripts"] = unit_paths
|
||||
if batch.group is not None:
|
||||
combined["group"] = batch.group
|
||||
return combined
|
||||
|
||||
_SYNC_FALLBACK_NOTES = {
|
||||
@@ -264,9 +278,10 @@ _BACKGROUND_NOTES = {
|
||||
"conversation as a new message when it finishes. Do not wait or poll — just continue."
|
||||
),
|
||||
"many": (
|
||||
"{n} subagents are running in parallel in the background. You and the user can keep working; they wait on "
|
||||
"each other and their consolidated results re-enter the conversation as a single message once ALL of them "
|
||||
"finish. Do not wait or poll — just continue."
|
||||
"{n} subagents are running in parallel in the background as {k} independent unit(s). You and the user can keep "
|
||||
"working; each unit's results re-enter the conversation as their own message as soon as THAT unit finishes "
|
||||
"(tasks sharing a `group` finish together; ungrouped tasks report individually), so act on each as it lands. "
|
||||
"Do not wait or poll — just continue."
|
||||
),
|
||||
"control_hint": (
|
||||
"While a child runs you can orchestrate it live with this same tool: delegate_task(action='list') to see live "
|
||||
@@ -279,29 +294,66 @@ _BACKGROUND_NOTES = {
|
||||
),
|
||||
}
|
||||
|
||||
def _dispatched_payload(dispatch: dict, goals: List[str], child_agents: List[Any], live_paths: List[str]) -> dict:
|
||||
"""Model-facing handle for an accepted background batch."""
|
||||
def _dispatched_payload(batch: _Batch, units: List[tuple[_Batch, str]]) -> dict:
|
||||
"""Model-facing handle for an accepted background call: one entry per async unit."""
|
||||
goals = [t["goal"] for t in batch.task_list]
|
||||
n = len(goals)
|
||||
payload = {
|
||||
"status": "dispatched", "mode": "background", "count": n, "delegation_id": dispatch["delegation_id"],
|
||||
"goals": goals, "note": _BACKGROUND_NOTES["one"] if n == 1 else _BACKGROUND_NOTES["many"].format(n=n),
|
||||
"status": "dispatched", "mode": "background", "count": n,
|
||||
"delegation_id": batch.live_deleg_id or units[0][1], "goals": goals,
|
||||
"note": _BACKGROUND_NOTES["one"] if n == 1 else _BACKGROUND_NOTES["many"].format(n=n, k=len(units)),
|
||||
}
|
||||
sids = [getattr(c, "_subagent_id", None) for c in child_agents]
|
||||
if len(units) > 1:
|
||||
payload["units"] = [
|
||||
{"delegation_id": uid, "group": unit.group, "task_indexes": [i for (i, _, _) in unit.children]}
|
||||
for unit, uid in units
|
||||
]
|
||||
sids = [getattr(c, "_subagent_id", None) for (_, _, c) in batch.children]
|
||||
if any(isinstance(s, str) and s for s in sids):
|
||||
payload["subagent_ids"] = sids
|
||||
payload["control_hint"] = _BACKGROUND_NOTES["control_hint"]
|
||||
if live_paths:
|
||||
payload["live_transcripts"] = list(live_paths)
|
||||
if batch.live_paths:
|
||||
payload["live_transcripts"] = list(batch.live_paths)
|
||||
payload["live_transcripts_hint"] = _BACKGROUND_NOTES["live_transcripts_hint"]
|
||||
return payload
|
||||
|
||||
def _dispatch_background(batch: _Batch) -> str:
|
||||
"""Dispatch the WHOLE batch as one async unit and return the tool result JSON. The runner joins on every child and
|
||||
yields ONE consolidated results block that re-enters the conversation as a single message when ALL children
|
||||
finish. Falls back to running synchronously (with an explanatory ``note``) when the session cannot receive
|
||||
detached completions or the async pool is at capacity."""
|
||||
from tools.delegate_tool import _get_max_async_children
|
||||
def _units_of(batch: _Batch) -> List[_Batch]:
|
||||
"""Partition the call's children into async units: one per distinct task ``group`` (first-appearance order) and
|
||||
one per ungrouped task. Each unit is a ``_Batch`` sharing the call's task_list/transcripts but owning a subset of
|
||||
``children``, so a unit joins only on itself and its completion re-enters the conversation on its own."""
|
||||
members: Dict[Any, List[tuple]] = {}
|
||||
for i, t, c in batch.children:
|
||||
g = t.get("group")
|
||||
key = ("g", str(g)) if g not in (None, "") else ("i", i)
|
||||
members.setdefault(key, []).append((i, t, c))
|
||||
return [replace(batch, children=ch, group=(key[1] if key[0] == "g" else None)) for key, ch in members.items()]
|
||||
|
||||
def _dispatch_unit(unit: _Batch, unit_id: Optional[str], slot_key: Optional[str], routing: dict) -> dict:
|
||||
"""Hand ONE unit to the async registry; the runner joins on that unit's children only."""
|
||||
from tools.async_delegation import dispatch_async_delegation_batch
|
||||
child_agents = [c for (_, _, c) in unit.children]
|
||||
|
||||
def _interrupt():
|
||||
for c in child_agents:
|
||||
_signal_child_stop(c, "Async delegation cancelled")
|
||||
|
||||
return dispatch_async_delegation_batch(
|
||||
# Call-wide goals: completion formatting indexes them by task_index.
|
||||
goals=[t["goal"] for t in unit.task_list], context=unit.context,
|
||||
toolsets=None, # metadata for the completion block only; subagents inherit the parent's toolsets
|
||||
role=unit.top_role, model=unit.creds["model"],
|
||||
runner=lambda: _execute_and_aggregate(unit, honor_parent_interrupt=False),
|
||||
interrupt_fn=_interrupt, delegation_id=unit_id, slot_key=slot_key,
|
||||
task_indexes=[i for (i, _, _) in unit.children] if len(unit.children) < len(unit.task_list) else None,
|
||||
progress_fn=lambda: _batch_progress_token(child_agents), **routing,
|
||||
)
|
||||
|
||||
def _dispatch_background(batch: _Batch) -> str:
|
||||
"""Dispatch the call as independent async units (see ``_units_of``) and return the tool result JSON. Every unit
|
||||
of one call shares ONE pool slot (``slot_key``), so grouping never changes capacity accounting. Falls back to
|
||||
running synchronously (with an explanatory ``note``) when the session cannot receive detached completions or the
|
||||
async pool is at capacity."""
|
||||
from tools.delegate_tool import _get_max_async_children
|
||||
wake_sid = _resolve_async_wake_sid(batch.origin_wake_sid)
|
||||
if wake_sid is None:
|
||||
logger.info("delegate_task: async delivery unsupported on this session runtime; running the batch synchronously instead.")
|
||||
@@ -309,40 +361,43 @@ def _dispatch_background(batch: _Batch) -> str:
|
||||
|
||||
parent_agent = batch.parent_agent
|
||||
session_key, origin_ui_session_id = _resolve_async_session_key(parent_agent, batch.origin_ui_session_id)
|
||||
child_agents = [c for (_, _, c) in batch.children]
|
||||
# The batch's lifecycle is owned by the async registry now: drop the children from the parent's
|
||||
# The children's lifecycle is owned by the async registry now: drop them from the parent's
|
||||
# interrupt-propagation list (_build_child_agent attached them, which is correct for sync runs).
|
||||
for c in child_agents:
|
||||
for (_, _, c) in batch.children:
|
||||
_detach_child(parent_agent, c)
|
||||
|
||||
def _batch_interrupt():
|
||||
# Cancellation path for the detached batch (owned by the async registry).
|
||||
for c in child_agents:
|
||||
_signal_child_stop(c, "Async delegation cancelled")
|
||||
|
||||
goals = [t["goal"] for t in batch.task_list]
|
||||
dispatch = dispatch_async_delegation_batch(
|
||||
goals=goals, context=batch.context,
|
||||
toolsets=None, # metadata for the completion block only; subagents inherit the parent's toolsets
|
||||
role=batch.top_role, model=batch.creds["model"], session_key=session_key,
|
||||
origin_ui_session_id=origin_ui_session_id, origin_session_id=wake_sid,
|
||||
parent_session_id=getattr(parent_agent, "session_id", None),
|
||||
runner=lambda: _execute_and_aggregate(batch, honor_parent_interrupt=False, detached=True),
|
||||
interrupt_fn=_batch_interrupt, max_async_children=_get_max_async_children(),
|
||||
# Reuse the live-transcript directory's id (when created) so the returned delegation_id matches
|
||||
# cache/delegation/live/<id>/.
|
||||
delegation_id=batch.live_deleg_id,
|
||||
progress_fn=lambda: _batch_progress_token(child_agents),
|
||||
routing = dict(
|
||||
session_key=session_key, origin_ui_session_id=origin_ui_session_id, origin_session_id=wake_sid,
|
||||
parent_session_id=getattr(parent_agent, "session_id", None), max_async_children=_get_max_async_children(),
|
||||
)
|
||||
if dispatch.get("status") == "dispatched":
|
||||
return json.dumps(_dispatched_payload(dispatch, goals, child_agents, batch.live_paths), ensure_ascii=False)
|
||||
# Pool at capacity / schedule failure: the async unit was never accepted, so just run inline (re-attaching to the
|
||||
# parent list is not needed).
|
||||
logger.info(
|
||||
"delegate_task: async pool at capacity (%s); running the whole batch synchronously instead.",
|
||||
dispatch.get("error", "rejected"),
|
||||
)
|
||||
return _run_sync_with_note(batch, "at_capacity")
|
||||
|
||||
units = _units_of(batch)
|
||||
dispatched: List[tuple[_Batch, str]] = []
|
||||
inline_results: List[dict] = []
|
||||
slot_key: Optional[str] = None
|
||||
for k, unit in enumerate(units):
|
||||
# One unit keeps the live-transcript directory's id so the returned delegation_id matches
|
||||
# cache/delegation/live/<id>/; several units suffix it (-1, -2, ...) and the call keeps the bare id.
|
||||
unit_id = batch.live_deleg_id if len(units) == 1 else (f"{batch.live_deleg_id}-{k + 1}" if batch.live_deleg_id else None)
|
||||
unit.unit_id = unit_id = unit_id or _new_delegation_id() # fixed before the runner can start
|
||||
dispatch = _dispatch_unit(unit, unit_id, slot_key, routing)
|
||||
if dispatch.get("status") == "dispatched":
|
||||
slot_key = slot_key or dispatch["delegation_id"]
|
||||
dispatched.append((unit, dispatch["delegation_id"]))
|
||||
continue
|
||||
if not dispatched:
|
||||
logger.info(
|
||||
"delegate_task: async pool at capacity (%s); running the whole batch synchronously instead.",
|
||||
dispatch.get("error", "rejected"),
|
||||
)
|
||||
return _run_sync_with_note(batch, "at_capacity")
|
||||
# Later units of an admitted call share its slot and cannot be capacity-rejected; a scheduler failure runs
|
||||
# the unit inline so no task is silently dropped.
|
||||
logger.warning("delegate_task: unit %d/%d not accepted (%s); running it inline.", k + 1, len(units), dispatch.get("error"))
|
||||
inline_results.extend(_execute_and_aggregate(unit, honor_parent_interrupt=False)["results"])
|
||||
payload = _dispatched_payload(batch, dispatched)
|
||||
if inline_results:
|
||||
payload["inline_results"] = inline_results
|
||||
return json.dumps(payload, ensure_ascii=False)
|
||||
|
||||
def _run_batch(batch: _Batch, background: bool) -> str:
|
||||
"""Tool result JSON: a dispatch handle (background) or the joined combined results."""
|
||||
|
||||
Reference in New Issue
Block a user