fix(agents): count model calls independently of reused message IDs (#8050)

* fix(agents): count model calls independently of reused message IDs

* docs(agents): clarify budget accounting and changelog

* docs: avoid concurrent changelog insertion conflicts

* docs(agents): condense TokenBudgetMiddleware guidance to fit chain budget

CHANGELOG.md

@@ -73,6 +73,11 @@ This section accumulates work toward the **2.3.0** milestone
   every user message on every model call, so one 100 KB paste blocked every
   thread for about 6 s per model call. The scan is now linear, and it redacts
   exactly the same addresses as before. ([#8006])
+
+- **agents:** Token budgets count separate accepted model calls even when a
+  provider reuses message IDs, so repeated IDs no longer bypass the configured
+  cap. Hard stops preserve already-completed structured-output tool results.
+  ([#8050](https://github.com/bytedance/deer-flow/pull/8050))
   
 - **subagents:** Managed subagent creation now rejects Windows reserved device
   names (`CON`, `PRN`, `AUX`, `NUL`, `COM1-9`, `LPT1-9`) at the shared

CHANGELOG_zh.md

@@ -98,6 +98,10 @@
   `TaskInfo`，即使没有注册生命周期贡献者也会放入。手动 `/compact` 仍收到分离存储，
   id 为 `None`。`deerflow-extension-api` 升级到 0.2.10，已有扩展不受影响。([#7831])
 
+- **智能体：** 即使模型提供方复用消息 ID，Token 预算也会分别统计已接受的模型调用，
+  避免重复 ID 绕过配置的上限；触发硬停止时保留已完成的结构化输出工具结果。
+  ([#8050](https://github.com/bytedance/deer-flow/pull/8050))
+
 - **渠道：** 钉钉 AI 卡片模式（`card_template_id`）下，卡片请求失败（例如缺少卡片权限或钉钉服务异常）时
   不再丢失回复。dingtalk-stream 会记录失败的卡片请求后按成功返回，因此文本形式的运行中提示和最终答案的
   `sampleMarkdown` 兜底都不会执行。现在渠道自行发送卡片请求，并在失败时回退。([#7781])

README.md

@@ -2797,6 +2797,7 @@ Manual compaction also returns HTTP 409 when execution is pending or the checkpo
 The chat header also shows a context-window gauge when the selected model has a positive `context_window` configured. It estimates the latest materialized checkpoint's message tokens and keeps the previous same-thread percentage visible while data refetches, independently of the cumulative token-usage setting.
 
 When `token_budget.enabled` and `token_usage.enabled` are both enabled, each lead-agent run budget includes completed subagent usage, including the final batch. Hidden goal continuations share the budget; a new user run starts fresh.
+The token budget counts separate accepted model calls even when their response message IDs repeat, without double-counting repeated observations or later subagent-usage increments.
 
 ### Sub-Agents
 

backend/packages/harness/deerflow/agents/middlewares/AGENTS.md

@@ -178,7 +178,7 @@ Before changing a later authorization phase, read the [authorization RFC](../../
    Isolate sibling warnings and keep lifecycle hooks topology-stable. Before changing this guard, read
    [Loop detection lifecycle](../../../../../docs/LOOP_DETECTION.md) for
    fallback identity, cleanup/LRU/reset, severity ordering, and test invariants.
-32. **TokenBudgetMiddleware** - `token_budget.enabled`: run-ID budgets survive continuations; missing/invalid IDs reset invocation state. With token usage tracking, place CompletedSubagentUsageMiddleware after the guard: reversed `after_model` hooks backfill terminal child usage before enforcement. Keep TokenUsageMiddleware earlier so current-step attribution reflects calls retained by all guards.
+32. **TokenBudgetMiddleware** - `token_budget.enabled`: run-ID budgets survive continuations; missing/invalid IDs reset invocation state. Confirm per-call usage anywhere in state (reused IDs replace earlier slots); retain child-usage deltas; clear unconfirmed calls on boundaries/errors; hard stops retain completed tool/result pairs. With token usage tracking, place CompletedSubagentUsageMiddleware after the guard: reversed `after_model` hooks backfill terminal child usage before enforcement. Keep TokenUsageMiddleware earlier so current-step attribution reflects calls retained by all guards.
 33. **ModelContentCompatibilityMiddleware** - Read-time sanitizer (`wrap_model_call`/`awrap_model_call`, request view only; never writes state): rewrites URL-sourced `file` blocks (any role, gated on the `url` key; an http(s) URL is preserved when the selected model's serializer accepts it — OpenAI Responses API or Anthropic) and non-fetchable-scheme (non-`http`/`https`/`data`) `image` blocks into text placeholders whose MIME/location metadata passes the shared `neutralize_untrusted_tags` guard, so pre-MCP-fix checkpoints heal on the next turn with no migration. Uses the conversion layer's `resource_placeholder_text`/`model_visible_location` for model-visible text; artifact capture scans state messages, so healing adds neither `resource_links` nor handles. See the MCP guide for conversion and the module docstring for healing; malformed URLs degrade to location-less placeholders instead of raising. Registered after DurableContext/summarization, before the date/coalescer tail.
 34. **Custom middlewares** - *(optional)* Any `custom_middlewares` passed to `build_middlewares` are injected here, before config-declared extensions and the terminal-response/safety/clarification tail
 35. **Configured extension middlewares** - `extensions.middlewares` in `config.yaml` or `extensions_config.json` optionally accepts `module.path:ClassName` strings or `{class, kwargs, apply_to}` objects. `deerflow.reflection.resolve_class` loads `AgentMiddleware` classes; import, class, and constructor errors fail agent creation. `kwargs` must be JSON-compatible; YAML dates/timestamps become ISO strings. Order: built-ins/custom and loop/token guards → extensions → terminal-response/safety/clarification tail. Optional `apply_to` selects both chains (default), lead, or subagent before import/construction. Trusted operator config only: paths instantiate arbitrary code. Gateway skill/MCP toggles preserve it in raw JSON; adding an API write path requires explicit trust-boundary review.

backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py

@@ -3,14 +3,15 @@ Tracks cumulative token usage (input, output, total) across model calls within
 a single agent run and enforces configurable soft-warning and hard-stop
 thresholds.
 Detection strategy:
-  1. After each model response, sum the `usage_metadata` of all `AIMessage`s
-     in the current thread history. This automatically captures tokens from
-     subagents because `TokenUsageMiddleware` retroactively adds them to the
-     history.
+  1. Confirm each observed model response in graph state and count its native
+     usage once. Scan per-message usage deltas for later subagent additions;
+     provider message IDs alone cannot identify distinct model calls.
+     Attempts discarded by outer retry middleware are intentionally not charged;
+     this budget tracks accepted responses, not all billable provider attempts.
   2. If the highest fraction (input, output, or total) >= warn_threshold,
      queue a warning.
-  3. If the highest fraction >= hard_stop_threshold, strip tool calls from
-     every provider surface (structured, raw, and content blocks).
+  3. If the highest fraction >= hard_stop_threshold, strip unexecuted calls
+     from every provider surface, preserving already-returned structured results.
 Warning injection uses the deferred pattern:
   - after_model queues the warning (does NOT mutate state).
   - wrap_model_call injects it as a HumanMessage at the next model call.
@@ -52,10 +53,11 @@ from typing import Any, override
 from langchain.agents import AgentState
 from langchain.agents.middleware import AgentMiddleware
 from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 from langgraph.runtime import Runtime
 
 from deerflow.agents.middlewares._bounded_dict import BoundedDict
+from deerflow.agents.middlewares.model_response import last_ai_message
 from deerflow.agents.middlewares.tool_call_metadata import clone_ai_message_with_tool_calls
 from deerflow.config.token_budget_config import TokenBudgetConfig
 
@@ -74,6 +76,21 @@ class TokenUsage:
     total: int = 0
 
 
+@dataclass(slots=True)
+class _PendingModelCall:
+    """A model call observation recorded by ``wrap_model_call`` and confirmed
+    by ``_apply``.  Holds an immutable usage snapshot and message identity so
+    each independent model call is counted once even when the provider reuses
+    the same ``AIMessage.id`` (which makes the diff-based scan see zero delta).
+    """
+
+    message: AIMessage | None  # short-lived ref for identity matching
+    message_id: str | None  # original provider ID (may differ from state)
+    input_tokens: int  # immutable snapshot — never mutated after construction
+    output_tokens: int
+    completed_call_ids: frozenset[str]
+
+
 class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
     """Enforce per-run token budget limits."""
 
@@ -87,6 +104,10 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
         self._pending_warnings: BoundedDict[str, list[str]] = BoundedDict(1000)
         self._seen_messages: BoundedDict[str, dict[str, tuple[int, int]]] = BoundedDict(1000)
         self._cumulative_usage: BoundedDict[str, TokenUsage] = BoundedDict(1000)
+        # Per-run pending model-call observations recorded by wrap_model_call
+        # and consumed by _apply.  Keyed by run_id and bounded so abandoned
+        # runs cannot leak.  See ``_PendingModelCall``.
+        self._pending_model_calls: BoundedDict[str, _PendingModelCall] = BoundedDict(1000)
         # Stop reason set when the hard-stop fires. NOT cleared by
         # ``_clear_run_state``/``after_agent`` so the executor can consume it
         # after the run returns; bounded so abandoned runs cannot leak.
@@ -110,6 +131,7 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
             self._cumulative_usage.clear()
             self._stop_reason.clear()
             self._fallback_run_ids.clear()
+            self._pending_model_calls.clear()
 
     def consume_stop_reason(self, run_id: str | None) -> str | None:
         """Pop and return the stop reason the hard-stop set for this run.
@@ -174,6 +196,7 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
             self._pending_warnings.pop(run_id, None)
             self._seen_messages.pop(run_id, None)
             self._cumulative_usage.pop(run_id, None)
+            self._pending_model_calls.pop(run_id, None)
 
     @override
     def before_agent(self, state: AgentState, runtime: Runtime) -> None:
@@ -182,11 +205,13 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
 
         # Mark all old messages from previous runs as 'seen' so they don't count toward THIS run's budget
         messages = state.get("messages", [])
-        if not messages:
-            return
-
         run_id = self._get_run_id(runtime)
         with self._lock:
+            # A new invocation must not reuse a pending observation from the
+            # previous invocation (which may have ended without ``after_model``).
+            self._pending_model_calls.pop(run_id, None)
+            if not messages:
+                return
             seen = self._seen_messages.setdefault(run_id, {})
             self._cumulative_usage.setdefault(run_id, TokenUsage())
 
@@ -210,8 +235,11 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
             # A Gateway run re-enters the graph for hidden goal continuations
             # under the same run_id, and they share this run's budget. Keep the
             # usage and warning state; before_agent rebuilds the seen map.
+            # Clear only the pending observation — a continuation must not
+            # reuse the previous invocation's unconfirmed model call.
             with self._lock:
                 self._seen_messages.pop(run_id, None)
+                self._pending_model_calls.pop(run_id, None)
             return
         self._clear_run_state(run_id)
         self._release_fallback_run_id(runtime)
@@ -235,9 +263,13 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
             return new_content
         return f"{content}\n\n{stop_msg}"
 
-    def _build_hard_stop_update(self, msg: AIMessage, stop_msg: str) -> dict[str, Any]:
+    def _build_hard_stop_update(self, msg: AIMessage, stop_msg: str, *, completed_call_ids: frozenset[str] = frozenset()) -> dict[str, Any]:
         """Build the state update dictionary for a hard stop."""
-        stopped_msg = clone_ai_message_with_tool_calls(msg, [], content=self._append_text(msg.content, stop_msg))
+        # Structured-output handlers can return an AIMessage and its completed
+        # ToolMessage together. Dropping those calls would orphan their results
+        # in the next model request; only unexecuted calls need suppression.
+        retained = [call for call in msg.tool_calls if isinstance(call.get("id"), str) and call["id"] in completed_call_ids] if completed_call_ids else []
+        stopped_msg = clone_ai_message_with_tool_calls(msg, retained, content=self._append_text(msg.content, stop_msg))
         return {"messages": [stopped_msg]}
 
     def _apply(self, state: AgentState, runtime: Runtime) -> dict | None:
@@ -245,19 +277,50 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
             return None
 
         messages = state.get("messages", [])
-        if not messages:
-            return None
-
-        last_msg = messages[-1]
-        if not isinstance(last_msg, AIMessage):
-            return None
-
         run_id = self._get_run_id(runtime)
 
         with self._lock:
+            pending = self._pending_model_calls.pop(run_id, None)
+            last = messages[-1] if messages else None
+            # The outer error handler may discard the observed response and
+            # publish a new fallback. Do not match that discarded response to
+            # an old same-ID message still present in history.
+            if isinstance(last, AIMessage) and last.additional_kwargs.get("deerflow_error_fallback") is True and pending is not None and not self._matches_pending(last, pending):
+                pending = None
+            # Reducers replace a reused ID at its original position (A,B,A),
+            # so neither the last item nor the last AIMessage identifies the
+            # current response. Prefer its object anywhere before trying IDs.
+            current = None
+            if pending is not None:
+                current = next((msg for msg in messages if isinstance(msg, AIMessage) and msg is pending.message), None)
+                if current is None:
+                    current = next((msg for msg in reversed(messages) if isinstance(msg, AIMessage) and self._matches_pending(msg, pending)), None)
+            confirmed = current is not None
+            if current is None:
+                current = last if isinstance(last, AIMessage) else None
+            if current is None:
+                return None
+
             seen = self._seen_messages.setdefault(run_id, {})
             usage_accum = self._cumulative_usage.setdefault(run_id, TokenUsage())
 
+            # Confirm the pending model-call observation recorded by
+            # ``wrap_model_call``.  When the provider reuses the same
+            # ``AIMessage.id`` across distinct calls, the diff-based scan
+            # below sees zero delta for the second call and silently drops its
+            # tokens.  The pending observation accounts for the call's full
+            # native usage once, independent of the diff.  The ``seen`` baseline
+            # is then set to this call's native usage so the scan does not
+            # re-count it; the scan still catches retroactive subagent-usage
+            # increments (usage that grows after the message enters state).
+            if confirmed and pending is not None:
+                usage_accum.input += pending.input_tokens
+                usage_accum.output += pending.output_tokens
+                usage_accum.total += pending.input_tokens + pending.output_tokens
+                seen_key = current.id or pending.message_id
+                if seen_key:
+                    seen[seen_key] = (pending.input_tokens, pending.output_tokens)
+
             for msg in messages:
                 if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"):
                     usage = msg.usage_metadata or {}
@@ -313,7 +376,8 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
                 if isinstance(ctx, dict):
                     ctx["stop_reason"] = "token_capped"
                 stop_text = _BUDGET_EXCEEDED_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget)
-                return self._build_hard_stop_update(last_msg, stop_text)
+                completed_call_ids = pending.completed_call_ids if confirmed and pending is not None else frozenset()
+                return self._build_hard_stop_update(current, stop_text, completed_call_ids=completed_call_ids)
 
             if highest_fraction >= self._config.warn_threshold and not self._warned.get(run_id, False):
                 self._warned[run_id] = True
@@ -369,20 +433,93 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
         new_messages = list(messages) + [warning_msg]
         return request.override(messages=new_messages)
 
+    def _record_pending_model_call(self, runtime: Runtime, response: ModelCallResult) -> None:
+        """Snapshot the last AIMessage's native usage after a successful model call.
+
+        The observation is confirmed later by ``_apply`` against the message
+        that actually enters graph state.  This decouples budget accounting from
+        the diff-based ``seen`` scan, which sees zero delta when the provider
+        reuses the same ``AIMessage.id`` across distinct calls.
+
+        The usage values are an immutable snapshot taken at call time; the
+        message reference is consumed in ``after_model`` and uses the same
+        bounded per-run storage as the other guards. Nothing is added to the
+        message — no public marker, no content/tool-call/metadata mutation.
+        """
+        if not self._config.enabled:
+            return
+        self._discard_pending_model_call(runtime)
+        message = last_ai_message(response)
+        if message is None:
+            return
+        usage = message.usage_metadata or {}
+        input_tokens = usage.get("input_tokens", 0)
+        output_tokens = usage.get("output_tokens", 0)
+        if not isinstance(input_tokens, int) or not isinstance(output_tokens, int):
+            return
+        returned = getattr(response, "result", ())
+        completed_call_ids: frozenset[str] = frozenset()
+        if isinstance(returned, (list, tuple)):
+            position = next((index for index in range(len(returned) - 1, -1, -1) if returned[index] is message), len(returned))
+            completed_call_ids = frozenset(result.tool_call_id for result in returned[position + 1 :] if isinstance(result, ToolMessage) and isinstance(result.tool_call_id, str))
+        run_id = self._get_run_id(runtime)
+        with self._lock:
+            self._pending_model_calls[run_id] = _PendingModelCall(
+                message=message,
+                message_id=message.id,
+                # Match the delta scan: malformed negative usage must never
+                # subtract tokens already charged to this run.
+                input_tokens=max(0, input_tokens),
+                output_tokens=max(0, output_tokens),
+                completed_call_ids=completed_call_ids,
+            )
+
+    @staticmethod
+    def _matches_pending(message: AIMessage, pending: _PendingModelCall) -> bool:
+        """Check whether a state message is the pending call's result.
+
+        Matching is by object identity first (the reducer may keep the same
+        object), then by provider message ID (covers ``model_copy`` and
+        in-place reducer replacement).  Content alone is never compared: two
+        independent calls can produce identical text.
+        The ID fallback assumes the observed response entered graph state; it
+        cannot distinguish a discarded response from historical same-ID state.
+        """
+        if message is pending.message:
+            return True
+        pending_id = pending.message_id or (pending.message.id if pending.message is not None else None)
+        if pending_id is not None and message.id == pending_id:
+            return True
+        return False
+
+    def _discard_pending_model_call(self, runtime: Runtime) -> None:
+        if self._config.enabled:
+            run_id = self._get_run_id(runtime)
+            with self._lock:
+                self._pending_model_calls.pop(run_id, None)
+
     @override
     def wrap_model_call(self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ModelCallResult:
         warnings = self._drain_pending_warnings(request.runtime)
         try:
-            return handler(self._inject_warnings(request, warnings))
-        except Exception:
-            self._restore_pending_warnings(request.runtime, warnings)
+            response = handler(self._inject_warnings(request, warnings))
+            self._record_pending_model_call(request.runtime, response)
+            return response
+        except BaseException as exc:
+            self._discard_pending_model_call(request.runtime)
+            if isinstance(exc, Exception):
+                self._restore_pending_warnings(request.runtime, warnings)
             raise
 
     @override
     async def awrap_model_call(self, request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]]) -> ModelCallResult:
         warnings = self._drain_pending_warnings(request.runtime)
         try:
-            return await handler(self._inject_warnings(request, warnings))
-        except Exception:
-            self._restore_pending_warnings(request.runtime, warnings)
+            response = await handler(self._inject_warnings(request, warnings))
+            self._record_pending_model_call(request.runtime, response)
+            return response
+        except BaseException as exc:
+            self._discard_pending_model_call(request.runtime)
+            if isinstance(exc, Exception):
+                self._restore_pending_warnings(request.runtime, warnings)
             raise

backend/tests/test_token_budget_reused_message_ids.py

@@ -0,0 +1,635 @@
+"""Regression tests for token budget accounting when a provider reuses
+``AIMessage.id`` across distinct model calls in the same Run.
+
+The bug: ``TokenBudgetMiddleware._apply`` keyed its ``seen`` map by
+``AIMessage.id`` and counted only the *diff* since the last observation of
+that ID.  When two independent model calls returned the same provider message
+ID, the second call's diff was zero and its tokens were silently dropped from
+the budget — the run could keep executing tools well past the cap.
+
+The fix: ``wrap_model_call`` records a short-lived *pending observation*
+(native usage snapshot + message identity) after each successful model call.
+``_apply`` confirms that observation against its response anywhere in state and,
+on match, accounts the call's full native usage once — independent of what the
+``seen`` diff says.  The existing diff scan is preserved for retroactive
+subagent-usage increments.
+
+These tests use the real ``create_agent`` graph with scripted model responses
+(no network, no paid models).  The graph and middleware are production code;
+only the model is fake.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import MagicMock
+
+import pytest
+from langchain.agents import create_agent
+from langchain.agents.middleware.types import ModelResponse
+from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langchain_core.tools import tool as as_tool
+
+from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
+from deerflow.config.token_budget_config import TokenBudgetConfig
+
+# --------------------------------------------------------------------------- #
+# Shared fakes and helpers
+# --------------------------------------------------------------------------- #
+
+
+class _ToolCallingFakeModel(FakeMessagesListChatModel):
+    def bind_tools(self, tools, *, tool_choice=None, **kwargs):
+        return self
+
+
+def _make_runtime(thread_id="test-thread", run_id="test-run"):
+    runtime = MagicMock()
+    runtime.context = {"thread_id": thread_id, "run_id": run_id}
+    return runtime
+
+
+def _make_request(messages, runtime):
+    """Build a minimal ModelRequest double for direct wrap_model_call tests."""
+    request = MagicMock()
+    request.messages = list(messages)
+    request.runtime = runtime
+
+    def override_fn(messages=None, **kwargs):
+        new_req = MagicMock()
+        new_req.messages = messages if messages is not None else request.messages
+        new_req.runtime = request.runtime
+        return new_req
+
+    request.override = override_fn
+    return request
+
+
+def _probe_tool(executed: list[int]):
+    @as_tool
+    def probe(step: int) -> str:
+        """Record a harmless diagnostic step."""
+        executed.append(step)
+        return "ok"
+
+    return probe
+
+
+def _call(step: int, *, msg_id: str, tokens: int = 600) -> AIMessage:
+    return AIMessage(
+        content=f"step {step}",
+        id=msg_id,
+        tool_calls=[{"name": "probe", "id": f"call-{step}", "args": {"step": step}}],
+        usage_metadata={"input_tokens": tokens - 50, "output_tokens": 50, "total_tokens": tokens},
+    )
+
+
+def _model_response(message: AIMessage) -> ModelResponse:
+    """Wrap a single AIMessage in a ModelResponse, as the real model handler does."""
+    return ModelResponse(result=[message])
+
+
+# --------------------------------------------------------------------------- #
+# Core bug: reused AIMessage.id under-counts the budget
+# --------------------------------------------------------------------------- #
+
+
+class TestReusedMessageIDBudget:
+    """Two model calls that reuse the same ``AIMessage.id`` must each count."""
+
+    @pytest.mark.parametrize("async_mode", [False, True])
+    def test_reused_id_counts_both_calls_and_caps(self, async_mode):
+        """cap=1000, two 600-token calls with the same provider ID.
+
+        Before fix: only 600 accounted, both tools run, no cap.
+        After fix:  1200 accounted, only step 1 runs, ``token_capped``.
+        """
+        executed: list[int] = []
+        probe = _probe_tool(executed)
+        model = _ToolCallingFakeModel(
+            responses=[
+                _call(1, msg_id="provider-reused-id"),
+                _call(2, msg_id="provider-reused-id"),
+                AIMessage(content="done", id="answer"),
+            ]
+        )
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+        graph = create_agent(model=model, tools=[probe], middleware=[mw])
+        ctx = {"thread_id": "budget-repro", "run_id": "budget-repro"}
+
+        if async_mode:
+            asyncio.run(
+                graph.ainvoke(
+                    {"messages": [HumanMessage("perform two steps")]},
+                    context=ctx,
+                    config={"recursion_limit": 20},
+                )
+            )
+        else:
+            graph.invoke(
+                {"messages": [HumanMessage("perform two steps")]},
+                context=ctx,
+                config={"recursion_limit": 20},
+            )
+
+        usage = mw._cumulative_usage.get("budget-repro")
+        assert usage is not None
+        assert usage.total == 1200, f"expected 1200, got {usage.total}"
+        assert executed == [1], f"only step 1 should execute, got {executed}"
+        assert mw.consume_stop_reason("budget-repro") == "token_capped"
+
+    def test_reused_id_same_text_same_usage_still_two_calls(self):
+        """Identical content, usage *and* ID must not be deduplicated: these
+        are two independent model calls that happen to look the same."""
+        executed: list[int] = []
+        probe = _probe_tool(executed)
+        identical = _call(1, msg_id="dup-id")
+        model = _ToolCallingFakeModel(responses=[identical, identical, AIMessage(content="done", id="answer")])
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+        graph = create_agent(model=model, tools=[probe], middleware=[mw])
+        graph.invoke(
+            {"messages": [HumanMessage("perform two steps")]},
+            context={"thread_id": "t", "run_id": "r"},
+            config={"recursion_limit": 20},
+        )
+        usage = mw._cumulative_usage.get("r")
+        assert usage is not None
+        assert usage.total == 1200
+
+    def test_reused_id_different_usage_counts_each_native_amount(self):
+        """When the reused-ID calls have *different* usage, each call's own
+        native usage must be counted — not the diff between them."""
+        executed: list[int] = []
+        probe = _probe_tool(executed)
+        model = _ToolCallingFakeModel(
+            responses=[
+                AIMessage(content="s1", id="shared", tool_calls=[{"name": "probe", "id": "c1", "args": {"step": 1}}], usage_metadata={"input_tokens": 400, "output_tokens": 0, "total_tokens": 400}),
+                AIMessage(content="s2", id="shared", tool_calls=[{"name": "probe", "id": "c2", "args": {"step": 2}}], usage_metadata={"input_tokens": 800, "output_tokens": 0, "total_tokens": 800}),
+                AIMessage(content="done", id="answer"),
+            ]
+        )
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10000))
+        graph = create_agent(model=model, tools=[probe], middleware=[mw])
+        graph.invoke(
+            {"messages": [HumanMessage("two steps")]},
+            context={"thread_id": "t", "run_id": "r"},
+            config={"recursion_limit": 20},
+        )
+        usage = mw._cumulative_usage.get("r")
+        assert usage is not None
+        assert usage.total == 1200  # 400 + 800
+
+
+# --------------------------------------------------------------------------- #
+# Preserved behavior: unique IDs must not be double-counted
+# --------------------------------------------------------------------------- #
+
+
+class TestUniqueIDBehaviorPreserved:
+    """When IDs are unique, the fix must not change accounting."""
+
+    @pytest.mark.parametrize("async_mode", [False, True])
+    def test_unique_ids_behave_as_before(self, async_mode):
+        executed: list[int] = []
+        probe = _probe_tool(executed)
+        model = _ToolCallingFakeModel(
+            responses=[
+                _call(1, msg_id="id-1"),
+                _call(2, msg_id="id-2"),
+                AIMessage(content="done", id="answer"),
+            ]
+        )
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+        graph = create_agent(model=model, tools=[probe], middleware=[mw])
+        ctx = {"thread_id": "t", "run_id": "r"}
+
+        if async_mode:
+            asyncio.run(
+                graph.ainvoke(
+                    {"messages": [HumanMessage("two steps")]},
+                    context=ctx,
+                    config={"recursion_limit": 20},
+                )
+            )
+        else:
+            graph.invoke(
+                {"messages": [HumanMessage("two steps")]},
+                context=ctx,
+                config={"recursion_limit": 20},
+            )
+
+        usage = mw._cumulative_usage.get("r")
+        assert usage is not None
+        assert usage.total == 1200
+        assert executed == [1]
+        assert mw.consume_stop_reason("r") == "token_capped"
+
+    def test_repeated_after_model_does_not_grow(self):
+        """Calling ``_apply`` twice on the same state must not double-count."""
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+        runtime = _make_runtime()
+        msg = AIMessage(id="m1", content="hi", usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500})
+        state = {"messages": [msg]}
+        mw._apply(state, runtime)
+        first = mw._cumulative_usage["test-run"].total
+        mw._apply(state, runtime)
+        second = mw._cumulative_usage["test-run"].total
+        assert first == second == 500
+
+
+# --------------------------------------------------------------------------- #
+# Historical messages + new Run sharing the same ID (direct wrap_model_call)
+# --------------------------------------------------------------------------- #
+
+
+class TestHistoricalSameIDNewRun:
+    """A historical AIMessage seeded by ``before_agent`` must not be counted
+    in the new Run's budget, even when the new call reuses its ID.
+
+    Uses ``wrap_model_call`` + ``_apply`` directly because the message reducer's
+    in-place replacement makes graph-level continuation with reused IDs behave
+    differently from the budget accounting we are testing here.
+    """
+
+    def test_historical_not_counted_new_call_fully_counted(self):
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+        runtime = _make_runtime(run_id="new-run")
+
+        # Historical message from a previous run — seeded by before_agent.
+        historical = AIMessage(
+            id="reused",
+            content="old",
+            usage_metadata={"input_tokens": 300, "output_tokens": 0, "total_tokens": 300},
+        )
+        mw.before_agent({"messages": [historical]}, runtime)
+        assert mw._seen_messages["new-run"].get("reused") == (300, 0)
+
+        # New model call reuses the same ID but with different usage.
+        new_call = AIMessage(
+            id="reused",
+            content="new",
+            usage_metadata={"input_tokens": 600, "output_tokens": 0, "total_tokens": 600},
+        )
+
+        # Simulate the model returning this response via wrap_model_call.
+        request = _make_request([historical], runtime)
+        mw.wrap_model_call(request, lambda req: _model_response(new_call))
+
+        # after_model sees the state with the new call as the last message.
+        mw._apply({"messages": [historical, new_call]}, runtime)
+
+        usage = mw._cumulative_usage["new-run"]
+        # Only the new call's 600 should be counted, not the historical 300
+        # (diff) or 900 (sum).
+        assert usage.total == 600
+
+    def test_historical_same_usage_new_call_zero_diff(self):
+        """Historical and new call share the same ID *and* the same usage.
+        The new call must still be counted (its native usage), not zero."""
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+        runtime = _make_runtime(run_id="r")
+
+        historical = AIMessage(
+            id="shared",
+            content="old",
+            usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500},
+        )
+        mw.before_agent({"messages": [historical]}, runtime)
+
+        new_call = AIMessage(
+            id="shared",
+            content="new",
+            usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500},
+        )
+        request = _make_request([historical], runtime)
+        mw.wrap_model_call(request, lambda req: _model_response(new_call))
+        mw._apply({"messages": [historical, new_call]}, runtime)
+
+        assert mw._cumulative_usage["r"].total == 500
+
+
+# --------------------------------------------------------------------------- #
+# No native ID: reducer assigns one, must be counted once
+# --------------------------------------------------------------------------- #
+
+
+class TestNoNativeID:
+    """A response without an ``id`` gets one from the reducer; it must be
+    counted exactly once."""
+
+    def test_no_id_counted_once_via_graph(self):
+        """Model returns an AIMessage without an ``id``.  The reducer assigns
+        one; the budget must count it exactly once."""
+        executed: list[int] = []
+        probe = _probe_tool(executed)
+        no_id_call = AIMessage(
+            content="step 1",
+            id=None,
+            tool_calls=[{"name": "probe", "id": "call-1", "args": {"step": 1}}],
+            usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500},
+        )
+        model = _ToolCallingFakeModel(responses=[no_id_call, AIMessage(content="done", id="answer")])
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+        graph = create_agent(model=model, tools=[probe], middleware=[mw])
+        graph.invoke(
+            {"messages": [HumanMessage("do step 1")]},
+            context={"thread_id": "no-id-t", "run_id": "no-id-r"},
+            config={"recursion_limit": 20},
+        )
+        usage = mw._cumulative_usage.get("no-id-r")
+        assert usage is not None
+        assert usage.total == 500
+
+
+# --------------------------------------------------------------------------- #
+# Pending confirmation: model_copy, retroactive usage
+# --------------------------------------------------------------------------- #
+
+
+class TestPendingConfirmation:
+    """Verify the observe/confirm mechanism handles edge cases."""
+
+    def test_model_copy_preserving_id_counted_once(self):
+        """When a middleware does ``model_copy()`` (preserving the ID) before
+        ``after_model``, the new cost must still be counted once."""
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+        runtime = _make_runtime()
+        original = AIMessage(
+            id="m1",
+            content="hello",
+            usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500},
+        )
+        copied = original.model_copy()
+        assert copied is not original
+        assert copied.id == original.id
+        mw.wrap_model_call(_make_request([], runtime), lambda request: _model_response(original))
+        state = {"messages": [copied]}
+        mw._apply(state, runtime)
+        usage = mw._cumulative_usage["test-run"]
+        assert usage.total == 500
+
+    def test_retroactive_subagent_usage_only_increments(self):
+        """Model returns 500 tokens; a retroactive subagent bump raises the
+        same message's usage to 800.  Total must be 800, not 500+800."""
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+        runtime = _make_runtime()
+        msg = AIMessage(
+            id="m1",
+            content="hi",
+            usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500},
+        )
+        mw._apply({"messages": [msg]}, runtime)
+        assert mw._cumulative_usage["test-run"].total == 500
+
+        bumped = AIMessage(
+            id="m1",
+            content="hi",
+            usage_metadata={"input_tokens": 800, "output_tokens": 0, "total_tokens": 800},
+        )
+        mw._apply({"messages": [bumped]}, runtime)
+        assert mw._cumulative_usage["test-run"].total == 800  # 500 + 300 increment
+
+
+# --------------------------------------------------------------------------- #
+# Lifecycle: disabled, reset, identity
+# --------------------------------------------------------------------------- #
+
+
+class TestLifecycleAndIdentity:
+    def test_disabled_config_no_side_effects(self):
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=False, max_tokens=1000))
+        runtime = _make_runtime()
+        msg = AIMessage(id="m1", content="hi", usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500})
+        result = mw._apply({"messages": [msg]}, runtime)
+        assert result is None
+        assert mw.wrap_model_call(_make_request([], runtime), lambda request: _model_response(msg)).result == [msg]
+        assert not mw._pending_model_calls
+        assert "test-run" not in mw._cumulative_usage
+
+    def test_reset_clears_all_state(self):
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+        runtime = _make_runtime()
+        msg = AIMessage(id="m1", content="hi", usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500})
+        mw._apply({"messages": [msg]}, runtime)
+        assert "test-run" in mw._cumulative_usage
+        mw.wrap_model_call(_make_request([], runtime), lambda request: _model_response(msg))
+        assert "test-run" in mw._pending_model_calls
+        mw.reset()
+        for store in (mw._cumulative_usage, mw._seen_messages, mw._pending_warnings, mw._warned, mw._pending_model_calls):
+            assert not store
+
+    def test_message_and_tool_call_identity_preserved(self):
+        """The fix must not mutate any AIMessage.id, tool-call ID, content,
+        or provider metadata — it only changes budget accounting."""
+        executed: list[int] = []
+        probe = _probe_tool(executed)
+        responses = [
+            AIMessage(
+                content="step 1",
+                id="provider-x",
+                tool_calls=[{"name": "probe", "id": "toolu_abc", "args": {"step": 1}}],
+                usage_metadata={"input_tokens": 550, "output_tokens": 50, "total_tokens": 600},
+                additional_kwargs={"provider": "anthropic"},
+            ),
+            AIMessage(content="final answer", id="provider-y", usage_metadata={"input_tokens": 100, "output_tokens": 0, "total_tokens": 100}),
+        ]
+        model = _ToolCallingFakeModel(responses=responses)
+        mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+        graph = create_agent(model=model, tools=[probe], middleware=[mw])
+        state = graph.invoke(
+            {"messages": [HumanMessage("do step 1")]},
+            context={"thread_id": "id-t", "run_id": "id-r"},
+            config={"recursion_limit": 20},
+        )
+        msgs = state["messages"]
+        ai_with_tool = next(m for m in msgs if isinstance(m, AIMessage) and m.tool_calls)
+        assert ai_with_tool.id == "provider-x"
+        assert ai_with_tool.tool_calls[0]["id"] == "toolu_abc"
+        assert ai_with_tool.additional_kwargs.get("provider") == "anthropic"
+
+
+@pytest.mark.parametrize("async_mode", [False, True])
+def test_nonconsecutive_reused_id_counts_all_calls_and_stops_current_response(async_mode):
+    executed = []
+    responses = [_call(step, msg_id=identity) for step, identity in enumerate(("a", "b", "a"), 1)]
+    responses.append(AIMessage(content="done", id="answer"))
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1500))
+    graph = create_agent(_ToolCallingFakeModel(responses=responses), tools=[_probe_tool(executed)], middleware=[mw])
+    kwargs = {"context": {"run_id": "r", "thread_id": "t"}, "config": {"recursion_limit": 20}}
+    state = asyncio.run(graph.ainvoke({"messages": [HumanMessage("three steps")]}, **kwargs)) if async_mode else graph.invoke({"messages": [HumanMessage("three steps")]}, **kwargs)
+
+    assert mw._cumulative_usage["r"].total == 1800
+    assert mw.consume_stop_reason("r") == "token_capped"
+    assert kwargs["context"]["stop_reason"] == "token_capped"
+    current = next(message for message in state["messages"] if message.id == "a")
+    previous = next(message for message in state["messages"] if message.id == "b")
+    assert "step 3" in current.content
+    assert "TOKEN BUDGET EXCEEDED" in current.content
+    assert current.tool_calls == []
+    assert previous.tool_calls[0]["id"] == "call-2"
+    assert executed == [1, 2]
+
+
+@pytest.mark.parametrize("async_mode", [False, True])
+def test_empty_history_clears_pending_but_keeps_same_run_budget(async_mode):
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+    runtime = _make_runtime()
+    message = _call(1, msg_id="a")
+    mw._apply({"messages": [message]}, runtime)
+    mw.wrap_model_call(_make_request([], runtime), lambda request: _model_response(message))
+    assert "test-run" in mw._pending_model_calls
+    if async_mode:
+        asyncio.run(mw.abefore_agent({"messages": []}, runtime))
+    else:
+        mw.before_agent({"messages": []}, runtime)
+    assert "test-run" not in mw._pending_model_calls
+    assert mw._cumulative_usage["test-run"].total == 600
+
+
+def test_pending_snapshot_and_retroactive_usage_are_counted_once():
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+    runtime = _make_runtime()
+    message = _call(1, msg_id="a", tokens=500)
+    mw.wrap_model_call(_make_request([], runtime), lambda request: _model_response(message))
+    message.usage_metadata["input_tokens"] += 300
+    message.usage_metadata["total_tokens"] += 300
+    assert mw._pending_model_calls["test-run"].input_tokens == 450
+    for _ in range(2):
+        mw._apply({"messages": [message]}, runtime)
+    assert mw._cumulative_usage["test-run"].total == 800
+    assert not mw._pending_model_calls
+
+
+def test_historical_message_after_tool_result_is_not_rewritten_without_pending():
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+    runtime = _make_runtime()
+    historical = _call(1, msg_id="old", tokens=1200)
+    result = mw._apply({"messages": [historical, ToolMessage(content="ok", tool_call_id="call-1")]}, runtime)
+    assert result is None
+    assert "test-run" not in mw._cumulative_usage
+
+
+@pytest.mark.parametrize("async_mode", [False, True])
+@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError])
+def test_handler_failure_discards_unconfirmed_response(async_mode, failure):
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+    runtime = _make_runtime()
+    request = _make_request([], runtime)
+    mw.wrap_model_call(request, lambda request: _model_response(_call(1, msg_id="a")))
+
+    def broken(request):
+        raise failure()
+
+    async def abroken(request):
+        raise failure()
+
+    with pytest.raises(failure):
+        if async_mode:
+            asyncio.run(mw.awrap_model_call(request, abroken))
+        else:
+            mw.wrap_model_call(request, broken)
+    assert not mw._pending_model_calls
+
+
+def test_response_without_ai_discards_previous_pending():
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+    runtime = _make_runtime()
+    request = _make_request([], runtime)
+    mw.wrap_model_call(request, lambda request: _model_response(_call(1, msg_id="a")))
+    mw.wrap_model_call(request, lambda request: ModelResponse(result=[]))
+    assert not mw._pending_model_calls
+
+
+def test_unrelated_error_fallback_does_not_confirm_historical_same_id():
+    from deerflow.agents.middlewares.llm_error_handling_middleware import LLMErrorHandlingMiddleware
+    from deerflow.config.app_config import AppConfig
+    from deerflow.config.sandbox_config import SandboxConfig
+
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+    runtime = _make_runtime()
+    historical = _call(1, msg_id="a")
+    mw.before_agent({"messages": [historical]}, runtime)
+    mw.wrap_model_call(_make_request([historical], runtime), lambda request: _model_response(_call(2, msg_id="a")))
+    fallback = LLMErrorHandlingMiddleware(app_config=AppConfig(sandbox=SandboxConfig(use="test")))._build_error_fallback_message("failed", error_type="RuntimeError", reason="other", detail="failed")
+    fallback.id = "fallback"
+    assert mw._apply({"messages": [historical, fallback]}, runtime) is None
+    assert mw._cumulative_usage["test-run"].total == 0
+    assert not mw._pending_model_calls
+
+
+@pytest.mark.parametrize("async_mode", [False, True])
+@pytest.mark.parametrize("always_empty", [False, True])
+def test_outer_empty_response_retry_accounts_only_the_accepted_result(monkeypatch, async_mode, always_empty):
+    from deerflow.agents.middlewares.llm_error_handling_middleware import LLMErrorHandlingMiddleware
+    from deerflow.config.app_config import AppConfig
+    from deerflow.config.sandbox_config import SandboxConfig
+
+    empty = AIMessage(content="", id="reused", usage_metadata={"input_tokens": 600, "output_tokens": 0, "total_tokens": 600})
+    accepted = AIMessage(content="accepted", id="reused", usage_metadata={"input_tokens": 600, "output_tokens": 0, "total_tokens": 600})
+    model = _ToolCallingFakeModel(responses=[empty] if always_empty else [empty, accepted])
+    errors = LLMErrorHandlingMiddleware(app_config=AppConfig(sandbox=SandboxConfig(use="test")))
+    monkeypatch.setattr(errors, "_build_retry_delay_ms", lambda *args: 0)
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+    graph = create_agent(model=model, middleware=[errors, mw])
+    kwargs = {"context": {"run_id": "r", "thread_id": "t"}}
+    state = asyncio.run(graph.ainvoke({"messages": [HumanMessage("answer")]}, **kwargs)) if async_mode else graph.invoke({"messages": [HumanMessage("answer")]}, **kwargs)
+    assert mw._cumulative_usage["r"].total == 600
+    assert not mw._pending_model_calls
+    assert mw.consume_stop_reason("r") is None
+    if always_empty:
+        assert state["messages"][-1].additional_kwargs["deerflow_error_fallback"] is True
+    else:
+        assert state["messages"][-1].content == "accepted"
+
+
+@pytest.mark.parametrize("run_id", ["r", None])
+def test_agent_cleanup_drops_pending_for_continuations_and_unidentified_runs(run_id):
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+    runtime = _make_runtime(run_id=run_id)
+    mw.wrap_model_call(_make_request([], runtime), lambda request: _model_response(_call(1, msg_id="a")))
+    assert mw._pending_model_calls
+    mw.after_agent({"messages": []}, runtime)
+    assert not mw._pending_model_calls
+
+
+@pytest.mark.parametrize(("input_tokens", "output_tokens"), [(-100, 0), (-100, 50), (0, -100), (50, -100)])
+def test_negative_provider_usage_cannot_refund_the_run_budget(input_tokens, output_tokens):
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
+    runtime = _make_runtime()
+    previous = _call(1, msg_id="previous")
+    mw._apply({"messages": [previous]}, runtime)
+    unknown = AIMessage(content="provider usage", id="unknown", usage_metadata={"input_tokens": input_tokens, "output_tokens": output_tokens, "total_tokens": input_tokens + output_tokens})
+    mw.wrap_model_call(_make_request([previous], runtime), lambda request: _model_response(unknown))
+    for _ in range(2):
+        mw._apply({"messages": [previous, unknown]}, runtime)
+    assert mw._cumulative_usage["test-run"].total == 600 + max(0, input_tokens) + max(0, output_tokens)
+    assert unknown.usage_metadata["input_tokens"] == input_tokens
+    assert unknown.usage_metadata["output_tokens"] == output_tokens
+
+
+@pytest.mark.parametrize("async_mode", [False, True])
+def test_structured_output_keeps_completed_tool_pair_at_budget_stop(async_mode):
+    from pydantic import BaseModel
+
+    from deerflow.models.claude_provider import ClaudeChatModel
+
+    class Answer(BaseModel):
+        answer: str
+
+    reply = AIMessage(content="", id="structured", tool_calls=[{"name": "Answer", "id": "schema-call", "args": {"answer": "ok"}}], usage_metadata={"input_tokens": 1200, "output_tokens": 0, "total_tokens": 1200})
+    mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
+    graph = create_agent(_ToolCallingFakeModel(responses=[reply]), middleware=[mw], response_format=Answer)
+    kwargs = {"context": {"run_id": "r", "thread_id": "t"}}
+    state = asyncio.run(graph.ainvoke({"messages": [HumanMessage("answer")]}, **kwargs)) if async_mode else graph.invoke({"messages": [HumanMessage("answer")]}, **kwargs)
+
+    assert state["structured_response"].answer == "ok"
+    assert mw._cumulative_usage["r"].total == 1200
+    assert mw.consume_stop_reason("r") == "token_capped"
+    message = next(message for message in state["messages"] if isinstance(message, AIMessage))
+    assert message.tool_calls[0]["id"] == "schema-call"
+    assert "TOKEN BUDGET EXCEEDED" in message.content
+    payload = ClaudeChatModel(model="claude-sonnet-4-5", anthropic_api_key="sk-ant-offline")._get_request_payload(state["messages"] + [HumanMessage("continue")])
+    uses = {block["id"] for turn in payload["messages"] if isinstance(turn["content"], list) for block in turn["content"] if block["type"] == "tool_use"}
+    results = {block["tool_use_id"] for turn in payload["messages"] if isinstance(turn["content"], list) for block in turn["content"] if block["type"] == "tool_result"}
+    assert uses == results == {"schema-call"}