Skip to content

[Bug]: run_task discards already-recorded token usage when a sub-task fails after a successful call #3097

Description

@AmirF194

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

  1. 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.
  2. Call run_task(worker, ..., parent_context=parent_ctx).
  3. 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions