From 5185c1a002754118f091c2fadb1056c47241693a Mon Sep 17 00:00:00 2001
From: Mark Licata <malicata@microsoft.com>
Date: Wed, 7 Oct 2026 08:33:43 -0700
Subject: [PATCH 1/7] Python: add AgentBudget for autonomous loop resource
 limits

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../packages/core/agent_framework/__init__.py |   2 +
 .../core/agent_framework/__init__.pyi         |   2 +
 .../core/agent_framework/_harness/_loop.py    | 113 ++++++++++++++++++
 .../core/tests/core/test_harness_loop.py      |  95 +++++++++++++++
 4 files changed, 212 insertions(+)

diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py
index 3953187578e..efacbf476c3 100644
--- a/python/packages/core/agent_framework/__init__.py
+++ b/python/packages/core/agent_framework/__init__.py
@@ -132,6 +132,7 @@
         "FileMemoryProvider",
     ),
     "._harness._loop": (
+        "AgentBudget",
         "AgentLoopMiddleware",
         "JudgeVerdict",
         "background_tasks_running",
@@ -432,6 +433,7 @@
     "USER_AGENT_KEY",
     "USER_AGENT_TELEMETRY_DISABLED_ENV_VAR",
     "Agent",
+    "AgentBudget",
     "AgentContext",
     "AgentEvalConverter",
     "AgentExecutor",
diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi
index e8d1bc6e411..5ba07b7e58a 100644
--- a/python/packages/core/agent_framework/__init__.pyi
+++ b/python/packages/core/agent_framework/__init__.pyi
@@ -90,6 +90,7 @@ from ._harness._file_access import (
 )
 from ._harness._file_memory import DEFAULT_FILE_MEMORY_INSTRUCTIONS, DEFAULT_FILE_MEMORY_SOURCE_ID, FileMemoryProvider
 from ._harness._loop import (
+    AgentBudget,
     AgentLoopMiddleware,
     JudgeVerdict,
     background_tasks_running,
@@ -387,6 +388,7 @@ __all__ = [
     "USER_AGENT_KEY",
     "USER_AGENT_TELEMETRY_DISABLED_ENV_VAR",
     "Agent",
+    "AgentBudget",
     "AgentContext",
     "AgentEvalConverter",
     "AgentExecutor",
diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py
index b4617a23e01..762f82e55e4 100644
--- a/python/packages/core/agent_framework/_harness/_loop.py
+++ b/python/packages/core/agent_framework/_harness/_loop.py
@@ -29,7 +29,10 @@
 from __future__ import annotations
 
 import inspect
+import math
+import time
 from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, TypeAlias
 
 from pydantic import BaseModel, Field
@@ -57,6 +60,7 @@
     from .._clients import SupportsChatGetResponse
 
 __all__ = [
+    "AgentBudget",
     "AgentLoopMiddleware",
     "JudgeVerdict",
     "background_tasks_running",
@@ -131,6 +135,58 @@ class JudgeVerdict(BaseModel):
 DEFAULT_JUDGE_MAX_ITERATIONS = 5
 
 
+@experimental(feature_id=ExperimentalFeature.HARNESS)
+@dataclass
+class AgentBudget:
+    """Budget constraints for an :class:`AgentLoopMiddleware` autonomous loop.
+
+    Both ``max_tokens`` and ``max_duration`` are optional with no default — they are
+    workload-specific and must be set explicitly for production autonomous sessions.
+    When a budget dimension is ``None``, it is unconstrained.
+
+    Examples:
+
+        .. code-block:: python
+
+            from agent_framework import AgentBudget, AgentLoopMiddleware, todos_remaining
+
+            agent = Agent(
+                client=client,
+                middleware=[
+                    AgentLoopMiddleware(
+                        todos_remaining(),
+                        budget=AgentBudget(max_tokens=50_000, max_duration=300.0),
+                    )
+                ],
+            )
+
+    Note:
+        Budget stops stamp ``additional_properties["loop_exit_reason"]`` with
+        ``"token_budget_exceeded"`` or ``"time_budget_exceeded"`` on the returned
+        :class:`~agent_framework.AgentResponse` (non-streaming only).
+    """
+
+    max_tokens: int | None = None
+    """Maximum total LLM token spend across all loop iterations.
+
+    Counts ``total_token_count`` from each iteration's usage details, falling back to
+    ``input_token_count + output_token_count`` when ``total_token_count`` is absent.
+    """
+
+    max_duration: float | None = None
+    """Maximum wall-clock time in seconds for the entire loop run.
+
+    Checked at the start of each iteration using :func:`time.monotonic`.
+    Catches agents blocked on slow tools that token budgets cannot detect.
+    """
+
+    def __post_init__(self) -> None:
+        if self.max_tokens is not None and self.max_tokens < 1:
+            raise ValueError("AgentBudget.max_tokens must be a positive integer.")
+        if self.max_duration is not None and (not math.isfinite(self.max_duration) or self.max_duration <= 0):
+            raise ValueError("AgentBudget.max_duration must be a positive, finite number of seconds.")
+
+
 # A callable invoked between iterations. It always receives the loop keyword arguments
 # (``iteration``, ``last_result``, ``messages``, ``original_messages``, ``session``, ``agent``,
 # ``progress``, ``feedback``). Callers declare only the keywords they need plus ``**kwargs`` to
@@ -285,6 +341,7 @@ def __init__(
         fresh_context: bool = False,
         return_final_only: bool = False,
         additional_instructions: str | None = None,
+        budget: AgentBudget | None = None,
     ) -> None:
         """Initialize the agent loop middleware.
 
@@ -342,6 +399,9 @@ def __init__(
                 messages, so it is preserved across ``fresh_context`` resets and (with a session)
                 persists server-side across iterations. Used by :meth:`with_judge` to tell the agent
                 about the criteria its response must satisfy, but available to any loop.
+            budget: Optional :class:`AgentBudget` bounding total token spend and/or wall-clock time.
+                Checked after each iteration, before ``should_continue`` is evaluated. When exceeded
+                the loop stops and (non-streaming only) stamps ``loop_exit_reason``.
 
         Raises:
             ValueError: If ``max_iterations`` is not ``None`` and is less than 1.
@@ -357,6 +417,7 @@ def __init__(
         self.fresh_context = fresh_context
         self.return_final_only = return_final_only
         self.additional_instructions = additional_instructions
+        self.budget: AgentBudget | None = budget
 
     @classmethod
     def with_judge(
@@ -370,6 +431,7 @@ def with_judge(
         max_iterations: int | None = DEFAULT_JUDGE_MAX_ITERATIONS,
         next_message: NextMessageCallable | None = None,
         fresh_context: bool = False,
+        budget: AgentBudget | None = None,
     ) -> Self:
         """Create a loop that continues until a judge chat client decides the request was answered.
 
@@ -424,6 +486,7 @@ def with_judge(
                 conversation; an attached session is snapshotted before the loop and restored to that
                 baseline between iterations. See :meth:`__init__` for the full semantics. Defaults to
                 ``False``.
+            budget: Optional :class:`AgentBudget`. See :meth:`__init__`.
         """
         judge_instructions = (instructions or DEFAULT_JUDGE_INSTRUCTIONS).replace(
             CRITERIA_PLACEHOLDER, _render_criteria_block(criteria)
@@ -440,6 +503,7 @@ def with_judge(
             next_message=next_message or judge_next_message,
             fresh_context=fresh_context,
             additional_instructions=_criteria_agent_instruction(criteria) if criteria else None,
+            budget=budget,
         )
 
     async def process(
@@ -538,6 +602,7 @@ async def _process_non_streaming(
         snapshot: dict[str, Any] | None,
     ) -> None:
         iteration = 0
+        cap_fired = False
         work_iterations = 0
         progress: list[str] = []
         # Aggregated transcript across iterations: each iteration's response messages plus the
@@ -545,6 +610,9 @@ async def _process_non_streaming(
         aggregated: list[Message] = []
         aggregated_usage: UsageDetails | None = None
         final_result: AgentResponse | None = None
+        tokens_used: int = 0
+        budget_start: float = time.monotonic()
+        budget_exit: str | None = None
         stamped_options = dict(context.options) if context.options is not None else {}
         stamped_options[_LOOP_ITERATION_TOKEN_KEY] = object()
         context.options = stamped_options
@@ -565,6 +633,13 @@ async def _process_non_streaming(
                 if result.usage_details is not None:
                     aggregated_usage = add_usage_details(aggregated_usage, result.usage_details)
 
+                # Accumulate token spend for budget enforcement.
+                if self.budget is not None and self.budget.max_tokens is not None and result.usage_details is not None:
+                    ud = result.usage_details
+                    tokens_used += ud.get("total_token_count") or (
+                        (ud.get("input_token_count") or 0) + (ud.get("output_token_count") or 0)
+                    )
+
                 # Escape hatch: if this iteration is asking for tool approval, stop and return the
                 # response so the caller can approve, instead of continuing or injecting next_message.
                 if self._has_pending_approval_request(result):
@@ -581,9 +656,21 @@ async def _process_non_streaming(
                 )
 
                 work_iterations += 1
+                # Budget pre-check — runs before should_continue to avoid invoking an expensive
+                # predicate/judge after a budget has already been exhausted.
+                if self.budget is not None:
+                    elapsed = time.monotonic() - budget_start
+                    if self.budget.max_duration is not None and elapsed >= self.budget.max_duration:
+                        budget_exit = "time_budget_exceeded"
+                        break
+                    if self.budget.max_tokens is not None and tokens_used >= self.budget.max_tokens:
+                        budget_exit = "token_budget_exceeded"
+                        break
                 # Decide whether to stop and capture any feedback from should_continue first, so the
                 # feedback is available to both the progress and next-message callables this iteration.
                 stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations)
+                if stop and self.max_iterations is not None and work_iterations >= self.max_iterations:
+                    cap_fired = True
                 loop_kwargs = self._build_loop_kwargs(
                     context=context,
                     iteration=iteration,
@@ -617,6 +704,11 @@ async def _process_non_streaming(
         finally:
             context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None)
 
+        if cap_fired:
+            final_result.additional_properties.setdefault("loop_exit_reason", "iteration_cap_reached")
+        if budget_exit is not None:
+            final_result.additional_properties.setdefault("loop_exit_reason", budget_exit)
+
         if not self.return_final_only:
             context.result = self._aggregate_response(final_result, aggregated, aggregated_usage)
         await self._fire_turn_scoped_after_providers(
@@ -638,8 +730,11 @@ def _process_streaming(
 
         async def _generator() -> Any:
             iteration = 0
+            cap_fired = False
             work_iterations = 0
             progress: list[str] = []
+            tokens_used: int = 0
+            budget_start: float = time.monotonic()
             stamped_options = dict(context.options) if context.options is not None else {}
             stamped_options[_LOOP_ITERATION_TOKEN_KEY] = object()
             context.options = stamped_options
@@ -683,9 +778,25 @@ async def _generator() -> Any:
                     )
 
                     work_iterations += 1
+                    # Budget pre-check (before should_continue). Streaming cannot stamp
+                    # ``loop_exit_reason`` on the response (the caller assembles it from updates), so
+                    # the loop simply stops, consistent with the iteration-cap limitation.
+                    if self.budget is not None:
+                        if self.budget.max_tokens is not None and final.usage_details:
+                            ud = final.usage_details
+                            tokens_used += ud.get("total_token_count") or (
+                                (ud.get("input_token_count") or 0) + (ud.get("output_token_count") or 0)
+                            )
+                        elapsed = time.monotonic() - budget_start
+                        if self.budget.max_duration is not None and elapsed >= self.budget.max_duration:
+                            return
+                        if self.budget.max_tokens is not None and tokens_used >= self.budget.max_tokens:
+                            return
                     # Decide whether to stop and capture any feedback from should_continue first, so the
                     # feedback is available to both the progress and next-message callables this iteration.
                     stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations)
+                    if stop and self.max_iterations is not None and work_iterations >= self.max_iterations:
+                        cap_fired = True
                     loop_kwargs = self._build_loop_kwargs(
                         context=context,
                         iteration=iteration,
@@ -721,6 +832,8 @@ async def _generator() -> Any:
                         yield self._message_to_update(message)
             finally:
                 context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None)
+                if cap_fired and holder["final"] is not None:
+                    holder["final"].additional_properties.setdefault("loop_exit_reason", "iteration_cap_reached")
                 await self._fire_turn_scoped_after_providers(context, holder["final"], original_messages)
 
         def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse:
diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py
index 3b87544d804..be2b392b271 100644
--- a/python/packages/core/tests/core/test_harness_loop.py
+++ b/python/packages/core/tests/core/test_harness_loop.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import time
 from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
 from typing import Any
 
@@ -10,6 +11,7 @@
 
 from agent_framework import (
     Agent,
+    AgentBudget,
     AgentContext,
     AgentMiddleware,
     AgentModeProvider,
@@ -29,6 +31,7 @@
     ResponseStream,
     TodoItem,
     TodoProvider,
+    UsageDetails,
     background_tasks_running,
     background_tasks_running_message,
     set_agent_mode,
@@ -61,8 +64,10 @@ def __init__(
         texts: list[str] | None = None,
         honor_response_format: bool = False,
         service_mode: bool = False,
+        usage_per_call: int | None = None,
     ) -> None:
         super().__init__()
+        self._usage_per_call = usage_per_call
         self.call_count: int = 0
         self.received_messages: list[list[str]] = []
         self.received_response_formats: list[Any] = []
@@ -108,6 +113,9 @@ async def _get() -> ChatResponse:
                 messages=Message(role="assistant", contents=[self._next_text(messages)]),
                 response_format=response_format if self._honor_response_format else None,
                 conversation_id=conversation_id,
+                usage_details=(
+                    UsageDetails(total_token_count=self._usage_per_call) if self._usage_per_call is not None else None
+                ),
             )
 
         return _get()
@@ -1697,3 +1705,90 @@ async def test_concurrent_same_agent_run_during_stream_pause_is_not_suppressed()
 
     # once for the loop boundary, once for the concurrent run's own boundary
     assert turn_scoped.after_calls == 2
+
+
+# region budget
+
+
+async def test_budget_stops_on_token_limit() -> None:
+    client = RecordingChatClient(usage_per_call=50)
+    agent = Agent(
+        client=client,
+        middleware=[AgentLoopMiddleware(always_continue, max_iterations=10, budget=AgentBudget(max_tokens=100))],
+    )
+
+    response = await agent.run("task")
+
+    assert client.call_count == 2
+    assert response.additional_properties.get("loop_exit_reason") == "token_budget_exceeded"
+
+
+async def test_budget_stops_on_duration_limit() -> None:
+    class _SlowClient(RecordingChatClient):
+        def _inner_get_response(self, **kwargs: Any) -> Any:
+            time.sleep(0.05)  # exceed the 1ms budget (Windows clock tick is ~15ms)
+            return super()._inner_get_response(**kwargs)
+
+    client = _SlowClient()
+    agent = Agent(
+        client=client,
+        middleware=[AgentLoopMiddleware(always_continue, max_iterations=10, budget=AgentBudget(max_duration=0.001))],
+    )
+
+    response = await agent.run("task")
+
+    assert client.call_count == 1
+    assert response.additional_properties.get("loop_exit_reason") == "time_budget_exceeded"
+
+
+async def test_budget_not_triggered_when_none() -> None:
+    client = RecordingChatClient(usage_per_call=50)
+    agent = Agent(client=client, middleware=[AgentLoopMiddleware(always_continue, max_iterations=3)])
+
+    response = await agent.run("task")
+
+    assert client.call_count == 3
+    assert response.additional_properties.get("loop_exit_reason") == "iteration_cap_reached"
+
+
+async def test_budget_token_check_before_should_continue() -> None:
+    client = RecordingChatClient(usage_per_call=50)
+    calls: list[int] = []
+
+    def should_continue(*, iteration: int, **kwargs: Any) -> bool:
+        calls.append(iteration)
+        return True
+
+    agent = Agent(
+        client=client,
+        middleware=[AgentLoopMiddleware(should_continue, max_iterations=10, budget=AgentBudget(max_tokens=100))],
+    )
+
+    await agent.run("task")
+
+    assert client.call_count == 2
+    # should_continue ran only for iteration 1; the budget stop pre-empted it on iteration 2.
+    assert calls == [1]
+
+
+async def test_budget_exit_reason_in_additional_properties() -> None:
+    client = RecordingChatClient(usage_per_call=60)
+    agent = Agent(
+        client=client,
+        middleware=[AgentLoopMiddleware(always_continue, budget=AgentBudget(max_tokens=50))],
+    )
+
+    response = await agent.run("task")
+
+    assert client.call_count == 1
+    assert response.additional_properties["loop_exit_reason"] == "token_budget_exceeded"
+
+
+def test_agent_budget_validates_max_tokens() -> None:
+    with pytest.raises(ValueError, match="max_tokens"):
+        AgentBudget(max_tokens=0)
+
+
+def test_agent_budget_validates_max_duration() -> None:
+    with pytest.raises(ValueError, match="max_duration"):
+        AgentBudget(max_duration=-1.0)

From 552d4b17ec20b0be3c51999aab8ed9054257436f Mon Sep 17 00:00:00 2001
From: Mark Licata <malicata@microsoft.com>
Date: Wed, 7 Oct 2026 09:00:05 -0700
Subject: [PATCH 2/7] Address review feedback: clarify max_duration semantics,
 test token fallback

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../core/agent_framework/_harness/_loop.py    | 11 ++--
 .../core/tests/core/test_harness_loop.py      | 52 +++++++++++++++++--
 2 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py
index 762f82e55e4..b7463a83061 100644
--- a/python/packages/core/agent_framework/_harness/_loop.py
+++ b/python/packages/core/agent_framework/_harness/_loop.py
@@ -176,8 +176,11 @@ class AgentBudget:
     max_duration: float | None = None
     """Maximum wall-clock time in seconds for the entire loop run.
 
-    Checked at the start of each iteration using :func:`time.monotonic`.
-    Catches agents blocked on slow tools that token budgets cannot detect.
+    Measured with :func:`time.monotonic` and checked between iterations, after each
+    iteration completes: once exceeded, no new iteration is started. It does not
+    interrupt an in-flight model or tool call, so a single slow or hung call can run
+    past the deadline. Callers that need a hard deadline should wrap the run in their
+    own cancellation/timeout (for example ``asyncio.timeout``).
     """
 
     def __post_init__(self) -> None:
@@ -401,7 +404,9 @@ def __init__(
                 about the criteria its response must satisfy, but available to any loop.
             budget: Optional :class:`AgentBudget` bounding total token spend and/or wall-clock time.
                 Checked after each iteration, before ``should_continue`` is evaluated. When exceeded
-                the loop stops and (non-streaming only) stamps ``loop_exit_reason``.
+                the loop stops and (non-streaming only) stamps ``loop_exit_reason``. Budgets do not
+                interrupt an in-flight model/tool call; use ``asyncio.timeout`` (or similar) around
+                the run for a hard wall-clock deadline.
 
         Raises:
             ValueError: If ``max_iterations`` is not ``None`` and is less than 1.
diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py
index be2b392b271..f0a74ecda6e 100644
--- a/python/packages/core/tests/core/test_harness_loop.py
+++ b/python/packages/core/tests/core/test_harness_loop.py
@@ -65,9 +65,11 @@ def __init__(
         honor_response_format: bool = False,
         service_mode: bool = False,
         usage_per_call: int | None = None,
+        usage_details: UsageDetails | None = None,
     ) -> None:
         super().__init__()
         self._usage_per_call = usage_per_call
+        self._usage_details = usage_details
         self.call_count: int = 0
         self.received_messages: list[list[str]] = []
         self.received_response_formats: list[Any] = []
@@ -79,6 +81,13 @@ def __init__(
             object.__setattr__(self, "STORES_BY_DEFAULT", True)
         self._conv_counter = 0
 
+    def _usage(self) -> UsageDetails | None:
+        if self._usage_details is not None:
+            return UsageDetails(**self._usage_details)
+        if self._usage_per_call is not None:
+            return UsageDetails(total_token_count=self._usage_per_call)
+        return None
+
     def _next_text(self, messages: Sequence[Message]) -> str:
         if self._texts:
             return self._texts.pop(0)
@@ -113,9 +122,7 @@ async def _get() -> ChatResponse:
                 messages=Message(role="assistant", contents=[self._next_text(messages)]),
                 response_format=response_format if self._honor_response_format else None,
                 conversation_id=conversation_id,
-                usage_details=(
-                    UsageDetails(total_token_count=self._usage_per_call) if self._usage_per_call is not None else None
-                ),
+                usage_details=self._usage(),
             )
 
         return _get()
@@ -126,8 +133,12 @@ def _stream(
         async def _gen() -> AsyncIterable[ChatResponseUpdate]:
             self.call_count += 1
             text = self._next_text(messages)
+            contents = [Content.from_text(text)]
+            usage = self._usage()
+            if usage is not None:
+                contents.append(Content.from_usage(usage))
             yield ChatResponseUpdate(
-                contents=[Content.from_text(text)],
+                contents=contents,
                 role="assistant",
                 finish_reason="stop",
                 conversation_id=conversation_id,
@@ -1723,7 +1734,38 @@ async def test_budget_stops_on_token_limit() -> None:
     assert response.additional_properties.get("loop_exit_reason") == "token_budget_exceeded"
 
 
-async def test_budget_stops_on_duration_limit() -> None:
+async def test_budget_token_limit_falls_back_to_input_plus_output_counts() -> None:
+    # No total_token_count: spend is input + output (30 + 20 = 50 per call), so 100 is hit on call 2.
+    client = RecordingChatClient(usage_details=UsageDetails(input_token_count=30, output_token_count=20))
+    agent = Agent(
+        client=client,
+        middleware=[AgentLoopMiddleware(always_continue, max_iterations=10, budget=AgentBudget(max_tokens=100))],
+    )
+
+    response = await agent.run("task")
+
+    assert client.call_count == 2
+    assert response.additional_properties.get("loop_exit_reason") == "token_budget_exceeded"
+
+
+async def test_budget_token_limit_falls_back_to_input_plus_output_counts_streaming() -> None:
+    client = RecordingChatClient(usage_details=UsageDetails(input_token_count=30, output_token_count=20))
+    agent = Agent(
+        client=client,
+        middleware=[AgentLoopMiddleware(always_continue, max_iterations=10, budget=AgentBudget(max_tokens=100))],
+    )
+
+    stream = agent.run("task", stream=True)
+    async for _ in stream:
+        pass
+    await stream.get_final_response()
+
+    assert client.call_count == 2
+
+
+async def test_budget_duration_stops_before_next_iteration() -> None:
+    # max_duration is checked between iterations: the in-flight call runs to completion,
+    # then no further iteration is started.
     class _SlowClient(RecordingChatClient):
         def _inner_get_response(self, **kwargs: Any) -> Any:
             time.sleep(0.05)  # exceed the 1ms budget (Windows clock tick is ~15ms)

From 84fdeadf2d18f69768df4e11579bd3f56f44b229 Mon Sep 17 00:00:00 2001
From: Mark Licata <malicata@microsoft.com>
Date: Wed, 7 Oct 2026 09:21:15 -0700
Subject: [PATCH 3/7] Fix zuban error in RecordingChatClient usage helper

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 python/packages/core/tests/core/test_harness_loop.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py
index f0a74ecda6e..52467ce357f 100644
--- a/python/packages/core/tests/core/test_harness_loop.py
+++ b/python/packages/core/tests/core/test_harness_loop.py
@@ -83,7 +83,7 @@ def __init__(
 
     def _usage(self) -> UsageDetails | None:
         if self._usage_details is not None:
-            return UsageDetails(**self._usage_details)
+            return self._usage_details.copy()
         if self._usage_per_call is not None:
             return UsageDetails(total_token_count=self._usage_per_call)
         return None

From 06d82d6e0b70a99080a3873987bfb96c7adf9d93 Mon Sep 17 00:00:00 2001
From: Mark Licata <malicata@microsoft.com>
Date: Tue, 13 Oct 2026 08:40:50 -0700
Subject: [PATCH 4/7] Python: Make AgentBudget a TypedDict

Convert AgentBudget from a dataclass to a total=False TypedDict so it can be
used for typing without importing it, and plain dicts work as budgets.
Validation moves into AgentLoopMiddleware.__init__ with the same ValueError
messages.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../core/agent_framework/_harness/_loop.py    | 58 ++++++++++---------
 .../core/tests/core/test_harness_loop.py      | 17 +++++-
 2 files changed, 45 insertions(+), 30 deletions(-)

diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py
index b7463a83061..00d9b53fbb9 100644
--- a/python/packages/core/agent_framework/_harness/_loop.py
+++ b/python/packages/core/agent_framework/_harness/_loop.py
@@ -32,11 +32,10 @@
 import math
 import time
 from collections.abc import Awaitable, Callable, Mapping, Sequence
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, TypeAlias
 
 from pydantic import BaseModel, Field
-from typing_extensions import Self
+from typing_extensions import Self, TypedDict
 
 from .._agents import (
     _LOOP_ITERATION_TOKEN_KEY,  # pyright: ignore[reportPrivateUsage] -- shared loop-turn marker, see _agents.py
@@ -136,26 +135,27 @@ class JudgeVerdict(BaseModel):
 
 
 @experimental(feature_id=ExperimentalFeature.HARNESS)
-@dataclass
-class AgentBudget:
+class AgentBudget(TypedDict, total=False):
     """Budget constraints for an :class:`AgentLoopMiddleware` autonomous loop.
 
     Both ``max_tokens`` and ``max_duration`` are optional with no default — they are
     workload-specific and must be set explicitly for production autonomous sessions.
-    When a budget dimension is ``None``, it is unconstrained.
+    When a budget dimension is omitted or ``None``, it is unconstrained. Because this is a
+    ``TypedDict``, a plain ``dict`` works too (for example ``budget={"max_tokens": 50_000}``).
+    Values are validated by :class:`AgentLoopMiddleware` when it is constructed.
 
     Examples:
 
         .. code-block:: python
 
-            from agent_framework import AgentBudget, AgentLoopMiddleware, todos_remaining
+            from agent_framework import AgentLoopMiddleware, todos_remaining
 
             agent = Agent(
                 client=client,
                 middleware=[
                     AgentLoopMiddleware(
                         todos_remaining(),
-                        budget=AgentBudget(max_tokens=50_000, max_duration=300.0),
+                        budget={"max_tokens": 50_000, "max_duration": 300.0},
                     )
                 ],
             )
@@ -166,14 +166,14 @@ class AgentBudget:
         :class:`~agent_framework.AgentResponse` (non-streaming only).
     """
 
-    max_tokens: int | None = None
+    max_tokens: int | None
     """Maximum total LLM token spend across all loop iterations.
 
     Counts ``total_token_count`` from each iteration's usage details, falling back to
     ``input_token_count + output_token_count`` when ``total_token_count`` is absent.
     """
 
-    max_duration: float | None = None
+    max_duration: float | None
     """Maximum wall-clock time in seconds for the entire loop run.
 
     Measured with :func:`time.monotonic` and checked between iterations, after each
@@ -183,12 +183,6 @@ class AgentBudget:
     own cancellation/timeout (for example ``asyncio.timeout``).
     """
 
-    def __post_init__(self) -> None:
-        if self.max_tokens is not None and self.max_tokens < 1:
-            raise ValueError("AgentBudget.max_tokens must be a positive integer.")
-        if self.max_duration is not None and (not math.isfinite(self.max_duration) or self.max_duration <= 0):
-            raise ValueError("AgentBudget.max_duration must be a positive, finite number of seconds.")
-
 
 # A callable invoked between iterations. It always receives the loop keyword arguments
 # (``iteration``, ``last_result``, ``messages``, ``original_messages``, ``session``, ``agent``,
@@ -402,17 +396,25 @@ def __init__(
                 messages, so it is preserved across ``fresh_context`` resets and (with a session)
                 persists server-side across iterations. Used by :meth:`with_judge` to tell the agent
                 about the criteria its response must satisfy, but available to any loop.
-            budget: Optional :class:`AgentBudget` bounding total token spend and/or wall-clock time.
-                Checked after each iteration, before ``should_continue`` is evaluated. When exceeded
-                the loop stops and (non-streaming only) stamps ``loop_exit_reason``. Budgets do not
-                interrupt an in-flight model/tool call; use ``asyncio.timeout`` (or similar) around
-                the run for a hard wall-clock deadline.
+            budget: Optional :class:`AgentBudget` (or a plain ``dict`` with the same keys) bounding total
+                token spend and/or wall-clock time. Checked after each iteration, before
+                ``should_continue`` is evaluated. When exceeded the loop stops and (non-streaming
+                only) stamps ``loop_exit_reason``. Budgets do not interrupt an in-flight model/tool
+                call; use ``asyncio.timeout`` (or similar) around the run for a hard wall-clock
+                deadline.
 
         Raises:
-            ValueError: If ``max_iterations`` is not ``None`` and is less than 1.
+            ValueError: If ``max_iterations`` is not ``None`` and is less than 1, or if ``budget``
+                contains a non-positive ``max_tokens`` or a non-positive/non-finite ``max_duration``.
         """
         if max_iterations is not None and max_iterations < 1:
             raise ValueError("max_iterations must be None or a positive integer (>= 1).")
+        budget_max_tokens = budget.get("max_tokens") if budget is not None else None
+        budget_max_duration = budget.get("max_duration") if budget is not None else None
+        if budget_max_tokens is not None and budget_max_tokens < 1:
+            raise ValueError("AgentBudget.max_tokens must be a positive integer.")
+        if budget_max_duration is not None and (not math.isfinite(budget_max_duration) or budget_max_duration <= 0):
+            raise ValueError("AgentBudget.max_duration must be a positive, finite number of seconds.")
 
         self.max_iterations: int | None = max_iterations
         self.should_continue: ShouldContinueCallable = should_continue
@@ -423,6 +425,8 @@ def __init__(
         self.return_final_only = return_final_only
         self.additional_instructions = additional_instructions
         self.budget: AgentBudget | None = budget
+        self._budget_max_tokens: int | None = budget_max_tokens
+        self._budget_max_duration: float | None = budget_max_duration
 
     @classmethod
     def with_judge(
@@ -639,7 +643,7 @@ async def _process_non_streaming(
                     aggregated_usage = add_usage_details(aggregated_usage, result.usage_details)
 
                 # Accumulate token spend for budget enforcement.
-                if self.budget is not None and self.budget.max_tokens is not None and result.usage_details is not None:
+                if self._budget_max_tokens is not None and result.usage_details is not None:
                     ud = result.usage_details
                     tokens_used += ud.get("total_token_count") or (
                         (ud.get("input_token_count") or 0) + (ud.get("output_token_count") or 0)
@@ -665,10 +669,10 @@ async def _process_non_streaming(
                 # predicate/judge after a budget has already been exhausted.
                 if self.budget is not None:
                     elapsed = time.monotonic() - budget_start
-                    if self.budget.max_duration is not None and elapsed >= self.budget.max_duration:
+                    if self._budget_max_duration is not None and elapsed >= self._budget_max_duration:
                         budget_exit = "time_budget_exceeded"
                         break
-                    if self.budget.max_tokens is not None and tokens_used >= self.budget.max_tokens:
+                    if self._budget_max_tokens is not None and tokens_used >= self._budget_max_tokens:
                         budget_exit = "token_budget_exceeded"
                         break
                 # Decide whether to stop and capture any feedback from should_continue first, so the
@@ -787,15 +791,15 @@ async def _generator() -> Any:
                     # ``loop_exit_reason`` on the response (the caller assembles it from updates), so
                     # the loop simply stops, consistent with the iteration-cap limitation.
                     if self.budget is not None:
-                        if self.budget.max_tokens is not None and final.usage_details:
+                        if self._budget_max_tokens is not None and final.usage_details:
                             ud = final.usage_details
                             tokens_used += ud.get("total_token_count") or (
                                 (ud.get("input_token_count") or 0) + (ud.get("output_token_count") or 0)
                             )
                         elapsed = time.monotonic() - budget_start
-                        if self.budget.max_duration is not None and elapsed >= self.budget.max_duration:
+                        if self._budget_max_duration is not None and elapsed >= self._budget_max_duration:
                             return
-                        if self.budget.max_tokens is not None and tokens_used >= self.budget.max_tokens:
+                        if self._budget_max_tokens is not None and tokens_used >= self._budget_max_tokens:
                             return
                     # Decide whether to stop and capture any feedback from should_continue first, so the
                     # feedback is available to both the progress and next-message callables this iteration.
diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py
index 52467ce357f..fe6464dc03f 100644
--- a/python/packages/core/tests/core/test_harness_loop.py
+++ b/python/packages/core/tests/core/test_harness_loop.py
@@ -1817,7 +1817,7 @@ async def test_budget_exit_reason_in_additional_properties() -> None:
     client = RecordingChatClient(usage_per_call=60)
     agent = Agent(
         client=client,
-        middleware=[AgentLoopMiddleware(always_continue, budget=AgentBudget(max_tokens=50))],
+        middleware=[AgentLoopMiddleware(always_continue, budget={"max_tokens": 50})],
     )
 
     response = await agent.run("task")
@@ -1828,9 +1828,20 @@ async def test_budget_exit_reason_in_additional_properties() -> None:
 
 def test_agent_budget_validates_max_tokens() -> None:
     with pytest.raises(ValueError, match="max_tokens"):
-        AgentBudget(max_tokens=0)
+        AgentLoopMiddleware(always_continue, budget=AgentBudget(max_tokens=0))
+    with pytest.raises(ValueError, match="max_tokens"):
+        AgentLoopMiddleware(always_continue, budget={"max_tokens": 0})
 
 
 def test_agent_budget_validates_max_duration() -> None:
     with pytest.raises(ValueError, match="max_duration"):
-        AgentBudget(max_duration=-1.0)
+        AgentLoopMiddleware(always_continue, budget=AgentBudget(max_duration=-1.0))
+    with pytest.raises(ValueError, match="max_duration"):
+        AgentLoopMiddleware(always_continue, budget={"max_duration": float("inf")})
+
+
+def test_agent_budget_accepts_empty_and_none_values() -> None:
+    middleware = AgentLoopMiddleware(always_continue, budget={})
+    assert middleware.budget == {}
+    middleware = AgentLoopMiddleware(always_continue, budget={"max_tokens": None, "max_duration": None})
+    assert middleware.budget == {"max_tokens": None, "max_duration": None}

From 63004d79b2fb21b3091b52489d9c15e7f1250c0e Mon Sep 17 00:00:00 2001
From: Mark Licata <malicata@microsoft.com>
Date: Tue, 13 Oct 2026 08:42:45 -0700
Subject: [PATCH 5/7] Python: Document AgentBudget max_tokens accounting, test
 judge exclusion

max_tokens charges the wrapped agent's total_token_count per iteration
(falling back to input + output). The with_judge judge call is never
charged. Sub-agent usage counts only when it is reported in the agent's
response usage_details; as_tool and background agents currently return
text only, so their usage is not included. Not configurable.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 python/packages/core/AGENTS.md                |  1 +
 .../core/agent_framework/_harness/_loop.py    | 26 +++++++++++---
 .../core/tests/core/test_harness_loop.py      | 35 +++++++++++++++++++
 3 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md
index 6f6913b71d0..5e9fbff46fe 100644
--- a/python/packages/core/AGENTS.md
+++ b/python/packages/core/AGENTS.md
@@ -309,6 +309,7 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID
 - **`todos_remaining(*, looping_modes=None)`** / **`todos_remaining_message`** - Helper factories for todo-driven loops (the Python counterpart of .NET's `TodoCompletionLoopEvaluator`), designed for `create_harness_agent` but usable with any agent that registers a `TodoProvider` via `context_providers`. They resolve the `TodoProvider`/`AgentModeProvider` from the *running agent* (`agent.context_providers`, via `_resolve_context_provider`) rather than taking the provider as an argument, so they can be wired directly into `loop_should_continue`/`loop_next_message`. `todos_remaining` returns a `should_continue` predicate that loops while any todo is open; pass `looping_modes=[...]` to gate looping to specific operating modes (case-insensitive; honors the `AgentModeProvider`'s `source_id`/`available_modes`), `looping_modes=None` (default) applies in every mode, and an empty sequence raises `ValueError`. `todos_remaining_message` is a `next_message` callable that lists the still-open todo titles and tells the agent to finish them, returning `None` when the session/agent/provider is unavailable or nothing is open (in which case the middleware's default `None` handling applies: reuse the previous iteration's messages verbatim under the default `fresh_context=False`, or `DEFAULT_NEXT_MESSAGE` only when `fresh_context=True`).
 - **`background_tasks_running()`** / **`background_tasks_running_message`** - Helper factories for background-agent-driven loops, mirroring the `todos_remaining` pair. They resolve the `BackgroundAgentsProvider` from the *running agent* (`agent.context_providers`, via `_resolve_context_provider`) rather than taking the provider as an argument, so they can be wired directly into `create_harness_agent`'s `loop_should_continue`/`loop_next_message`. `background_tasks_running` returns a `should_continue` predicate that loops while the provider's persisted state shows any task with `status == RUNNING` (pair it with `max_iterations` so the loop is bounded even if a task's persisted status is never refreshed). `background_tasks_running_message` is a `next_message` callable that lists the still-running tasks (`#<id> (<agent_name>): <description>`) and tells the agent to wait for them to finish and retrieve their results, returning `None` when the session/agent/provider is unavailable or no task is running.
   - **Approval escape hatch** - `_has_pending_approval_request(result)` checks whether an iteration's response carries a pending tool-approval request (any content with `type == "function_approval_request"`). Both the streaming and non-streaming loops stop and return that response to the caller *before* evaluating `should_continue`/`max_iterations` or injecting `next_message`, so the loop is HITL-safe even when wrapped outermost around a `ToolApprovalMiddleware` (mirrors the C# `LoopAgent`'s `HasPendingApprovalRequests`).
+  - **Budget** - `budget=AgentBudget` (a `total=False` `TypedDict`; plain dicts work) adds `max_tokens` / `max_duration` limits, validated in `AgentLoopMiddleware.__init__` and checked after each iteration before `should_continue`. `max_tokens` charges each iteration the wrapped agent response's `total_token_count`, falling back to `input_token_count + output_token_count`; the `with_judge` judge call is never charged. Sub-agent usage counts only when it appears in the agent's response `usage_details` (`as_tool` and background agents return text only, so their usage is currently not included). The accounting is not configurable.
   - **Harness integration** - `create_harness_agent` enables the loop when a `loop_should_continue` callable is passed; it prepends `AgentLoopMiddleware(loop_should_continue, max_iterations=loop_max_iterations, next_message=loop_next_message)` ahead of `ToolApprovalMiddleware` so the loop is the outermost middleware (each iteration is a full agent run including tool approval, and the escape hatch hands pending approvals back to the caller). `loop_next_message` and `loop_max_iterations` only take effect together with `loop_should_continue` (with no `loop_should_continue` there is no loop, so they are ignored); `loop_max_iterations` defaults to the loop's default cap (`None` → unbounded).
 
 ### Workflows (`_workflows/`)
diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py
index 00d9b53fbb9..f75e2a25fc3 100644
--- a/python/packages/core/agent_framework/_harness/_loop.py
+++ b/python/packages/core/agent_framework/_harness/_loop.py
@@ -167,10 +167,21 @@ class AgentBudget(TypedDict, total=False):
     """
 
     max_tokens: int | None
-    """Maximum total LLM token spend across all loop iterations.
+    """Maximum token spend (input + output) of the wrapped agent across all loop iterations.
 
-    Counts ``total_token_count`` from each iteration's usage details, falling back to
-    ``input_token_count + output_token_count`` when ``total_token_count`` is absent.
+    Each iteration is charged the ``total_token_count`` from the agent response's usage details,
+    falling back to ``input_token_count + output_token_count`` when ``total_token_count`` is
+    absent. That usage covers every model call the agent makes during the iteration, including
+    each round of tool calling.
+
+    Calls the loop framework makes on its own behalf are never charged: the judge call made by
+    :meth:`AgentLoopMiddleware.with_judge` does not count toward this budget.
+
+    Sub-agents are only counted when their usage is reported in the wrapped agent's response
+    usage details. Agents exposed as tools via :meth:`~agent_framework.Agent.as_tool` and
+    background agents return only text, so their token usage is currently *not* included.
+
+    This accounting is fixed and not configurable.
     """
 
     max_duration: float | None
@@ -397,7 +408,11 @@ def __init__(
                 persists server-side across iterations. Used by :meth:`with_judge` to tell the agent
                 about the criteria its response must satisfy, but available to any loop.
             budget: Optional :class:`AgentBudget` (or a plain ``dict`` with the same keys) bounding total
-                token spend and/or wall-clock time. Checked after each iteration, before
+                token spend and/or wall-clock time. ``max_tokens`` counts the wrapped agent's
+                ``total_token_count`` (falling back to input + output tokens) per iteration; the
+                judge call made by :meth:`with_judge` is never counted, and sub-agent usage counts
+                only when it is reported in the agent's response usage details (see
+                :attr:`AgentBudget.max_tokens`). Not configurable. Checked after each iteration, before
                 ``should_continue`` is evaluated. When exceeded the loop stops and (non-streaming
                 only) stamps ``loop_exit_reason``. Budgets do not interrupt an in-flight model/tool
                 call; use ``asyncio.timeout`` (or similar) around the run for a hard wall-clock
@@ -495,7 +510,8 @@ def with_judge(
                 conversation; an attached session is snapshotted before the loop and restored to that
                 baseline between iterations. See :meth:`__init__` for the full semantics. Defaults to
                 ``False``.
-            budget: Optional :class:`AgentBudget`. See :meth:`__init__`.
+            budget: Optional :class:`AgentBudget`. See :meth:`__init__`. The judge call's own token
+                usage is never charged against ``max_tokens``.
         """
         judge_instructions = (instructions or DEFAULT_JUDGE_INSTRUCTIONS).replace(
             CRITERIA_PLACEHOLDER, _render_criteria_block(criteria)
diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py
index fe6464dc03f..9322dbf964c 100644
--- a/python/packages/core/tests/core/test_harness_loop.py
+++ b/python/packages/core/tests/core/test_harness_loop.py
@@ -1826,6 +1826,41 @@ async def test_budget_exit_reason_in_additional_properties() -> None:
     assert response.additional_properties["loop_exit_reason"] == "token_budget_exceeded"
 
 
+async def test_budget_excludes_judge_call_usage() -> None:
+    # The agent spends 40 tokens per iteration; the judge spends 1000 per call and never says done.
+    # If judge usage were charged, the 100-token budget would trip after iteration 1's judge call
+    # (stopping at 2 agent calls). Only agent usage counts, so the budget trips on iteration 3.
+    agent_client = RecordingChatClient(usage_per_call=40)
+    judge_client = RecordingChatClient(texts=["VERDICT: MORE"] * 5, usage_per_call=1000)
+    agent = Agent(
+        client=agent_client,
+        middleware=[AgentLoopMiddleware.with_judge(judge_client, max_iterations=10, budget={"max_tokens": 100})],
+    )
+
+    response = await agent.run("task")
+
+    assert agent_client.call_count == 3
+    assert judge_client.call_count == 2
+    assert response.additional_properties["loop_exit_reason"] == "token_budget_exceeded"
+
+
+async def test_budget_excludes_judge_call_usage_streaming() -> None:
+    agent_client = RecordingChatClient(usage_per_call=40)
+    judge_client = RecordingChatClient(texts=["VERDICT: MORE"] * 5, usage_per_call=1000)
+    agent = Agent(
+        client=agent_client,
+        middleware=[AgentLoopMiddleware.with_judge(judge_client, max_iterations=10, budget={"max_tokens": 100})],
+    )
+
+    stream = agent.run("task", stream=True)
+    async for _ in stream:
+        pass
+    await stream.get_final_response()
+
+    assert agent_client.call_count == 3
+    assert judge_client.call_count == 2
+
+
 def test_agent_budget_validates_max_tokens() -> None:
     with pytest.raises(ValueError, match="max_tokens"):
         AgentLoopMiddleware(always_continue, budget=AgentBudget(max_tokens=0))

From 3a91bbda659c2a46fda05a70abba57d44a014ba1 Mon Sep 17 00:00:00 2001
From: Mark Licata <malicata@microsoft.com>
Date: Tue, 13 Oct 2026 10:02:06 -0700
Subject: [PATCH 6/7] Python: Count sub-agent token usage toward
 AgentBudget.max_tokens

Sub-agent usage now rolls into the parent's response usage_details, streaming
and non-streaming, so AgentBudget.max_tokens includes it.

- Agent.as_tool reports the sub-agent's usage_details through a private
  FunctionInvocationContext metadata key (also before the approval/user-input
  raises). The function invocation loop moves it onto the function_result and
  folds it into the parent response exactly once (Phase 1/2/3 non-streaming,
  a usage content on the tool update when streaming), including the
  error-result, middleware-termination and user-input exception paths.
  Nested A->B->C counts each layer once.
- BackgroundAgentsProvider stores each finished run's usage on
  BackgroundTaskInfo (usage_details/usage_reported, serialized) and charges it
  once, to the parent run that first observes the completion: through the
  observing tool's context, or a run-scoped chat middleware on the run's first
  model response when before_run observes it.
- invoke_agent spans keep reporting each agent's own usage: the
  response-usage fallback subtracts delegated usage.
- Documented gaps: failed/cancelled sub-agent runs, background tasks never
  observed or abandoned at release, and custom tools calling agent.run
  directly are not counted.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 python/packages/core/AGENTS.md                |   2 +-
 .../packages/core/agent_framework/_agents.py  |  11 +
 .../_harness/_background_agents.py            | 160 +++++++++++++-
 .../core/agent_framework/_harness/_loop.py    |  19 +-
 .../packages/core/agent_framework/_tools.py   | 176 +++++++++++++--
 .../core/agent_framework/observability.py     |  54 ++++-
 .../packages/core/tests/core/test_agents.py   | 206 ++++++++++++++++++
 .../core/test_function_invocation_logic.py    |  83 +++++++
 .../core/test_harness_background_agents.py    | 148 +++++++++++++
 .../core/tests/core/test_harness_loop.py      |  48 ++++
 .../core/tests/core/test_observability.py     |  95 ++++++++
 11 files changed, 963 insertions(+), 39 deletions(-)

diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md
index 5e9fbff46fe..2c2d996f7ce 100644
--- a/python/packages/core/AGENTS.md
+++ b/python/packages/core/AGENTS.md
@@ -309,7 +309,7 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID
 - **`todos_remaining(*, looping_modes=None)`** / **`todos_remaining_message`** - Helper factories for todo-driven loops (the Python counterpart of .NET's `TodoCompletionLoopEvaluator`), designed for `create_harness_agent` but usable with any agent that registers a `TodoProvider` via `context_providers`. They resolve the `TodoProvider`/`AgentModeProvider` from the *running agent* (`agent.context_providers`, via `_resolve_context_provider`) rather than taking the provider as an argument, so they can be wired directly into `loop_should_continue`/`loop_next_message`. `todos_remaining` returns a `should_continue` predicate that loops while any todo is open; pass `looping_modes=[...]` to gate looping to specific operating modes (case-insensitive; honors the `AgentModeProvider`'s `source_id`/`available_modes`), `looping_modes=None` (default) applies in every mode, and an empty sequence raises `ValueError`. `todos_remaining_message` is a `next_message` callable that lists the still-open todo titles and tells the agent to finish them, returning `None` when the session/agent/provider is unavailable or nothing is open (in which case the middleware's default `None` handling applies: reuse the previous iteration's messages verbatim under the default `fresh_context=False`, or `DEFAULT_NEXT_MESSAGE` only when `fresh_context=True`).
 - **`background_tasks_running()`** / **`background_tasks_running_message`** - Helper factories for background-agent-driven loops, mirroring the `todos_remaining` pair. They resolve the `BackgroundAgentsProvider` from the *running agent* (`agent.context_providers`, via `_resolve_context_provider`) rather than taking the provider as an argument, so they can be wired directly into `create_harness_agent`'s `loop_should_continue`/`loop_next_message`. `background_tasks_running` returns a `should_continue` predicate that loops while the provider's persisted state shows any task with `status == RUNNING` (pair it with `max_iterations` so the loop is bounded even if a task's persisted status is never refreshed). `background_tasks_running_message` is a `next_message` callable that lists the still-running tasks (`#<id> (<agent_name>): <description>`) and tells the agent to wait for them to finish and retrieve their results, returning `None` when the session/agent/provider is unavailable or no task is running.
   - **Approval escape hatch** - `_has_pending_approval_request(result)` checks whether an iteration's response carries a pending tool-approval request (any content with `type == "function_approval_request"`). Both the streaming and non-streaming loops stop and return that response to the caller *before* evaluating `should_continue`/`max_iterations` or injecting `next_message`, so the loop is HITL-safe even when wrapped outermost around a `ToolApprovalMiddleware` (mirrors the C# `LoopAgent`'s `HasPendingApprovalRequests`).
-  - **Budget** - `budget=AgentBudget` (a `total=False` `TypedDict`; plain dicts work) adds `max_tokens` / `max_duration` limits, validated in `AgentLoopMiddleware.__init__` and checked after each iteration before `should_continue`. `max_tokens` charges each iteration the wrapped agent response's `total_token_count`, falling back to `input_token_count + output_token_count`; the `with_judge` judge call is never charged. Sub-agent usage counts only when it appears in the agent's response `usage_details` (`as_tool` and background agents return text only, so their usage is currently not included). The accounting is not configurable.
+  - **Budget** - `budget=AgentBudget` (a `total=False` `TypedDict`; plain dicts work) adds `max_tokens` / `max_duration` limits, validated in `AgentLoopMiddleware.__init__` and checked after each iteration before `should_continue`. `max_tokens` charges each iteration the wrapped agent response's `total_token_count`, falling back to `input_token_count + output_token_count`; the `with_judge` judge call is never charged. Sub-agent usage counts: `as_tool` reports the sub-agent's `usage_details` through a private `FunctionInvocationContext.metadata` key (`_TOOL_USAGE_CONTEXT_KEY`), which the function invocation loop moves onto the function_result and folds into the parent response exactly once (non-streaming `aggregated_usage`; streaming via a `usage` content on the tool update), so nested A->B->C counts each layer once. `BackgroundAgentsProvider` stores each finished run's usage on `BackgroundTaskInfo` (`usage_details`/`usage_reported`) and charges it once, to the parent run that first observes the completion (via the observing tool's context, or a run-scoped chat middleware when `before_run` observes it). `invoke_agent` telemetry spans still report each agent's own usage. Gaps: failed/cancelled sub-agent runs produce no response (tokens lost), background tasks never observed as finished or abandoned at `release_session` are not counted, and custom tools that call `agent.run` directly are not counted. The accounting is not configurable.
   - **Harness integration** - `create_harness_agent` enables the loop when a `loop_should_continue` callable is passed; it prepends `AgentLoopMiddleware(loop_should_continue, max_iterations=loop_max_iterations, next_message=loop_next_message)` ahead of `ToolApprovalMiddleware` so the loop is the outermost middleware (each iteration is a full agent run including tool approval, and the escape hatch hands pending approvals back to the caller). `loop_next_message` and `loop_max_iterations` only take effect together with `loop_should_continue` (with no `loop_should_continue` there is no loop, so they are ignored); `loop_max_iterations` defaults to the loop's default cap (`None` → unbounded).
 
 ### Workflows (`_workflows/`)
diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py
index 5215c14880a..465185bb6e9 100644
--- a/python/packages/core/agent_framework/_agents.py
+++ b/python/packages/core/agent_framework/_agents.py
@@ -687,6 +687,11 @@ def as_tool(
             A FunctionTool that can be used as a tool by other agents.
 
         Note:
+            The sub-agent's response ``usage_details`` (including its own sub-agents) is rolled into
+            the calling agent's response ``usage_details`` when the tool runs inside the function
+            invocation loop, also when the sub-agent stops on an approval or user-input request. A
+            sub-agent run that fails produces no response, so its usage is not counted.
+
             Child function approvals are not propagated into the calling agent.
             Configure ToolApprovalMiddleware with runtime auto-approval rules on
             the child for immediate policy decisions. Use a workflow when approval
@@ -747,6 +752,8 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str:
                 ctx: the function invocation context used
                 **kwargs: only used to dynamically load the argument that is defined for this tool.
             """
+            from ._tools import _record_tool_usage  # pyright: ignore[reportPrivateUsage]
+
             parent_session = ctx.session
             session = AgentSession()
             child_approval_source_ids = _tool_approval_source_ids(self.middleware)
@@ -833,6 +840,10 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str:
                         if isawaitable(callback_result):
                             await callback_result
                 final_response = await stream.get_final_response()
+                # Report the sub-agent's token usage (including its own sub-agents) to the calling loop,
+                # which folds it into the parent response's usage_details exactly once. Recorded before
+                # the approval/user-input raises below so that usage is not lost on those paths.
+                _record_tool_usage(ctx, final_response.usage_details)
                 approval_requests = [
                     request
                     for request in final_response.user_input_requests
diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py
index 6cd70a34037..239f20b004b 100644
--- a/python/packages/core/agent_framework/_harness/_background_agents.py
+++ b/python/packages/core/agent_framework/_harness/_background_agents.py
@@ -11,18 +11,19 @@
 
 import asyncio
 import logging
-from collections.abc import Awaitable, MutableMapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, ClassVar, cast
 
 from .._agents import SupportsAgentRun
 from .._feature_stage import ExperimentalFeature, experimental
+from .._middleware import ChatContext, ChatMiddleware, FunctionInvocationContext
 from .._serialization import SerializationMixin
 from .._sessions import AgentSession, ContextProvider, SessionContext
 from .._telemetry import FeatureIndex, mark_feature_used
-from .._tools import tool
-from .._types import AgentResponse, Message
+from .._tools import _record_tool_usage, tool  # pyright: ignore[reportPrivateUsage]
+from .._types import AgentResponse, ChatResponse, ChatResponseUpdate, Content, Message, UsageDetails, add_usage_details
 
 logger = logging.getLogger(__name__)
 
@@ -65,7 +66,18 @@ class BackgroundTaskInfo(SerializationMixin):
     status: BackgroundTaskStatus
     result_text: str | None
     error_text: str | None
-    __slots__ = ("agent_name", "description", "error_text", "id", "result_text", "status")
+    usage_details: UsageDetails | None
+    usage_reported: bool
+    __slots__ = (
+        "agent_name",
+        "description",
+        "error_text",
+        "id",
+        "result_text",
+        "status",
+        "usage_details",
+        "usage_reported",
+    )
 
     def __init__(
         self,
@@ -75,14 +87,30 @@ def __init__(
         status: BackgroundTaskStatus = BackgroundTaskStatus.RUNNING,
         result_text: str | None = None,
         error_text: str | None = None,
+        usage_details: UsageDetails | None = None,
+        usage_reported: bool = False,
     ) -> None:
-        """Initialize a background task info entry."""
+        """Initialize a background task info entry.
+
+        Args:
+            id: The task ID.
+            agent_name: The name of the background agent running the task.
+            description: The task description.
+            status: The task status.
+            result_text: The text output of a completed task.
+            error_text: The error of a failed task.
+            usage_details: Token usage of the task's latest completed run.
+            usage_reported: Whether ``usage_details`` has already been charged to a parent run, so the
+                usage of each completed run is charged exactly once.
+        """
         self.id = id
         self.agent_name = agent_name
         self.description = description
         self.status = status
         self.result_text = result_text
         self.error_text = error_text
+        self.usage_details = usage_details
+        self.usage_reported = usage_reported
 
     def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
         """Serialize for session state persistence."""
@@ -97,6 +125,9 @@ def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True)
             data["result_text"] = self.result_text
         if not exclude_none or self.error_text is not None:
             data["error_text"] = self.error_text
+        if not exclude_none or self.usage_details is not None:
+            data["usage_details"] = dict(self.usage_details) if self.usage_details is not None else None
+            data["usage_reported"] = self.usage_reported
         return data
 
     @classmethod
@@ -109,9 +140,18 @@ def from_dict(cls, data: MutableMapping[str, Any], **kwargs: Any) -> BackgroundT
             status=BackgroundTaskStatus(data["status"]),
             result_text=data.get("result_text"),
             error_text=data.get("error_text"),
+            usage_details=_usage_from_state(data.get("usage_details")),
+            usage_reported=bool(data.get("usage_reported", False)),
         )
 
 
+def _usage_from_state(value: Any) -> UsageDetails | None:
+    """Restore persisted usage details, ignoring malformed state."""
+    if not isinstance(value, Mapping) or not value:
+        return None
+    return cast(UsageDetails, dict(cast(Mapping[str, Any], value)))
+
+
 @dataclass
 class _RuntimeState:
     """Non-serializable per-session runtime state for background tasks."""
@@ -227,11 +267,67 @@ def _finalize_task(
             task_info.status = BackgroundTaskStatus.FAILED
             task_info.error_text = str(exception)
         else:
+            response = completed_task.result()
             task_info.status = BackgroundTaskStatus.COMPLETED
-            task_info.result_text = completed_task.result().text
+            task_info.result_text = response.text
+            # Held until a parent run observes the completion; see _take_unreported_usage.
+            task_info.usage_details = (
+                cast(UsageDetails, dict(response.usage_details)) if response.usage_details else None
+            )
+            task_info.usage_reported = False
     runtime.in_flight_tasks.pop(task_info.id, None)
 
 
+def _take_unreported_usage(tasks: Sequence[BackgroundTaskInfo]) -> UsageDetails | None:
+    """Sum the usage of finished tasks not yet charged to a parent run and mark it charged.
+
+    Each completed run's usage is charged exactly once, to the parent run that first observes the
+    completion. Failed or cancelled runs produce no response, so they carry no usage.
+    """
+    total: UsageDetails | None = None
+    for task_info in tasks:
+        if task_info.usage_details and not task_info.usage_reported:
+            total = add_usage_details(total, task_info.usage_details)
+            task_info.usage_reported = True
+    return total
+
+
+class _BackgroundUsageChargeMiddleware(ChatMiddleware):
+    """Charge background task usage that ``before_run`` observed to the current run's response.
+
+    ``before_run`` has no tool invocation context to report usage through, so this run-scoped chat
+    middleware folds the usage into the first successful model response of the run instead. The
+    function invocation loop then folds that response's usage into the run's ``usage_details``.
+    """
+
+    def __init__(self, take_usage: Callable[[], UsageDetails | None]) -> None:
+        self._take_usage = take_usage
+
+    def _charge(self) -> UsageDetails | None:
+        usage = self._take_usage()
+        if usage:
+            from ..observability import _record_delegated_usage  # pyright: ignore[reportPrivateUsage]
+
+            # The usage is not this agent's own; keep it out of the invoke_agent span's fallback usage.
+            _record_delegated_usage(usage)
+        return usage
+
+    async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+        if not context.stream:
+            await call_next()
+            if isinstance(context.result, ChatResponse) and (usage := self._charge()):
+                context.result.usage_details = add_usage_details(context.result.usage_details, usage)
+            return
+
+        def _attach_usage(update: ChatResponseUpdate) -> ChatResponseUpdate:
+            if usage := self._charge():
+                update.contents = [*update.contents, Content.from_usage(usage)]
+            return update
+
+        context.stream_update_transforms.append(_attach_usage)
+        await call_next()
+
+
 def _refresh_task_state(
     session: AgentSession, state: dict[str, Any], runtime: _RuntimeState, *, source_id: str
 ) -> list[BackgroundTaskInfo]:
@@ -283,6 +379,14 @@ class BackgroundAgentsProvider(ContextProvider):
     - ``background_agents_continue_task`` — Send follow-up input to a completed task's session to resume work.
     - ``background_agents_clear_completed_task`` — Remove a completed task and release its session.
 
+    Token usage:
+        Each completed background run's token usage is charged exactly once to the parent agent's
+        response ``usage_details``: to the parent run that first observes the task finished, through
+        the observing tool call, or through the run's first model response when ``before_run`` observes
+        it. Failed or cancelled runs produce no response, so their tokens are not counted, and tasks
+        that are never observed as finished (including tasks abandoned by :meth:`release_session`)
+        are not counted either.
+
     Security considerations:
         The agents passed to the constructor are delegated arbitrary work by the parent agent — the
         parent sends them text input (which may include content derived from the parent's own
@@ -473,6 +577,18 @@ async def before_run(
         runtime = self._get_runtime(session)
         source_id = self.source_id
 
+        def take_unreported_usage(tasks: Sequence[BackgroundTaskInfo]) -> UsageDetails | None:
+            usage = _take_unreported_usage(tasks)
+            if usage is not None:
+                _save_tasks(provider_state, list(tasks))
+                _save_provider_state(session, provider_state, source_id=source_id)
+            return usage
+
+        def charge_observed_usage(tasks: Sequence[BackgroundTaskInfo], ctx: FunctionInvocationContext | None) -> None:
+            """Charge finished tasks' usage to the parent run through the observing tool's context."""
+            if ctx is not None:
+                _record_tool_usage(ctx, take_unreported_usage(tasks))
+
         @tool(name="background_agents_start_task", approval_mode="never_require")
         def background_agents_start_task(agent_name: str, input: str, description: str) -> str:
             """Start a background task on a named agent. Returns a confirmation with the task ID."""
@@ -514,7 +630,9 @@ def background_agents_start_task(agent_name: str, input: str, description: str)
         background_agents_start_task._invoke_sync_on_event_loop = True  # pyright: ignore[reportPrivateUsage]
 
         @tool(name="background_agents_wait_for_first_completion", approval_mode="never_require")
-        async def background_agents_wait_for_first_completion(task_ids: list[int]) -> str:
+        async def background_agents_wait_for_first_completion(
+            task_ids: list[int], ctx: FunctionInvocationContext | None = None
+        ) -> str:
             """Wait until the first of the specified tasks completes or the configured timeout expires.
 
             Returns the completed task's ID. On timeout, the tasks remain running and this tool
@@ -536,6 +654,7 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st
             if not waitable:
                 # Refresh state to catch any that completed.
                 tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
+                charge_observed_usage(tasks, ctx)
                 already_complete = next(
                     (t for t in tasks if t.id in task_ids and t.status != BackgroundTaskStatus.RUNNING), None
                 )
@@ -573,14 +692,16 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st
                     _finalize_task(task_info, completed_task, runtime)
                     _save_tasks(provider_state, tasks)
                     _save_provider_state(session, provider_state, source_id=source_id)
+                    charge_observed_usage(tasks, ctx)
 
             status_str = task_info.status.value if task_info else "Unknown"
             return f"Task {completed_id} finished with status: {status_str}."
 
         @tool(name="background_agents_get_task_results", approval_mode="never_require")
-        def background_agents_get_task_results(task_id: int) -> str:
+        def background_agents_get_task_results(task_id: int, ctx: FunctionInvocationContext | None = None) -> str:
             """Get the text output of a background task by its ID."""
             tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
+            charge_observed_usage(tasks, ctx)
             task_info = next((t for t in tasks if t.id == task_id), None)
 
             if task_info is None:
@@ -597,9 +718,10 @@ def background_agents_get_task_results(task_id: int) -> str:
             return f"Task {task_id} has status: {task_info.status.value}."
 
         @tool(name="background_agents_get_all_tasks", approval_mode="never_require")
-        def background_agents_get_all_tasks() -> str:
+        def background_agents_get_all_tasks(ctx: FunctionInvocationContext | None = None) -> str:
             """List all background tasks with their IDs, statuses, agent names, and descriptions."""
             tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
+            charge_observed_usage(tasks, ctx)
 
             if not tasks:
                 return "No tasks."
@@ -610,12 +732,16 @@ def background_agents_get_all_tasks() -> str:
             return "\n".join(lines)
 
         @tool(name="background_agents_continue_task", approval_mode="never_require")
-        def background_agents_continue_task(task_id: int, text: str) -> str:
+        def background_agents_continue_task(
+            task_id: int, text: str, ctx: FunctionInvocationContext | None = None
+        ) -> str:
             """Send follow-up input to a completed or failed task to resume its work."""
             if runtime.closed:
                 return "Error: Session is being released; cannot continue a background task."
 
             tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
+            # Charge the previous run's usage before a new run replaces it.
+            charge_observed_usage(tasks, ctx)
             task_info = next((t for t in tasks if t.id == task_id), None)
 
             if task_info is None:
@@ -647,6 +773,8 @@ def background_agents_continue_task(task_id: int, text: str) -> str:
             task_info.status = BackgroundTaskStatus.RUNNING
             task_info.result_text = None
             task_info.error_text = None
+            task_info.usage_details = None
+            task_info.usage_reported = False
             _save_tasks(provider_state, tasks)
 
             _save_provider_state(session, provider_state, source_id=source_id)
@@ -655,12 +783,13 @@ def background_agents_continue_task(task_id: int, text: str) -> str:
         background_agents_continue_task._invoke_sync_on_event_loop = True  # pyright: ignore[reportPrivateUsage]
 
         @tool(name="background_agents_clear_completed_task", approval_mode="never_require")
-        def background_agents_clear_completed_task(task_id: int) -> str:
+        def background_agents_clear_completed_task(task_id: int, ctx: FunctionInvocationContext | None = None) -> str:
             """Remove a completed or failed task and release its session to free memory."""
             if runtime.closed:
                 return "Error: Session is being released; cannot clear tasks."
 
             tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
+            charge_observed_usage(tasks, ctx)
             task_info = next((t for t in tasks if t.id == task_id), None)
 
             if task_info is None:
@@ -697,6 +826,15 @@ def background_agents_clear_completed_task(task_id: int) -> str:
         # Include current task status as context message if there are tasks.
         # Refresh first to get accurate statuses for any tasks that completed between turns.
         tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
+        if any(t.usage_details and not t.usage_reported for t in tasks):
+            # Completions observed here have no tool context to report through, so charge their usage to
+            # this run's first model response instead.
+            context.extend_middleware(
+                self.source_id,
+                _BackgroundUsageChargeMiddleware(
+                    lambda: take_unreported_usage(_get_tasks(provider_state)),
+                ),
+            )
         if tasks:
             status_lines = ["### Current background tasks"]
             for t in tasks:
diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py
index f75e2a25fc3..df4e01b35ac 100644
--- a/python/packages/core/agent_framework/_harness/_loop.py
+++ b/python/packages/core/agent_framework/_harness/_loop.py
@@ -177,9 +177,13 @@ class AgentBudget(TypedDict, total=False):
     Calls the loop framework makes on its own behalf are never charged: the judge call made by
     :meth:`AgentLoopMiddleware.with_judge` does not count toward this budget.
 
-    Sub-agents are only counted when their usage is reported in the wrapped agent's response
-    usage details. Agents exposed as tools via :meth:`~agent_framework.Agent.as_tool` and
-    background agents return only text, so their token usage is currently *not* included.
+    Sub-agent usage counts too: an agent exposed as a tool via :meth:`~agent_framework.Agent.as_tool`
+    rolls its response usage (including its own sub-agents) into the calling agent's usage details,
+    and a :class:`~agent_framework.BackgroundAgentsProvider` task's usage is charged once, to the
+    parent run that first observes the task finished. Remaining gaps: a sub-agent run that fails or
+    is cancelled produces no response, so its tokens are not counted; background tasks that are never
+    observed as finished (including tasks abandoned when their session is released) are not counted;
+    and custom tools that call ``agent.run`` directly do not report the nested run's usage.
 
     This accounting is fixed and not configurable.
     """
@@ -410,9 +414,9 @@ def __init__(
             budget: Optional :class:`AgentBudget` (or a plain ``dict`` with the same keys) bounding total
                 token spend and/or wall-clock time. ``max_tokens`` counts the wrapped agent's
                 ``total_token_count`` (falling back to input + output tokens) per iteration; the
-                judge call made by :meth:`with_judge` is never counted, and sub-agent usage counts
-                only when it is reported in the agent's response usage details (see
-                :attr:`AgentBudget.max_tokens`). Not configurable. Checked after each iteration, before
+                judge call made by :meth:`with_judge` is never counted, while sub-agent usage from
+                ``as_tool`` and background agents is included, with the gaps listed on
+                :attr:`AgentBudget.max_tokens`. Not configurable. Checked after each iteration, before
                 ``should_continue`` is evaluated. When exceeded the loop stops and (non-streaming
                 only) stamps ``loop_exit_reason``. Budgets do not interrupt an in-flight model/tool
                 call; use ``asyncio.timeout`` (or similar) around the run for a hard wall-clock
@@ -511,7 +515,8 @@ def with_judge(
                 baseline between iterations. See :meth:`__init__` for the full semantics. Defaults to
                 ``False``.
             budget: Optional :class:`AgentBudget`. See :meth:`__init__`. The judge call's own token
-                usage is never charged against ``max_tokens``.
+                usage is never charged against ``max_tokens``; the agent's sub-agent usage (``as_tool``
+                and background agents) is, as described on :attr:`AgentBudget.max_tokens`.
         """
         judge_instructions = (instructions or DEFAULT_JUDGE_INSTRUCTIONS).replace(
             CRITERIA_PLACEHOLDER, _render_criteria_block(criteria)
diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py
index b008a62733b..b4808e6bc0e 100644
--- a/python/packages/core/agent_framework/_tools.py
+++ b/python/packages/core/agent_framework/_tools.py
@@ -127,6 +127,13 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None)
 _APPROVAL_REQUEST_ID_KEY: Final[str] = "_approval_request_id"
 _FUNCTION_INVOCATION_BUDGET_STATE_KEY: Final[str] = "_function_invocation_budget_state"
 _FUNCTION_RESULT_CARRIER_CONTEXT_KEY: Final[str] = "_function_result_carrier"
+# Private FunctionInvocationContext.metadata key through which a tool reports token usage it spent on the
+# caller's behalf (for example a delegated sub-agent run). The function invocation loop folds it into the
+# calling response's usage_details exactly once.
+_TOOL_USAGE_CONTEXT_KEY: Final[str] = "_tool_usage_details"
+# Private exception attribute that carries tool usage across control-flow exceptions which escape
+# _auto_invoke_function without producing a function_result (user input requests, middleware termination).
+_TOOL_USAGE_EXCEPTION_ATTRIBUTE: Final[str] = "_agent_framework_tool_usage_details"
 _FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY: Final[str] = "_function_result_payload_budget"
 _FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY: Final[str] = "_function_result_payload_budget"
 _APPROVED_ARGUMENTS_CONTEXT_KEY: Final[str] = "_approved_function_arguments"
@@ -1968,6 +1975,43 @@ def _replacement_approval_request(
     )
 
 
+def _record_tool_usage(context: FunctionInvocationContext, usage: UsageDetails | None) -> None:  # pyright: ignore[reportUnusedFunction]
+    """Accumulate token usage a tool spent on the caller's behalf into its invocation context.
+
+    The function invocation loop moves the accumulated usage onto the tool's result and folds it
+    into the calling response's ``usage_details`` exactly once.
+    """
+    if not usage:
+        return
+    from ._types import add_usage_details
+
+    existing = context.metadata.get(_TOOL_USAGE_CONTEXT_KEY)
+    context.metadata[_TOOL_USAGE_CONTEXT_KEY] = add_usage_details(
+        cast("UsageDetails | None", existing) if isinstance(existing, Mapping) else None,
+        usage,
+    )
+
+
+def _pop_tool_usage(context: FunctionInvocationContext | None) -> UsageDetails | None:
+    """Remove and return the tool usage accumulated in an invocation context."""
+    if context is None:
+        return None
+    usage = context.metadata.pop(_TOOL_USAGE_CONTEXT_KEY, None)
+    return cast("UsageDetails", usage) if isinstance(usage, Mapping) and usage else None
+
+
+def _attach_tool_usage_to_exception(exception: BaseException, context: FunctionInvocationContext | None) -> None:
+    """Carry tool usage on a control-flow exception that escapes without a function_result."""
+    if (usage := _pop_tool_usage(context)) is not None:
+        setattr(exception, _TOOL_USAGE_EXCEPTION_ATTRIBUTE, usage)
+
+
+def _tool_usage_from_exception(exception: BaseException) -> UsageDetails | None:
+    """Return the tool usage carried by a control-flow exception, if any."""
+    usage = getattr(exception, _TOOL_USAGE_EXCEPTION_ATTRIBUTE, None)
+    return cast("UsageDetails", usage) if isinstance(usage, Mapping) and usage else None
+
+
 def _finalize_function_result(
     *,
     call_id: str,
@@ -1997,6 +2041,9 @@ def _finalize_function_result(
         exception=exception,
         additional_properties=additional_properties,
     )
+    # Usage a tool spent on the caller's behalf (e.g. a delegated sub-agent run) travels on the result so
+    # the function invocation loop can fold it into the calling response exactly once.
+    function_result.usage_details = _pop_tool_usage(context)
     if carrier is None or function_result.items is None:
         return function_result
 
@@ -2149,7 +2196,12 @@ async def _auto_invoke_function(
                 base_additional_properties=function_call_content.additional_properties,
                 context=direct_context,
             )
-        except (MiddlewareFailure, UserInputRequiredException):
+        except UserInputRequiredException as exc:
+            # The tool ran (e.g. a sub-agent that ended on a user-input request), so keep the
+            # usage it spent with the request contents surfaced by the loop.
+            _attach_tool_usage_to_exception(exc, direct_context)
+            raise
+        except MiddlewareFailure:
             # Explicit control-flow signals escape the loop; only ordinary exceptions
             # are absorbed into tool-error results below.
             raise
@@ -2275,6 +2327,8 @@ def ensure_short_circuit_arguments_are_authorized() -> None:
                     base_additional_properties=function_call_content.additional_properties,
                     context=middleware_context,
                 )
+        # Usage not already moved onto a finalized result travels with the exception.
+        _attach_tool_usage_to_exception(term_exc, middleware_context)
         raise
     except _FunctionArgumentsChangedAfterApproval as exc:
         raise MiddlewareTermination(
@@ -2283,7 +2337,10 @@ def ensure_short_circuit_arguments_are_authorized() -> None:
         ) from exc
     except _FunctionArgumentValidationError as exc:
         return _function_argument_validation_error_result(function_call_content, exc, config, middleware_context)
-    except (MiddlewareFailure, UserInputRequiredException):
+    except UserInputRequiredException as exc:
+        _attach_tool_usage_to_exception(exc, middleware_context)
+        raise
+    except MiddlewareFailure:
         # MiddlewareFailure is the loop's explicit fail-closed escape: middleware that
         # must abort the run (enforcement layers, guardrails) raises it instead of
         # relying on the tool-error conversion below, and it propagates to the caller.
@@ -2334,7 +2391,7 @@ async def _execute_single_function_call(
 ) -> tuple[list[Content], bool]:
     from ._middleware import MiddlewareTermination
     from ._sessions import _suspend_run_persistence_gate  # pyright: ignore[reportPrivateUsage]
-    from ._types import Content
+    from ._types import Content, add_usage_details
 
     source_function_call = _underlying_function_call(function_call)
     source_tool = tool_map.get(source_function_call.name) if source_function_call.name is not None else None
@@ -2362,16 +2419,20 @@ async def _execute_single_function_call(
             )
         return [result], False
     except MiddlewareTermination as exc:
+        tool_usage = _tool_usage_from_exception(exc)
         if isinstance(exc.result, Content):
+            if tool_usage is not None:
+                exc.result.usage_details = add_usage_details(exc.result.usage_details, tool_usage)
             return [exc.result], True
         source_function_call = _underlying_function_call(function_call)
-        return [
-            Content.from_function_result(
-                call_id=source_function_call.call_id,  # type: ignore[arg-type]
-                result=exc.result,
-            )
-        ], True
+        terminated_result = Content.from_function_result(
+            call_id=source_function_call.call_id,  # type: ignore[arg-type]
+            result=exc.result,
+        )
+        terminated_result.usage_details = tool_usage
+        return [terminated_result], True
     except UserInputRequiredException as exc:
+        tool_usage = _tool_usage_from_exception(exc)
         source_function_call = _underlying_function_call(function_call)
         call_id = source_function_call.call_id
         propagated_contents = [item for item in exc.contents if isinstance(item, Content)] if exc.contents else []
@@ -2379,15 +2440,19 @@ async def _execute_single_function_call(
             item.call_id = call_id
             if not item.id:
                 item.id = call_id
+            # Request contents can come from a nested agent run whose own loop already counted the usage
+            # they carry; the usage this tool reports below already includes it, so count it only once.
+            item.usage_details = None
         if propagated_contents:
+            propagated_contents[0].usage_details = tool_usage
             return propagated_contents, False
-        return [
-            Content.from_function_result(
-                call_id=call_id,  # type: ignore[arg-type]
-                result="Tool requires user input but no request details were provided.",
-                exception="UserInputRequiredException",
-            )
-        ], False
+        missing_details_result = Content.from_function_result(
+            call_id=call_id,  # type: ignore[arg-type]
+            result="Tool requires user input but no request details were provided.",
+            exception="UserInputRequiredException",
+        )
+        missing_details_result.usage_details = tool_usage
+        return [missing_details_result], False
 
 
 async def _try_execute_function_call_groups(
@@ -2620,6 +2685,17 @@ async def execute_local_calls(function_calls_to_execute: Sequence[Content]) -> l
     return [result_contents for result_contents, _ in execution_results], should_terminate
 
 
+def _sum_content_usage(contents: Iterable[Content]) -> UsageDetails | None:
+    """Sum the usage_details carried by top-level contents, or return None when none carry usage."""
+    from ._types import add_usage_details
+
+    total: UsageDetails | None = None
+    for content in contents:
+        if content.usage_details:
+            total = add_usage_details(total, content.usage_details)
+    return total
+
+
 @dataclass
 class _FunctionExecutionBatch:
     """Results from one ordered batch of function-call executions."""
@@ -2648,6 +2724,15 @@ def executed_call_count(self) -> int:
             if not any(content.type in {"function_approval_request", "function_call"} for content in result_group)
         )
 
+    @property
+    def usage_details(self) -> UsageDetails | None:
+        """Token usage the executed tools spent on the caller's behalf (e.g. delegated sub-agent runs).
+
+        Each tool reports such usage once on its own result content (see ``_TOOL_USAGE_CONTEXT_KEY``),
+        so summing the fresh results of one batch counts every tool exactly once.
+        """
+        return _sum_content_usage(self.contents)
+
     @property
     def had_errors(self) -> bool:
         """Whether any execution produced an error result."""
@@ -4551,6 +4636,26 @@ class _FunctionProcessingResult:
     function_call_count: int = 0
     response_messages: tuple[Message, ...] = ()
     streaming_updates: tuple[ChatResponseUpdate, ...] = ()
+    usage_details: UsageDetails | None = None
+    """Token usage the executed tools spent on the caller's behalf, to fold into the calling response."""
+
+    def streaming_updates_with_usage(self) -> tuple[ChatResponseUpdate, ...]:
+        """Return the streaming updates with this batch's tool usage attached as a usage content.
+
+        The usage rides on the last update so response finalization (``_process_update``) folds it
+        into ``usage_details`` without creating an extra message. The update is copied because
+        terminal-content updates can share their contents list with ``response_messages``.
+        """
+        if not self.usage_details:
+            return self.streaming_updates
+        from ._types import ChatResponseUpdate, Content
+
+        usage_content = Content.from_usage(self.usage_details)
+        if not self.streaming_updates:
+            return (ChatResponseUpdate(role="tool", contents=[usage_content]),)
+        last_update = copy.copy(self.streaming_updates[-1])
+        last_update.contents = [*last_update.contents, usage_content]
+        return (*self.streaming_updates[:-1], last_update)
 
 
 _FunctionCallExecutor: TypeAlias = Callable[..., Awaitable[_FunctionExecutionBatch]]
@@ -4783,6 +4888,7 @@ async def _resolve_deferred_mixed_pause_calls(
         should_terminate = False
         executed_function_count = 0
         had_errors = True
+        usage_details: UsageDetails | None = None
     else:
         try:
             execution = await execute_function_calls(
@@ -4797,6 +4903,7 @@ async def _resolve_deferred_mixed_pause_calls(
         should_terminate = execution.should_terminate
         executed_function_count = execution.executed_call_count
         had_errors = execution.had_errors
+        usage_details = execution.usage_details
 
     _merge_deferred_mixed_pause_results(
         prepared_messages,
@@ -4834,6 +4941,7 @@ async def _resolve_deferred_mixed_pause_calls(
         function_call_count=executed_function_count,
         response_messages=response_messages,
         streaming_updates=streaming_updates,
+        usage_details=usage_details,
     )
 
 
@@ -5054,6 +5162,7 @@ async def _resolve_approval_responses(
     execution_result_groups: list[list[Content]] = []
     should_terminate = False
     reached_error_limit = False
+    usage_details: UsageDetails | None = None
     if responses_to_execute and not (options and options.get("tool_choice") == "none"):
         try:
             execution = await execute_function_calls(
@@ -5069,6 +5178,7 @@ async def _resolve_approval_responses(
             raise
         execution_result_groups = execution.result_groups
         should_terminate = execution.should_terminate
+        usage_details = execution.usage_details
         errors_in_a_row, reached_error_limit = _update_consecutive_error_count(
             errors_in_a_row,
             had_errors=execution.had_errors,
@@ -5129,6 +5239,7 @@ async def _resolve_approval_responses(
         function_call_count=executed_function_count,
         response_messages=response_messages,
         streaming_updates=streaming_updates,
+        usage_details=usage_details,
     )
 
 
@@ -5180,6 +5291,7 @@ async def _process_model_function_calls(
         for message in response.messages
         for content in message.contents
     )
+    processing_result.usage_details = execution.usage_details
     if execution.should_terminate or has_pending_computer_call:
         processing_result.action = "return"
     if processing_result.action == "return":
@@ -5194,6 +5306,29 @@ async def _process_model_function_calls(
     return processing_result
 
 
+def _fold_tool_usage(aggregated: UsageDetails | None, tool_usage: UsageDetails | None) -> UsageDetails | None:
+    """Fold usage that tools spent on the caller's behalf into the loop's aggregated usage.
+
+    The usage is also recorded as delegated for the active ``invoke_agent`` span, whose fallback
+    usage attributes report only the agent's own usage.
+    """
+    if not tool_usage:
+        return aggregated
+    from ._types import add_usage_details
+
+    _record_delegated_usage(tool_usage)
+    return add_usage_details(aggregated, tool_usage)
+
+
+def _record_delegated_usage(tool_usage: UsageDetails | None) -> None:
+    """Record tool-reported usage as delegated for the active ``invoke_agent`` telemetry span."""
+    if not tool_usage:
+        return
+    from .observability import _record_delegated_usage as _record_span_delegated_usage  # pyright: ignore[reportPrivateUsage]
+
+    _record_span_delegated_usage(tool_usage)
+
+
 OptionsCoT = TypeVar(
     "OptionsCoT",
     bound=TypedDict,  # type: ignore[valid-type]
@@ -5414,6 +5549,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non
         )
         function_call_messages.extend(approval_processing.response_messages)
         errors_in_a_row = approval_processing.errors_in_a_row
+        aggregated_usage = _fold_tool_usage(aggregated_usage, approval_processing.usage_details)
         total_function_calls = _record_function_calls(
             budget_state,
             total_function_calls,
@@ -5505,6 +5641,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non
                     error_exception="KeyError" if is_unknown_call else "MiddlewareFailure",
                 )
                 raise
+            aggregated_usage = _fold_tool_usage(aggregated_usage, function_processing.usage_details)
             total_function_calls = _record_function_calls(
                 budget_state,
                 total_function_calls,
@@ -5640,7 +5777,9 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non
             total_function_calls,
             approval_processing.function_call_count,
         )
-        for update in approval_processing.streaming_updates:
+        # The usage rides on the yielded updates; record it as delegated for the agent span.
+        _record_delegated_usage(approval_processing.usage_details)
+        for update in approval_processing.streaming_updates_with_usage():
             yield update
         if approval_processing.action == "return":
             _clear_budget_state_from_session(invocation_session, budget_state)
@@ -5793,7 +5932,8 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non
                 total_function_calls,
                 function_processing.function_call_count,
             )
-            for update in function_processing.streaming_updates:
+            _record_delegated_usage(function_processing.usage_details)
+            for update in function_processing.streaming_updates_with_usage():
                 yield update
             if function_processing.action != "continue" and function_processing.action != "stop":
                 # "return" action: model produced a terminal response.
diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py
index 0290b916eec..5f289df98e7 100644
--- a/python/packages/core/agent_framework/observability.py
+++ b/python/packages/core/agent_framework/observability.py
@@ -130,6 +130,14 @@
     "inner_accumulated_usage", default=None
 )
 
+# Tracks token usage that tools reported on the agent's behalf (for example delegated sub-agent runs) within an
+# agent invoke. That usage is part of the agent response's usage_details, but the invoke_agent span reports only
+# the agent's own usage, so the response-usage fallback subtracts it. The dict is mutated in place so tool usage
+# recorded from a copied context still reaches the owning span.
+_INNER_DELEGATED_USAGE: Final[contextvars.ContextVar[dict[str, int] | None]] = contextvars.ContextVar(
+    "inner_delegated_usage", default=None
+)
+
 # Allows protocol adapters to supply an application-managed conversation identity for one execution
 # without putting that value into a service-owned continuation field.
 _TELEMETRY_CONVERSATION_ID: Final[contextvars.ContextVar[str | None]] = contextvars.ContextVar(
@@ -2394,6 +2402,8 @@ def _trace_agent_invocation(
             inner_response_telemetry_captured_fields: set[str] = set()
             inner_response_telemetry_captured_fields_token: contextvars.Token[set[str] | None] | None = None
             inner_accumulated_usage_token: contextvars.Token[UsageDetails | None] | None = None
+            inner_delegated_usage: dict[str, int] = {}
+            inner_delegated_usage_token: contextvars.Token[dict[str, int] | None] | None = None
             # Agent Framework's agents run in-process (the actual network call happens on a nested
             # chat span), so invoke_agent spans use the default INTERNAL kind.
             span = _start_streaming_span(attributes, OtelAttr.AGENT_NAME)
@@ -2463,6 +2473,7 @@ async def _finalize_stream() -> None:
                         capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD
                         not in inner_response_telemetry_captured_fields,
                         capture_usage=INNER_USAGE_CAPTURED_FIELD not in inner_response_telemetry_captured_fields,
+                        delegated_usage=inner_delegated_usage,
                     )
                     _apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields)
                     _capture_response(span=span, attributes=response_attributes, duration=duration)
@@ -2487,6 +2498,8 @@ async def _finalize_stream() -> None:
                         INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token)
                     if inner_accumulated_usage_token is not None:
                         INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
+                    if inner_delegated_usage_token is not None:
+                        _INNER_DELEGATED_USAGE.reset(inner_delegated_usage_token)
                     _close_span()
 
             def _inner_telemetry_pull_context() -> contextlib.AbstractContextManager[Any]:
@@ -2497,11 +2510,13 @@ def _inner_telemetry_pull_context() -> contextlib.AbstractContextManager[Any]:
                 # underlying iterator is pulled, so inner chat completion spans created during the
                 # pull can still accumulate usage / mark captured fields.
                 nonlocal inner_response_telemetry_captured_fields_token, inner_accumulated_usage_token
+                nonlocal inner_delegated_usage_token
                 if inner_response_telemetry_captured_fields_token is None:
                     inner_response_telemetry_captured_fields_token = INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.set(
                         inner_response_telemetry_captured_fields
                     )
                     inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({})
+                    inner_delegated_usage_token = _INNER_DELEGATED_USAGE.set(inner_delegated_usage)
                 return _activate_span(span)
 
             # The pull context manager attaches the span around each underlying iterator pull so
@@ -2533,6 +2548,8 @@ async def _run() -> AgentResponse[Any]:
                 inner_response_telemetry_captured_fields
             )
             inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({})
+            inner_delegated_usage: dict[str, int] = {}
+            inner_delegated_usage_token = _INNER_DELEGATED_USAGE.set(inner_delegated_usage)
             try:
                 with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span:
                     try:
@@ -2554,6 +2571,7 @@ async def _run() -> AgentResponse[Any]:
                                 capture_usage=(
                                     INNER_USAGE_CAPTURED_FIELD not in inner_response_telemetry_captured_fields
                                 ),
+                                delegated_usage=inner_delegated_usage,
                             )
                             _apply_accumulated_usage(
                                 response_attributes,
@@ -2581,6 +2599,7 @@ async def _run() -> AgentResponse[Any]:
             finally:
                 INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token)
                 INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
+                _INNER_DELEGATED_USAGE.reset(inner_delegated_usage_token)
 
         return _run()
 
@@ -3497,6 +3516,32 @@ def _mark_inner_response_telemetry_captured(
             INNER_ACCUMULATED_USAGE.set(add_usage_details(accumulated, response.usage_details))
 
 
+def _record_delegated_usage(usage: Mapping[str, Any] | None) -> None:  # pyright: ignore[reportUnusedFunction]
+    """Record usage that tools reported on the active agent's behalf (e.g. a delegated sub-agent run).
+
+    The function invocation loop folds such usage into the agent response's ``usage_details``; the
+    ``invoke_agent`` span reports only the agent's own usage, so its response-usage fallback subtracts it.
+    """
+    delegated = _INNER_DELEGATED_USAGE.get()
+    if delegated is None or not usage:
+        return
+    for key, value in usage.items():
+        if isinstance(value, int) and not isinstance(value, bool):
+            delegated[key] = delegated.get(key, 0) + value
+
+
+def _subtract_usage(usage: Mapping[str, Any], delegated: Mapping[str, int]) -> dict[str, Any]:
+    """Return ``usage`` minus ``delegated`` token counts, dropping counts that are not positive."""
+    own: dict[str, Any] = {}
+    for key, value in usage.items():
+        if isinstance(value, int) and not isinstance(value, bool) and key in delegated:
+            value -= delegated[key]
+            if value <= 0:
+                continue
+        own[key] = value
+    return own
+
+
 def _apply_accumulated_usage(attributes: dict[str, Any], captured_fields: set[str]) -> None:
     """Apply accumulated usage from inner chat spans to the invoke_agent span attributes."""
     if INNER_USAGE_CAPTURED_FIELD not in captured_fields:
@@ -3541,8 +3586,13 @@ def _get_response_attributes(
     *,
     capture_response_id: bool = True,
     capture_usage: bool = True,
+    delegated_usage: Mapping[str, int] | None = None,
 ) -> dict[str, Any]:
-    """Get the response attributes from a response."""
+    """Get the response attributes from a response.
+
+    ``delegated_usage`` is token usage that tools reported on the agent's behalf (e.g. sub-agent runs). It is
+    part of ``response.usage_details`` but excluded here so the span reports only the agent's own usage.
+    """
     if capture_response_id and response.response_id:
         attributes[OtelAttr.RESPONSE_ID] = response.response_id
     finish_reason = _get_response_finish_reason(response)
@@ -3551,7 +3601,7 @@ def _get_response_attributes(
     if model := getattr(response, "model", None):
         attributes[OtelAttr.RESPONSE_MODEL] = model
     if capture_usage and (usage := response.usage_details):
-        _apply_usage_attributes(attributes, usage)
+        _apply_usage_attributes(attributes, _subtract_usage(usage, delegated_usage) if delegated_usage else usage)
     return attributes
 
 
diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py
index 5bd4b4d8754..07f16c9c77e 100644
--- a/python/packages/core/tests/core/test_agents.py
+++ b/python/packages/core/tests/core/test_agents.py
@@ -4802,3 +4802,209 @@ async def test_psc_flag_on_storing_without_conversation_id_warns_every_call(
     # Two service calls -> the warning is emitted twice (one per call, not deduped).
     missing_id_warnings = [r for r in caplog.records if "returned no conversation_id" in r.message]
     assert len(missing_id_warnings) == 2
+
+
+# region as_tool sub-agent usage roll-up
+
+
+def _usage(input_tokens: int, output_tokens: int) -> dict[str, int]:
+    return {
+        "input_token_count": input_tokens,
+        "output_token_count": output_tokens,
+        "total_token_count": input_tokens + output_tokens,
+    }
+
+
+def _script_turns(client: MockBaseChatClient, *turns: tuple[list[Content], dict[str, int]]) -> None:
+    """Script the same model turns for streaming and non-streaming calls, each turn with its usage."""
+    client.run_responses = [
+        ChatResponse(messages=Message(role="assistant", contents=list(contents)), usage_details=usage)  # type: ignore[arg-type]
+        for contents, usage in turns
+    ]
+    client.streaming_responses = [
+        [ChatResponseUpdate(role="assistant", contents=[*contents, Content.from_usage(usage)])]  # type: ignore[arg-type]
+        for contents, usage in turns
+    ]
+
+
+def _delegate_call(call_id: str, name: str = "delegate") -> Content:
+    return Content.from_function_call(call_id=call_id, name=name, arguments={"task": "work"})
+
+
+async def _run_agent(agent: Agent, stream: bool) -> AgentResponse:
+    if stream:
+        return await agent.run("Go.", stream=True).get_final_response()
+    return await agent.run("Go.")
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_as_tool_sub_agent_usage_rolls_into_parent_usage(stream: bool) -> None:
+    child_client = MockBaseChatClient()
+    _script_turns(child_client, ([Content.from_text("child answer")], _usage(10, 5)))
+    child = Agent(client=child_client, name="Child")
+
+    parent_client = MockBaseChatClient()
+    _script_turns(
+        parent_client,
+        ([_delegate_call("call-1")], _usage(3, 2)),
+        ([Content.from_text("parent answer")], _usage(4, 1)),
+    )
+    parent = Agent(client=parent_client, name="Parent", tools=[child.as_tool(name="delegate")])
+
+    response = await _run_agent(parent, stream)
+
+    assert response.text.endswith("parent answer")
+    assert response.usage_details == _usage(17, 8)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_as_tool_nested_sub_agent_usage_counts_each_layer_once(stream: bool) -> None:
+    c_client = MockBaseChatClient()
+    _script_turns(c_client, ([Content.from_text("c answer")], _usage(10, 1)))
+    agent_c = Agent(client=c_client, name="AgentC")
+
+    b_client = MockBaseChatClient()
+    _script_turns(
+        b_client,
+        ([_delegate_call("b-call", name="c_tool")], _usage(20, 2)),
+        ([Content.from_text("b answer")], _usage(30, 3)),
+    )
+    agent_b = Agent(client=b_client, name="AgentB", tools=[agent_c.as_tool(name="c_tool")])
+
+    a_client = MockBaseChatClient()
+    _script_turns(
+        a_client,
+        ([_delegate_call("a-call", name="b_tool")], _usage(100, 10)),
+        ([Content.from_text("a answer")], _usage(200, 20)),
+    )
+    agent_a = Agent(client=a_client, name="AgentA", tools=[agent_b.as_tool(name="b_tool")])
+
+    response = await _run_agent(agent_a, stream)
+
+    # A spends 300/30 itself, B 50/5 and C 10/1; each layer is counted exactly once.
+    assert response.usage_details == _usage(360, 36)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_as_tool_parallel_sub_agent_calls_sum_usage(stream: bool) -> None:
+    child_client = MockBaseChatClient()
+    _script_turns(
+        child_client,
+        ([Content.from_text("first child answer")], _usage(10, 1)),
+        ([Content.from_text("second child answer")], _usage(20, 2)),
+    )
+    child = Agent(client=child_client, name="Child")
+
+    parent_client = MockBaseChatClient()
+    _script_turns(
+        parent_client,
+        ([_delegate_call("call-1"), _delegate_call("call-2")], _usage(1, 1)),
+        ([Content.from_text("parent answer")], _usage(2, 2)),
+    )
+    parent = Agent(client=parent_client, name="Parent", tools=[child.as_tool(name="delegate")])
+
+    response = await _run_agent(parent, stream)
+
+    assert response.usage_details == _usage(33, 6)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_as_tool_sub_agent_usage_survives_user_input_request(stream: bool) -> None:
+    """A sub-agent that ends on a user-input request raises out of the tool; its usage still rolls up."""
+
+    @tool(name="guarded_write", approval_mode="always_require")
+    def guarded_write() -> str:
+        raise AssertionError("Unapproved child tool must not execute.")
+
+    child_client = MockBaseChatClient()
+    _script_turns(
+        child_client,
+        (
+            [
+                Content.from_function_call(call_id="child-write", name="guarded_write", arguments={}),
+                Content.from_oauth_consent_request(consent_link="https://example.com/consent"),
+            ],
+            _usage(10, 5),
+        ),
+    )
+    child = Agent(client=child_client, name="Child", tools=[guarded_write])
+
+    parent_client = MockBaseChatClient()
+    _script_turns(parent_client, ([_delegate_call("call-1")], _usage(3, 2)))
+    parent = Agent(client=parent_client, name="Parent", tools=[child.as_tool(name="delegate")])
+
+    response = await _run_agent(parent, stream)
+
+    assert [request.type for request in response.user_input_requests] == ["oauth_consent_request"]
+    assert response.usage_details == _usage(13, 7)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_as_tool_sub_agent_usage_survives_unresolved_approval_error(stream: bool) -> None:
+    """A sub-agent stopped on an approval request becomes a tool error; its usage still rolls up."""
+
+    @tool(name="guarded_write", approval_mode="always_require")
+    def guarded_write() -> str:
+        raise AssertionError("Unapproved child tool must not execute.")
+
+    child_client = MockBaseChatClient()
+    _script_turns(
+        child_client,
+        ([Content.from_function_call(call_id="child-write", name="guarded_write", arguments={})], _usage(10, 5)),
+    )
+    child = Agent(client=child_client, name="Child", tools=[guarded_write])
+
+    parent_client = MockBaseChatClient()
+    _script_turns(
+        parent_client,
+        ([_delegate_call("call-1")], _usage(3, 2)),
+        ([Content.from_text("parent answer")], _usage(4, 1)),
+    )
+    parent = Agent(client=parent_client, name="Parent", tools=[child.as_tool(name="delegate")])
+
+    response = await _run_agent(parent, stream)
+
+    function_results = [
+        content for message in response.messages for content in message.contents if content.type == "function_result"
+    ]
+    assert len(function_results) == 1
+    assert function_results[0].exception is not None
+    assert response.usage_details == _usage(17, 8)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_as_tool_sub_agent_usage_kept_when_middleware_overrides_result(stream: bool) -> None:
+    from agent_framework import FunctionMiddleware
+
+    class _OverrideResult(FunctionMiddleware):
+        async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+            context.result = "overridden"
+
+    child_client = MockBaseChatClient()
+    _script_turns(child_client, ([Content.from_text("child answer")], _usage(10, 5)))
+    child = Agent(client=child_client, name="Child")
+
+    parent_client = MockBaseChatClient()
+    _script_turns(
+        parent_client,
+        ([_delegate_call("call-1")], _usage(3, 2)),
+        ([Content.from_text("parent answer")], _usage(4, 1)),
+    )
+    parent = Agent(
+        client=parent_client,
+        name="Parent",
+        tools=[child.as_tool(name="delegate")],
+        middleware=[_OverrideResult()],
+    )
+
+    response = await _run_agent(parent, stream)
+
+    function_results = [
+        content for message in response.messages for content in message.contents if content.type == "function_result"
+    ]
+    assert [result.result for result in function_results] == ["overridden"]
+    assert response.usage_details == _usage(17, 8)
+
+
+# endregion
diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py
index 653617e1e60..086e43f9af8 100644
--- a/python/packages/core/tests/core/test_function_invocation_logic.py
+++ b/python/packages/core/tests/core/test_function_invocation_logic.py
@@ -11824,3 +11824,86 @@ async def test_unbound_local_approval_response_is_filtered_before_mixed_batch_va
     assert executed == []
     # The warning proves the response was filtered before the completeness check reported it missing.
     assert any("tool-approval response" in record.message for record in caplog.records)
+
+
+# region tool-reported usage
+
+
+_TOOL_REPORTED_USAGE = {"input_token_count": 7, "output_token_count": 3, "total_token_count": 10}
+
+
+async def _get_tool_usage_response(client: Any, tool_obj: FunctionTool, *, streaming: bool) -> ChatResponse:
+    messages = [Message(role="user", contents=["run"])]
+    if streaming:
+        stream = client.get_response(messages, options={"tools": [tool_obj]}, stream=True)
+        async for _ in stream:
+            pass
+        return await stream.get_final_response()
+    return await client.get_response(messages, options={"tools": [tool_obj]})
+
+
+def _script_tool_call_turn(client: Any, *, streaming: bool) -> None:
+    call = Content.from_function_call(call_id="usage-call", name="delegating_tool", arguments={})
+    model_usage = {"input_token_count": 2, "output_token_count": 1, "total_token_count": 3}
+    if streaming:
+        client.streaming_responses = [
+            [ChatResponseUpdate(role="assistant", contents=[call, Content.from_usage(model_usage)])]  # type: ignore[arg-type]
+        ]
+    else:
+        client.run_responses = [
+            ChatResponse(messages=Message(role="assistant", contents=[call]), usage_details=model_usage)  # type: ignore[arg-type]
+        ]
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+async def test_tool_reported_usage_counted_once_across_max_iterations_final_call(
+    chat_client_base: SupportsChatGetResponse, streaming: bool
+) -> None:
+    """Usage a tool reports through the private context key is folded in once, also when Phase 3 runs."""
+    from agent_framework._tools import _TOOL_USAGE_CONTEXT_KEY
+
+    @tool(name="delegating_tool", approval_mode="never_require")
+    def delegating_tool(ctx: FunctionInvocationContext) -> str:
+        ctx.metadata[_TOOL_USAGE_CONTEXT_KEY] = dict(_TOOL_REPORTED_USAGE)
+        return "delegated"
+
+    chat_client_base.function_invocation_configuration["max_iterations"] = 1  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    _script_tool_call_turn(chat_client_base, streaming=streaming)
+
+    response = await _get_tool_usage_response(chat_client_base, delegating_tool, streaming=streaming)
+
+    # The max-iterations Phase 3 call ran (and reported no usage of its own).
+    assert chat_client_base.call_count == 2  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert response.usage_details == {"input_token_count": 9, "output_token_count": 4, "total_token_count": 13}
+    function_results = [
+        content for message in response.messages for content in message.contents if content.type == "function_result"
+    ]
+    assert len(function_results) == 1
+    assert function_results[0].usage_details == _TOOL_REPORTED_USAGE
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+async def test_tool_reported_usage_survives_tool_exception(
+    chat_client_base: SupportsChatGetResponse, streaming: bool
+) -> None:
+    """Usage reported before a tool raises is kept on the error result and folded into the response."""
+    from agent_framework._tools import _TOOL_USAGE_CONTEXT_KEY
+
+    @tool(name="delegating_tool", approval_mode="never_require")
+    def delegating_tool(ctx: FunctionInvocationContext) -> str:
+        ctx.metadata[_TOOL_USAGE_CONTEXT_KEY] = dict(_TOOL_REPORTED_USAGE)
+        raise RuntimeError("delegated work failed after spending tokens")
+
+    _script_tool_call_turn(chat_client_base, streaming=streaming)
+
+    response = await _get_tool_usage_response(chat_client_base, delegating_tool, streaming=streaming)
+
+    function_results = [
+        content for message in response.messages for content in message.contents if content.type == "function_result"
+    ]
+    assert len(function_results) == 1
+    assert function_results[0].exception is not None
+    assert response.usage_details == {"input_token_count": 9, "output_token_count": 4, "total_token_count": 13}
+
+
+# endregion
diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py
index 88e776a0c81..fac33cea670 100644
--- a/python/packages/core/tests/core/test_harness_background_agents.py
+++ b/python/packages/core/tests/core/test_harness_background_agents.py
@@ -9,11 +9,15 @@
 import pytest
 
 from agent_framework import (
+    Agent,
     AgentResponse,
     AgentSession,
     BackgroundAgentsProvider,
     BackgroundTaskInfo,
     BackgroundTaskStatus,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
     Message,
     ServiceSessionId,
 )
@@ -920,3 +924,147 @@ async def fake_drain(*args: Any, **kwargs: Any) -> None:
 
     assert drain_called is False
     assert session.session_id not in provider._runtime
+
+
+# --- Usage Accounting Tests ---
+
+
+def _usage(input_tokens: int, output_tokens: int) -> dict[str, int]:
+    return {
+        "input_token_count": input_tokens,
+        "output_token_count": output_tokens,
+        "total_token_count": input_tokens + output_tokens,
+    }
+
+
+class _UsageReportingFakeAgent(_FakeAgent):
+    """Fake background agent whose response reports token usage."""
+
+    def __init__(self, name: str, usage: dict[str, int]) -> None:
+        super().__init__(name)
+        self._usage = usage
+
+    async def run(
+        self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any
+    ) -> AgentResponse[Any]:
+        response = await super().run(messages, stream=stream, session=session, **kwargs)
+        response.usage_details = dict(self._usage)  # type: ignore[assignment]
+        return response
+
+
+def _script_parent_turns(client: Any, *turns: list[Content]) -> None:
+    """Script parent model turns that each spend 1 input and 1 output token."""
+    client.run_responses = [
+        ChatResponse(messages=Message(role="assistant", contents=contents), usage_details=_usage(1, 1))  # type: ignore[arg-type]
+        for contents in turns
+    ]
+    client.streaming_responses = [
+        [ChatResponseUpdate(role="assistant", contents=[*contents, Content.from_usage(_usage(1, 1))])]  # type: ignore[arg-type]
+        for contents in turns
+    ]
+
+
+def _start_task_call(call_id: str) -> Content:
+    return Content.from_function_call(
+        call_id=call_id,
+        name="background_agents_start_task",
+        arguments={"agent_name": "worker", "input": "do it", "description": "work"},
+    )
+
+
+async def _run_parent(agent: Agent, session: AgentSession, *, stream: bool = False) -> AgentResponse[Any]:
+    if stream:
+        return await agent.run("Go.", session=session, stream=True).get_final_response()
+    return await agent.run("Go.", session=session)
+
+
+async def _await_background_tasks(provider: BackgroundAgentsProvider, session: AgentSession) -> None:
+    await asyncio.gather(*provider._runtime[session.session_id].in_flight_tasks.values())
+
+
+def _stored_task(session: AgentSession) -> BackgroundTaskInfo:
+    return BackgroundTaskInfo.from_dict(session.state["background_agents"]["tasks"][0])
+
+
+def test_task_info_serializes_usage() -> None:
+    info = BackgroundTaskInfo(
+        id=1,
+        agent_name="worker",
+        description="work",
+        status=BackgroundTaskStatus.COMPLETED,
+        result_text="done",
+        usage_details=_usage(5, 2),  # type: ignore[arg-type]
+        usage_reported=True,
+    )
+    restored = BackgroundTaskInfo.from_dict(info.to_dict())
+    assert restored.usage_details == _usage(5, 2)
+    assert restored.usage_reported is True
+    assert "usage_details" not in BackgroundTaskInfo(id=2, agent_name="worker", description="work").to_dict()
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_background_task_usage_charged_once_through_observing_tool(stream: bool) -> None:
+    from .conftest import MockBaseChatClient
+
+    provider = _make_provider(_UsageReportingFakeAgent("worker", _usage(50, 5)))
+    parent_client = MockBaseChatClient()
+    parent = Agent(client=parent_client, name="Parent", context_providers=[provider])
+    session = AgentSession()
+
+    _script_parent_turns(
+        parent_client,
+        [_start_task_call("start")],
+        [
+            Content.from_function_call(
+                call_id="wait", name="background_agents_wait_for_first_completion", arguments={"task_ids": [1]}
+            )
+        ],
+        [Content.from_text("waited")],
+    )
+    first = await _run_parent(parent, session, stream=stream)
+
+    # The wait tool observed the completion, so this run is charged the background run's usage.
+    assert first.usage_details == _usage(53, 8)
+    assert _stored_task(session).usage_reported is True
+
+    _script_parent_turns(
+        parent_client,
+        [
+            Content.from_function_call(
+                call_id="results", name="background_agents_get_task_results", arguments={"task_id": 1}
+            )
+        ],
+        [Content.from_text("read results")],
+    )
+    second = await _run_parent(parent, session, stream=stream)
+
+    # Later observations do not charge the same background run again.
+    assert second.usage_details == _usage(2, 2)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_background_task_usage_charged_once_to_run_whose_before_run_observes_completion(
+    stream: bool,
+) -> None:
+    from .conftest import MockBaseChatClient
+
+    provider = _make_provider(_UsageReportingFakeAgent("worker", _usage(50, 5)))
+    parent_client = MockBaseChatClient()
+    parent = Agent(client=parent_client, name="Parent", context_providers=[provider])
+    session = AgentSession()
+
+    _script_parent_turns(parent_client, [_start_task_call("start")], [Content.from_text("started")])
+    first = await _run_parent(parent, session, stream=stream)
+    assert first.usage_details == _usage(2, 2)
+
+    # The task finishes between runs, so the next run's before_run is the first to observe it.
+    await _await_background_tasks(provider, session)
+
+    _script_parent_turns(parent_client, [Content.from_text("noticed completion")])
+    second = await _run_parent(parent, session, stream=stream)
+    assert second.usage_details == _usage(51, 6)
+    assert _stored_task(session).usage_reported is True
+
+    _script_parent_turns(parent_client, [Content.from_text("nothing new")])
+    third = await _run_parent(parent, session, stream=stream)
+    assert third.usage_details == _usage(1, 1)
diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py
index 9322dbf964c..be7891599b0 100644
--- a/python/packages/core/tests/core/test_harness_loop.py
+++ b/python/packages/core/tests/core/test_harness_loop.py
@@ -1861,6 +1861,54 @@ async def test_budget_excludes_judge_call_usage_streaming() -> None:
     assert judge_client.call_count == 2
 
 
+@pytest.mark.parametrize("stream", [False, True])
+async def test_budget_counts_sub_agent_usage_from_as_tool(stream: bool) -> None:
+    # The parent spends 10 tokens per model call; the as_tool sub-agent spends 200. Only the
+    # roll-up of the sub-agent's usage into the parent's usage_details trips the 100-token
+    # budget after iteration 1; without it the loop would run to the iteration cap.
+    child = Agent(client=RecordingChatClient(usage_per_call=200), name="researcher")
+    parent_client = MockBaseChatClient()
+    delegate_call = Content.from_function_call(call_id="delegate", name="researcher", arguments={"task": "dig"})
+    parent_client.run_responses = [
+        ChatResponse(
+            messages=Message(role="assistant", contents=[delegate_call]),
+            usage_details=UsageDetails(total_token_count=10),
+        ),
+        ChatResponse(
+            messages=Message(role="assistant", contents=["summary"]), usage_details=UsageDetails(total_token_count=10)
+        ),
+    ]
+    parent_client.streaming_responses = [
+        [
+            ChatResponseUpdate(
+                role="assistant", contents=[delegate_call, Content.from_usage(UsageDetails(total_token_count=10))]
+            )
+        ],
+        [
+            ChatResponseUpdate(
+                role="assistant",
+                contents=[Content.from_text("summary"), Content.from_usage(UsageDetails(total_token_count=10))],
+            )
+        ],
+    ]
+    agent = Agent(
+        client=parent_client,
+        tools=[child.as_tool()],
+        middleware=[AgentLoopMiddleware(always_continue, max_iterations=10, budget=AgentBudget(max_tokens=100))],
+    )
+
+    if stream:
+        response = await agent.run("task", stream=True).get_final_response()
+    else:
+        response = await agent.run("task")
+
+    assert parent_client.call_count == 2
+    assert response.usage_details is not None
+    assert response.usage_details.get("total_token_count") == 220
+    if not stream:
+        assert response.additional_properties.get("loop_exit_reason") == "token_budget_exceeded"
+
+
 def test_agent_budget_validates_max_tokens() -> None:
     with pytest.raises(ValueError, match="max_tokens"):
         AgentLoopMiddleware(always_continue, budget=AgentBudget(max_tokens=0))
diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py
index fb4a5b1d38c..067228f60fc 100644
--- a/python/packages/core/tests/core/test_observability.py
+++ b/python/packages/core/tests/core/test_observability.py
@@ -6527,6 +6527,101 @@ class _InstrumentedAgent(AgentTelemetryLayer, RawAgent):
     assert agent_span.attributes.get(OtelAttr.OUTPUT_TOKENS) == 100  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
 
 
+def _script_usage_turns(client: Any, *turns: tuple[list[Content], UsageDetails]) -> None:
+    client.run_responses = [
+        ChatResponse(messages=Message(role="assistant", contents=contents), usage_details=usage)
+        for contents, usage in turns
+    ]
+    client.streaming_responses = [
+        [ChatResponseUpdate(role="assistant", contents=[*contents, Content.from_usage(usage)])]
+        for contents, usage in turns
+    ]
+
+
+def _invoke_agent_spans_by_name(spans: Sequence[Any]) -> dict[str, Any]:
+    return {
+        span.attributes.get(OtelAttr.AGENT_NAME): span.attributes
+        for span in spans
+        if span.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.AGENT_INVOKE_OPERATION
+    }
+
+
+async def _run_parent_with_sub_agent(client_type: type[Any], stream: bool) -> dict[str, Any]:
+    child_client = client_type()
+    _script_usage_turns(
+        child_client,
+        ([Content.from_text("child answer")], UsageDetails(input_token_count=1000, output_token_count=500)),
+    )
+    child = Agent(client=child_client, name="child_agent")
+
+    parent_client = client_type()
+    delegate_call = Content.from_function_call(call_id="delegate", name="delegate", arguments={"task": "dig"})
+    _script_usage_turns(
+        parent_client,
+        ([delegate_call], UsageDetails(input_token_count=30, output_token_count=20)),
+        ([Content.from_text("parent answer")], UsageDetails(input_token_count=40, output_token_count=10)),
+    )
+    parent = Agent(client=parent_client, name="parent_agent", tools=[child.as_tool(name="delegate")])
+
+    if stream:
+        response = await parent.run("Go.", stream=True).get_final_response()
+    else:
+        response = await parent.run("Go.")
+    return dict(response.usage_details or {})
+
+
+@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
+@pytest.mark.parametrize("stream", [False, True])
+async def test_agent_invoke_span_reports_own_usage_when_sub_agent_usage_rolls_up(
+    span_exporter: InMemorySpanExporter, stream: bool
+) -> None:
+    """Sub-agent usage rolls into the parent's response, but each invoke_agent span reports its own usage."""
+    from tests.core.conftest import MockBaseChatClient
+
+    span_exporter.clear()
+    usage = await _run_parent_with_sub_agent(MockBaseChatClient, stream)
+
+    assert usage["input_token_count"] == 1070
+    assert usage["output_token_count"] == 530
+    spans = _invoke_agent_spans_by_name(span_exporter.get_finished_spans())
+    assert spans["parent_agent"].get(OtelAttr.INPUT_TOKENS) == 70
+    assert spans["parent_agent"].get(OtelAttr.OUTPUT_TOKENS) == 30
+    assert spans["child_agent"].get(OtelAttr.INPUT_TOKENS) == 1000
+    assert spans["child_agent"].get(OtelAttr.OUTPUT_TOKENS) == 500
+
+
+@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
+@pytest.mark.parametrize("stream", [False, True])
+async def test_agent_invoke_span_usage_fallback_excludes_sub_agent_usage(
+    span_exporter: InMemorySpanExporter, stream: bool
+) -> None:
+    """Without chat spans, the invoke_agent span falls back to response usage minus sub-agent usage."""
+    from agent_framework import ChatMiddlewareLayer, FunctionInvocationLayer
+    from tests.core.conftest import MockBaseChatClient
+
+    class _NoChatTelemetryClient(FunctionInvocationLayer[Any], ChatMiddlewareLayer[Any], BaseChatClient[Any]):
+        def __init__(self) -> None:
+            super().__init__(middleware=[])
+            self.run_responses: list[ChatResponse] = []
+            self.streaming_responses: list[list[ChatResponseUpdate]] = []
+            self.call_count = 0
+
+        _inner_get_response = MockBaseChatClient._inner_get_response
+        _get_non_streaming_response = MockBaseChatClient._get_non_streaming_response
+        _get_streaming_response = MockBaseChatClient._get_streaming_response
+
+    span_exporter.clear()
+    usage = await _run_parent_with_sub_agent(_NoChatTelemetryClient, stream)
+
+    assert usage["input_token_count"] == 1070
+    assert usage["output_token_count"] == 530
+    spans = _invoke_agent_spans_by_name(span_exporter.get_finished_spans())
+    assert spans["parent_agent"].get(OtelAttr.INPUT_TOKENS) == 70
+    assert spans["parent_agent"].get(OtelAttr.OUTPUT_TOKENS) == 30
+    assert spans["child_agent"].get(OtelAttr.INPUT_TOKENS) == 1000
+    assert spans["child_agent"].get(OtelAttr.OUTPUT_TOKENS) == 500
+
+
 # region Test span nesting (parent-child relationships)
 
 

From 98fb54c7a6c056be7389eeea4c22a958d601357f Mon Sep 17 00:00:00 2001
From: Mark Licata <malicata@microsoft.com>
Date: Tue, 13 Oct 2026 10:19:16 -0700
Subject: [PATCH 7/7] Fix test typing for sub-agent usage tests

Build real UsageDetails values instead of plain dicts, type the run
helpers with Agent[Any], and give the no-chat-telemetry observability
client its own _inner_get_response matching BaseChatClient's signature
instead of reassigning MockBaseChatClient's methods.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../packages/core/tests/core/test_agents.py   | 21 +++++++------
 .../core/test_function_invocation_logic.py    |  9 +++---
 .../core/test_harness_background_agents.py    | 25 ++++++++-------
 .../core/tests/core/test_observability.py     | 31 ++++++++++++++++---
 4 files changed, 55 insertions(+), 31 deletions(-)

diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py
index 07f16c9c77e..d096c896bfe 100644
--- a/python/packages/core/tests/core/test_agents.py
+++ b/python/packages/core/tests/core/test_agents.py
@@ -45,6 +45,7 @@
     ToolApprovalMiddleware,
     ToolResultCompactionStrategy,
     TruncationStrategy,
+    UsageDetails,
     VectorStoreHistoryProvider,
     chat_middleware,
     enqueue_messages,
@@ -4807,22 +4808,22 @@ async def test_psc_flag_on_storing_without_conversation_id_warns_every_call(
 # region as_tool sub-agent usage roll-up
 
 
-def _usage(input_tokens: int, output_tokens: int) -> dict[str, int]:
-    return {
-        "input_token_count": input_tokens,
-        "output_token_count": output_tokens,
-        "total_token_count": input_tokens + output_tokens,
-    }
+def _usage(input_tokens: int, output_tokens: int) -> UsageDetails:
+    return UsageDetails(
+        input_token_count=input_tokens,
+        output_token_count=output_tokens,
+        total_token_count=input_tokens + output_tokens,
+    )
 
 
-def _script_turns(client: MockBaseChatClient, *turns: tuple[list[Content], dict[str, int]]) -> None:
+def _script_turns(client: MockBaseChatClient, *turns: tuple[list[Content], UsageDetails]) -> None:
     """Script the same model turns for streaming and non-streaming calls, each turn with its usage."""
     client.run_responses = [
-        ChatResponse(messages=Message(role="assistant", contents=list(contents)), usage_details=usage)  # type: ignore[arg-type]
+        ChatResponse(messages=Message(role="assistant", contents=list(contents)), usage_details=usage)
         for contents, usage in turns
     ]
     client.streaming_responses = [
-        [ChatResponseUpdate(role="assistant", contents=[*contents, Content.from_usage(usage)])]  # type: ignore[arg-type]
+        [ChatResponseUpdate(role="assistant", contents=[*contents, Content.from_usage(usage)])]
         for contents, usage in turns
     ]
 
@@ -4831,7 +4832,7 @@ def _delegate_call(call_id: str, name: str = "delegate") -> Content:
     return Content.from_function_call(call_id=call_id, name=name, arguments={"task": "work"})
 
 
-async def _run_agent(agent: Agent, stream: bool) -> AgentResponse:
+async def _run_agent(agent: Agent[Any], stream: bool) -> AgentResponse:
     if stream:
         return await agent.run("Go.", stream=True).get_final_response()
     return await agent.run("Go.")
diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py
index 5fdedf2b601..5ecb540bde4 100644
--- a/python/packages/core/tests/core/test_function_invocation_logic.py
+++ b/python/packages/core/tests/core/test_function_invocation_logic.py
@@ -32,6 +32,7 @@
     ResponseInvalidatedException,
     ResponseStream,
     SupportsChatGetResponse,
+    UsageDetails,
     chat_middleware,
     tool,
 )
@@ -13033,7 +13034,7 @@ async def test_unbound_local_approval_response_is_filtered_before_mixed_batch_va
 # region tool-reported usage
 
 
-_TOOL_REPORTED_USAGE = {"input_token_count": 7, "output_token_count": 3, "total_token_count": 10}
+_TOOL_REPORTED_USAGE = UsageDetails(input_token_count=7, output_token_count=3, total_token_count=10)
 
 
 async def _get_tool_usage_response(client: Any, tool_obj: FunctionTool, *, streaming: bool) -> ChatResponse:
@@ -13048,14 +13049,14 @@ async def _get_tool_usage_response(client: Any, tool_obj: FunctionTool, *, strea
 
 def _script_tool_call_turn(client: Any, *, streaming: bool) -> None:
     call = Content.from_function_call(call_id="usage-call", name="delegating_tool", arguments={})
-    model_usage = {"input_token_count": 2, "output_token_count": 1, "total_token_count": 3}
+    model_usage = UsageDetails(input_token_count=2, output_token_count=1, total_token_count=3)
     if streaming:
         client.streaming_responses = [
-            [ChatResponseUpdate(role="assistant", contents=[call, Content.from_usage(model_usage)])]  # type: ignore[arg-type]
+            [ChatResponseUpdate(role="assistant", contents=[call, Content.from_usage(model_usage)])]
         ]
     else:
         client.run_responses = [
-            ChatResponse(messages=Message(role="assistant", contents=[call]), usage_details=model_usage)  # type: ignore[arg-type]
+            ChatResponse(messages=Message(role="assistant", contents=[call]), usage_details=model_usage)
         ]
 
 
diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py
index fac33cea670..497f400a300 100644
--- a/python/packages/core/tests/core/test_harness_background_agents.py
+++ b/python/packages/core/tests/core/test_harness_background_agents.py
@@ -20,6 +20,7 @@
     Content,
     Message,
     ServiceSessionId,
+    UsageDetails,
 )
 from agent_framework._sessions import SessionContext
 
@@ -929,18 +930,18 @@ async def fake_drain(*args: Any, **kwargs: Any) -> None:
 # --- Usage Accounting Tests ---
 
 
-def _usage(input_tokens: int, output_tokens: int) -> dict[str, int]:
-    return {
-        "input_token_count": input_tokens,
-        "output_token_count": output_tokens,
-        "total_token_count": input_tokens + output_tokens,
-    }
+def _usage(input_tokens: int, output_tokens: int) -> UsageDetails:
+    return UsageDetails(
+        input_token_count=input_tokens,
+        output_token_count=output_tokens,
+        total_token_count=input_tokens + output_tokens,
+    )
 
 
 class _UsageReportingFakeAgent(_FakeAgent):
     """Fake background agent whose response reports token usage."""
 
-    def __init__(self, name: str, usage: dict[str, int]) -> None:
+    def __init__(self, name: str, usage: UsageDetails) -> None:
         super().__init__(name)
         self._usage = usage
 
@@ -948,18 +949,18 @@ async def run(
         self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any
     ) -> AgentResponse[Any]:
         response = await super().run(messages, stream=stream, session=session, **kwargs)
-        response.usage_details = dict(self._usage)  # type: ignore[assignment]
+        response.usage_details = self._usage.copy()
         return response
 
 
 def _script_parent_turns(client: Any, *turns: list[Content]) -> None:
     """Script parent model turns that each spend 1 input and 1 output token."""
     client.run_responses = [
-        ChatResponse(messages=Message(role="assistant", contents=contents), usage_details=_usage(1, 1))  # type: ignore[arg-type]
+        ChatResponse(messages=Message(role="assistant", contents=contents), usage_details=_usage(1, 1))
         for contents in turns
     ]
     client.streaming_responses = [
-        [ChatResponseUpdate(role="assistant", contents=[*contents, Content.from_usage(_usage(1, 1))])]  # type: ignore[arg-type]
+        [ChatResponseUpdate(role="assistant", contents=[*contents, Content.from_usage(_usage(1, 1))])]
         for contents in turns
     ]
 
@@ -972,7 +973,7 @@ def _start_task_call(call_id: str) -> Content:
     )
 
 
-async def _run_parent(agent: Agent, session: AgentSession, *, stream: bool = False) -> AgentResponse[Any]:
+async def _run_parent(agent: Agent[Any], session: AgentSession, *, stream: bool = False) -> AgentResponse[Any]:
     if stream:
         return await agent.run("Go.", session=session, stream=True).get_final_response()
     return await agent.run("Go.", session=session)
@@ -993,7 +994,7 @@ def test_task_info_serializes_usage() -> None:
         description="work",
         status=BackgroundTaskStatus.COMPLETED,
         result_text="done",
-        usage_details=_usage(5, 2),  # type: ignore[arg-type]
+        usage_details=_usage(5, 2),
         usage_reported=True,
     )
     restored = BackgroundTaskInfo.from_dict(info.to_dict())
diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py
index c4425674055..7080e292a7f 100644
--- a/python/packages/core/tests/core/test_observability.py
+++ b/python/packages/core/tests/core/test_observability.py
@@ -6663,18 +6663,39 @@ async def test_agent_invoke_span_usage_fallback_excludes_sub_agent_usage(
 ) -> None:
     """Without chat spans, the invoke_agent span falls back to response usage minus sub-agent usage."""
     from agent_framework import ChatMiddlewareLayer, FunctionInvocationLayer
-    from tests.core.conftest import MockBaseChatClient
 
     class _NoChatTelemetryClient(FunctionInvocationLayer[Any], ChatMiddlewareLayer[Any], BaseChatClient[Any]):
+        """Replays scripted responses like MockBaseChatClient, but without ChatTelemetryLayer."""
+
         def __init__(self) -> None:
             super().__init__(middleware=[])
             self.run_responses: list[ChatResponse] = []
             self.streaming_responses: list[list[ChatResponseUpdate]] = []
-            self.call_count = 0
 
-        _inner_get_response = MockBaseChatClient._inner_get_response
-        _get_non_streaming_response = MockBaseChatClient._get_non_streaming_response
-        _get_streaming_response = MockBaseChatClient._get_streaming_response
+        def _inner_get_response(
+            self,
+            *,
+            messages: Sequence[Message],
+            stream: bool,
+            options: Mapping[str, Any],
+            **kwargs: Any,
+        ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+            if stream:
+                scripted = self.streaming_responses.pop(0)
+
+                async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+                    for update in scripted:
+                        yield update
+
+                def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+                    return ChatResponse.from_updates(updates)
+
+                return ResponseStream(_stream(), finalizer=_finalize)
+
+            async def _get() -> ChatResponse:
+                return self.run_responses.pop(0)
+
+            return _get()
 
     span_exporter.clear()
     usage = await _run_parent_with_sub_agent(_NoChatTelemetryClient, stream)