Describe the bug
run_task's exception handler discards a sub-task's already-recorded token usage instead of reporting it.
On the success path, TaskResult.usage (and the UsageEvent rolled up onto the parent stream) comes from (await reply.usage()).total, which aggregates every real UsageEvent already recorded on the child task_stream's history (ag2/agent.py, Agent.usage()). On the failure path (ag2/tools/subagents/run_task.py, the except Exception as e: block, ~line 136), TaskResult.usage is hardcoded to Usage() (all None) and no rollup UsageEvent is emitted, without ever reading the child stream's history.
So a sub-task that makes one or more successful (billable) model calls and then fails on a later step, timeout, tool error, or a bad final turn, loses that already-incurred usage entirely from TaskResult.usage and from the parent's UsageReport, even though the tokens were spent and the events recording them are sitting right there on result.stream.history.
Steps to reproduce
- Build an
Agent whose TestConfig scripts one successful ModelResponse (with usage= set, e.g. via a tool call) and then runs out of scripted responses on the follow-up call, so agent.ask() raises after the first call already succeeded.
- Call
run_task(worker, ..., parent_context=parent_ctx).
result.completed is False and result.usage is all-None, while result.stream.history.get_events() still contains the real UsageEvent/model-response usage from step 1.
Minimal reproducible example
import asyncio
from ag2 import Agent, Context, MemoryStream, tool
from ag2.events import ModelResponse, ToolCallEvent, ToolCallsEvent, Usage
from ag2.testing import TestConfig
from ag2.tools.subagents.run_task import run_task
from ag2.usage import UsageReport
@tool
def noop() -> str:
"""A tool that does nothing."""
return "ok"
async def main():
# First (and only) scripted response succeeds and is billable; the
# follow-up call needed to process the tool result has no scripted
# response left, so it fails.
worker_config = TestConfig(
ModelResponse(
tool_calls=ToolCallsEvent(calls=[ToolCallEvent(name="noop", arguments="{}")]),
usage=Usage(prompt_tokens=250, completion_tokens=50),
),
)
worker = Agent("worker", config=worker_config, tools=[noop])
parent_ctx = Context(stream=MemoryStream(), dependencies={}, variables={})
result = await run_task(worker, "do something", parent_context=parent_ctx)
print("completed:", result.completed)
print("result.usage:", result.usage)
child_report = UsageReport.from_events(list(await result.stream.history.get_events()))
print("child stream's actually-recorded usage:", child_report.total)
asyncio.run(main())
Output on current main (2e66147c44b1e982eac64833229f1842a815defd):
completed: False
result.usage: Usage(prompt_tokens=None, completion_tokens=None, total_tokens=None, cache_read_input_tokens=None, cache_creation_input_tokens=None, thinking_tokens=None)
child stream's actually-recorded usage: Usage(prompt_tokens=250, completion_tokens=None ... total_tokens=None, cache_read_input_tokens=None, cache_creation_input_tokens=None, thinking_tokens=None)
(the child report line above is trimmed for the issue; the full run shows prompt_tokens=250, completion_tokens=50 recorded and present, while result.usage reports none of it)
The existing suite (test/tools/test_task.py, 35/35 passing on this HEAD) has no test covering a failure that happens after a prior successful billable call; TestSubtaskUsageRollup only covers the success path and test_failure/test_failure_event fail on the very first call, so they never reach this state.
Model Used
None, reproduces entirely against the test double (ag2.testing.TestConfig), no live LLM call needed.
Expected Behavior
TaskResult.usage on a failed sub-task should reflect whatever usage the child stream actually recorded before the failure (e.g. UsageReport.from_events(await task_stream.history.get_events()).total), the same way the success path already computes it, instead of always reporting Usage().
Additional Information
- Package Name & Version: ag2, current
main at 2e66147c44b1e982eac64833229f1842a815defd (pyproject version 1.0.0b0)
- Operating System: Linux (reproduced in a clean
python:3.12-slim Docker container)
- Python Version: 3.12.3
- Happy to send a PR for this (read the child stream's history in the
except branch the same way the success path does) if useful.
Describe the bug
run_task's exception handler discards a sub-task's already-recorded token usage instead of reporting it.On the success path,
TaskResult.usage(and theUsageEventrolled up onto the parent stream) comes from(await reply.usage()).total, which aggregates every realUsageEventalready recorded on the child task_stream's history (ag2/agent.py,Agent.usage()). On the failure path (ag2/tools/subagents/run_task.py, theexcept Exception as e:block, ~line 136),TaskResult.usageis hardcoded toUsage()(allNone) and no rollupUsageEventis emitted, without ever reading the child stream's history.So a sub-task that makes one or more successful (billable) model calls and then fails on a later step, timeout, tool error, or a bad final turn, loses that already-incurred usage entirely from
TaskResult.usageand from the parent'sUsageReport, even though the tokens were spent and the events recording them are sitting right there onresult.stream.history.Steps to reproduce
AgentwhoseTestConfigscripts one successfulModelResponse(withusage=set, e.g. via a tool call) and then runs out of scripted responses on the follow-up call, soagent.ask()raises after the first call already succeeded.run_task(worker, ..., parent_context=parent_ctx).result.completedisFalseandresult.usageis all-None, whileresult.stream.history.get_events()still contains the realUsageEvent/model-response usage from step 1.Minimal reproducible example
Output on current
main(2e66147c44b1e982eac64833229f1842a815defd):(the child report line above is trimmed for the issue; the full run shows
prompt_tokens=250, completion_tokens=50recorded and present, whileresult.usagereports none of it)The existing suite (
test/tools/test_task.py, 35/35 passing on this HEAD) has no test covering a failure that happens after a prior successful billable call;TestSubtaskUsageRolluponly covers the success path andtest_failure/test_failure_eventfail on the very first call, so they never reach this state.Model Used
None, reproduces entirely against the test double (
ag2.testing.TestConfig), no live LLM call needed.Expected Behavior
TaskResult.usageon a failed sub-task should reflect whatever usage the child stream actually recorded before the failure (e.g.UsageReport.from_events(await task_stream.history.get_events()).total), the same way the success path already computes it, instead of always reportingUsage().Additional Information
mainat2e66147c44b1e982eac64833229f1842a815defd(pyproject version1.0.0b0)python:3.12-slimDocker container)exceptbranch the same way the success path does) if useful.