feat: add stop_run State key and TokenBudgetHook (#12411)

docs-website/docs/pipeline-components/agents-1/hooks.mdx

@@ -44,6 +44,7 @@ Registering a hook under an unknown hook point raises a `ValueError` at construc
 The Agent manages a few state keys that hooks interact with. Like the run-metadata keys (`step_count`, `token_usage`, `tool_call_counts`), they are reserved — using any of them in your own `state_schema` raises a `ValueError`. See [State](./state.mdx#schema-definition) for the full list:
 
 - `continue_run`: Set by an `on_exit` hook to keep the Agent running.
+- `stop_run`: Set by any hook to stop the run, read before each LLM call and used as the `exit_reason`.
 - `tools`: The tools available in the current step, for hooks to inspect.
 - `hook_context`: Request-scoped resources passed to `Agent.run(hook_context={...})` / `run_async(hook_context={...})`. Hooks read it with `state.data["hook_context"]` or `state.data.get("hook_context")` — use it for per-request resources such as a user ID, a WebSocket, or a database client. Avoid the plain `state.get("hook_context")` here: `State.get` returns a deep copy of the value, which often fails for the kinds of resources stored in this dict (such as a WebSocket or a database client).
 - `context_tokens`: An approximate count of the tokens currently in the context window, refreshed after each LLM call with that reply's prompt-plus-completion tokens (read it with `state.get("context_tokens")`). Unlike `token_usage`, which accumulates over the run, this is replaced on every call. It's `0` until the first reply that reports usage and doesn't count messages appended after the latest call. A `before_llm` hook can read it to trigger context compaction once it crosses a threshold.

haystack/components/agents/agent.py

@@ -82,6 +82,7 @@ _RUN_METADATA_STATE_KEYS: dict[str, dict[str, Any]] = {
 # Internal state keys the Agent manages for run control and hooks. Like run-metadata keys they are reserved and cannot
 # be redefined by users, but unlike them they are NOT exposed as Agent inputs or outputs (purely internal state):
 # - `continue_run`: set by an `on_exit` hook to keep the Agent running instead of stopping (re-read each exit attempt).
+# - `stop_run`: set by a hook to stop the run, read before each LLM call and used as the `exit_reason`.
 # - `tools`: the flattened tools available in the current step, so a hook can inspect them (e.g. HITL confirmation).
 # - `hook_context`: per-run request-scoped resources passed to `run`/`run_async` for hooks to read.
 # - `context_tokens`: approximate current context-window size, refreshed after each LLM call, for hooks to read
@@ -89,6 +90,7 @@ _RUN_METADATA_STATE_KEYS: dict[str, dict[str, Any]] = {
 #   exposed as an output because it is a best-effort snapshot; see `_record_context_tokens`.
 _INTERNAL_STATE_KEYS: dict[str, dict[str, Any]] = {
     "continue_run": {"type": bool, "handler": replace_values},
+    "stop_run": {"type": str, "handler": replace_values},
     "tools": {"type": list, "handler": replace_values},
     "hook_context": {"type": dict[str, Any], "handler": replace_values},
     "context_tokens": {"type": int, "handler": replace_values},
@@ -838,9 +840,10 @@ class Agent:
               `meta["usage"]`.
             - "tool_call_counts": Mapping of tool name to the number of times that tool was invoked.
             - "exit_reason": Why the Agent stopped, useful for routing the output downstream (e.g. with a
-              `ConditionalRouter`). One of: `"text"` (the model returned a reply with no tool calls), the name of
+              `ConditionalRouter`). This is `"text"` when the model returned a reply with no tool calls, the name of
               the tool that satisfied a tool exit condition (in which case `last_message` is that tool's result),
-              or `"max_agent_steps"` (the Agent hit `max_agent_steps` before meeting an exit condition).
+              `"max_agent_steps"` when the Agent hit `max_agent_steps`, or a custom reason a hook supplied through
+              the `stop_run` state key.
             - Any additional keys defined in the `state_schema`.
         """
         agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs}
@@ -922,9 +925,10 @@ class Agent:
               `meta["usage"]`.
             - "tool_call_counts": Mapping of tool name to the number of times that tool was invoked.
             - "exit_reason": Why the Agent stopped, useful for routing the output downstream (e.g. with a
-              `ConditionalRouter`). One of: `"text"` (the model returned a reply with no tool calls), the name of
+              `ConditionalRouter`). This is `"text"` when the model returned a reply with no tool calls, the name of
               the tool that satisfied a tool exit condition (in which case `last_message` is that tool's result),
-              or `"max_agent_steps"` (the Agent hit `max_agent_steps` before meeting an exit condition).
+              `"max_agent_steps"` when the Agent hit `max_agent_steps`, or a custom reason a hook supplied through
+              the `stop_run` state key.
             - Any additional keys defined in the `state_schema`.
         """
         agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs}
@@ -978,6 +982,10 @@ class Agent:
             exe_context.state.set("tools", current_tools, handler_override=replace_values)
 
             _run_hooks(hooks=self.hooks, hook_point=BEFORE_LLM, state=exe_context.state)
+            # A hook requested a stop: end the run at the step boundary, before spending another LLM call.
+            if (reason := exe_context.state.data.get("stop_run")) is not None:
+                exe_context.state.set("exit_reason", reason)
+                return False
             chat_generator_inputs = {
                 "messages": exe_context.state.data["messages"],
                 **exe_context.chat_generator_inputs,
@@ -1041,6 +1049,10 @@ class Agent:
             exe_context.state.set("tools", current_tools, handler_override=replace_values)
 
             await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_LLM, state=exe_context.state)
+            # A hook requested a stop: end the run at the step boundary, before spending another LLM call.
+            if (reason := exe_context.state.data.get("stop_run")) is not None:
+                exe_context.state.set("exit_reason", reason)
+                return False
             chat_generator_inputs = {
                 "messages": exe_context.state.data["messages"],
                 **exe_context.chat_generator_inputs,

haystack/hooks/budget/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+from typing import TYPE_CHECKING
+
+from lazy_imports import LazyImporter
+
+_import_structure = {"hooks": ["TokenBudgetHook"]}
+
+if TYPE_CHECKING:
+    from .hooks import TokenBudgetHook as TokenBudgetHook
+else:
+    sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)

haystack/hooks/budget/hooks.py

@@ -0,0 +1,100 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+from haystack import logging
+from haystack.components.agents.state.state import State
+from haystack.components.agents.utils import _INPUT_TOKEN_KEYS, _OUTPUT_TOKEN_KEYS, _first_numeric
+from haystack.core.serialization import default_from_dict, default_to_dict
+from haystack.dataclasses import ChatMessage
+from haystack.utils.experimental import _experimental
+
+logger = logging.getLogger(__name__)
+
+
+_FINAL_MESSAGE_TEXT = "The Agent stopped because the token budget was exceeded."
+
+
+@_experimental
+class TokenBudgetHook:
+    """
+    Stop an Agent run when its token usage reaches a configured budget.
+
+    The hook runs at the `before_llm` hook point and checks the cumulative token usage recorded in the Agent state.
+    When the budget is reached, the run ends before the next LLM call with the exit reason `"token_budget_exceeded"`.
+
+    Only calls made by the Agent's chat generator contribute to `token_usage`; calls made by tools or other hooks are
+    not included.
+
+    <!-- test-ignore -->
+    ```python
+    from haystack.components.agents import Agent
+    from haystack.components.generators.chat import OpenAIChatGenerator
+    from haystack.hooks.budget import TokenBudgetHook
+
+    agent = Agent(
+        chat_generator=OpenAIChatGenerator(),
+        tools=[web_search],
+        hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100_000)]},
+    )
+
+    result = agent.run(messages=[...])
+    ```
+    """
+
+    allowed_hook_points = ("before_llm",)
+
+    def __init__(self, *, max_total_tokens: int, add_final_message: bool = False) -> None:
+        """
+        Create a token budget hook.
+
+        :param max_total_tokens: Maximum cumulative token usage before the Agent is stopped.
+        :param add_final_message: Whether to append an assistant message explaining why the Agent stopped.
+        :raises ValueError: If `max_total_tokens` is less than 1.
+        """
+        if max_total_tokens < 1:
+            raise ValueError(f"`max_total_tokens` must be a positive number of tokens, got {max_total_tokens}.")
+        self.max_total_tokens = max_total_tokens
+        self.add_final_message = add_final_message
+
+    def run(self, state: State) -> None:
+        """
+        Stop the Agent if its cumulative token usage has reached the budget.
+
+        :param state: Agent state containing the cumulative token usage.
+        """
+        usage = state.data.get("token_usage") or {}
+        # Not every chat generator reports `total_tokens`, so fall back to summing the input and output keys across
+        # the known naming conventions.
+        total_tokens = _first_numeric(usage, ("total_tokens",))
+        if not total_tokens:
+            total_tokens = _first_numeric(usage, _INPUT_TOKEN_KEYS) + _first_numeric(usage, _OUTPUT_TOKEN_KEYS)
+        if total_tokens >= self.max_total_tokens:
+            logger.warning(
+                "Agent reached its token budget of {max_total_tokens} ({total_tokens} used); requesting a stop.",
+                max_total_tokens=self.max_total_tokens,
+                total_tokens=total_tokens,
+            )
+            state.set("stop_run", "token_budget_exceeded")
+            if self.add_final_message:
+                state.set("messages", [ChatMessage.from_assistant(_FINAL_MESSAGE_TEXT)])
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialize this hook to a dictionary.
+
+        :returns: Serialized representation of the hook.
+        """
+        return default_to_dict(self, max_total_tokens=self.max_total_tokens, add_final_message=self.add_final_message)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "TokenBudgetHook":
+        """
+        Create a hook from its serialized representation.
+
+        :param data: Serialized hook data.
+        :returns: The deserialized hook.
+        """
+        return default_from_dict(cls, data=data)

pydoc/hooks_api.yml

@@ -1,6 +1,6 @@
 loaders:
   - search_path: [../haystack/hooks]
-    modules: ["protocol", "from_function", "compaction/hooks", "compaction/sliding_window",
+    modules: ["protocol", "from_function", "budget/hooks", "compaction/hooks", "compaction/sliding_window",
       "compaction/summarization", "compaction/tool_result_pruning",
       "compaction/types/protocol", "human_in_the_loop/dataclasses", "human_in_the_loop/hooks",
       "human_in_the_loop/policies", "human_in_the_loop/strategies", "human_in_the_loop/user_interfaces",

releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml

@@ -0,0 +1,19 @@
+---
+features:
+  - |
+    Added ``stop_run``, a control flag that lets an Agent hook end a run cleanly instead of raising an error. The
+    hook sets it to a reason of your choice (``state.set("stop_run", "my_reason")``), and the Agent stops without
+    another model call and reports that reason in ``exit_reason``.
+  - |
+    Added ``TokenBudgetHook`` (experimental), a ready-made hook that caps how many tokens an Agent run may spend.
+    The run ends with ``exit_reason`` set to ``"token_budget_exceeded"``, keeping the messages collected so far.
+
+    .. code-block:: python
+
+        from haystack.hooks.budget import TokenBudgetHook
+
+        agent = Agent(
+            chat_generator=OpenAIChatGenerator(model="gpt-5.4-nano"),
+            tools=[search],
+            hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100_000)]},
+        )

test/components/agents/test_agent.py

@@ -221,6 +221,7 @@ class TestAgentInit:
             "tool_call_counts": {"type": dict[str, int], "handler": replace_values},
             "exit_reason": {"type": str, "handler": replace_values},
             "continue_run": {"type": bool, "handler": replace_values},
+            "stop_run": {"type": str, "handler": replace_values},
             "tools": {"type": list, "handler": replace_values},
             "hook_context": {"type": dict[str, Any], "handler": replace_values},
             "context_tokens": {"type": int, "handler": replace_values},
@@ -248,7 +249,7 @@ class TestAgentInit:
             assert internal_key not in agent.__haystack_output__._sockets_dict
 
     def test_reserved_state_schema_keys_raise(self, weather_tool):
-        for reserved in ("step_count", "token_usage", "context_tokens", "tool_call_counts", "exit_reason"):
+        for reserved in ("step_count", "token_usage", "context_tokens", "tool_call_counts", "exit_reason", "stop_run"):
             with pytest.raises(ValueError, match="reserved for Agent internal state"):
                 Agent(
                     chat_generator=MockChatGenerator("Hello"),
@@ -1227,6 +1228,7 @@ class TestAgentTracing:
                     "type": "bool",
                     "handler": "haystack.components.agents.state.state_utils.replace_values",
                 },
+                "stop_run": {"type": "str", "handler": "haystack.components.agents.state.state_utils.replace_values"},
                 "tools": {"type": "list", "handler": "haystack.components.agents.state.state_utils.replace_values"},
                 "hook_context": {
                     "type": "dict[str, typing.Any]",

test/components/agents/test_agent_hooks.py

@@ -129,6 +129,11 @@ def always_continue(state: State) -> None:
     state.set("continue_run", True)
 
 
+@hook
+def stop_for_budget(state: State) -> None:
+    state.set("stop_run", "budget_exceeded")
+
+
 @hook
 def critique(state: State) -> None:
     # Push back on the first final answer to force one more loop.
@@ -527,6 +532,120 @@ class TestOnExitHook:
         assert agent.chat_generator.run.call_count == 1
 
 
+class TestStopRun:
+    def test_before_llm_stop_skips_the_current_call(self):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_llm": [stop_for_budget]})
+        agent.chat_generator.run = MagicMock(
+            return_value={"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]}
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 0
+        assert result["step_count"] == 0
+        assert result["tool_call_counts"]["save"] == 0
+        assert result["exit_reason"] == "budget_exceeded"
+
+    @pytest.mark.parametrize("hook_point", ["before_tool", "after_tool"])
+    def test_mid_step_stop_skips_the_next_call(self, hook_point):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={hook_point: [stop_for_budget]})
+        agent.chat_generator.run = MagicMock(
+            return_value={"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]}
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 1
+        assert result["step_count"] == 1
+        assert result["tool_call_counts"]["save"] == 1
+        assert result["last_message"].tool_call_result is not None
+        assert result["exit_reason"] == "budget_exceeded"
+        assert "stop_run" not in result
+
+    def test_after_run_stop_is_ignored(self):
+        agent = _agent(MockChatGenerator(), hooks={"after_run": [stop_for_budget]})
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert result["exit_reason"] == "text"
+
+    @pytest.mark.parametrize("hook_point", ["after_tool", "on_exit"])
+    def test_exit_reason_is_the_natural_exit_in_the_same_step(self, hook_point):
+        agent = _agent(
+            MockChatGenerator(),
+            tools=[final_answer],
+            exit_conditions=["final_answer"],
+            hooks={hook_point: [stop_for_budget]},
+        )
+        agent.chat_generator.run = MagicMock(
+            return_value={
+                "replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "a"})])]
+            }
+        )
+        result = agent.run(messages=[ChatMessage.from_user("q")])
+        assert result["exit_reason"] == "final_answer"
+
+    def test_exit_reason_is_the_stop_when_continue_run_cancels_the_exit(self):
+        agent = _agent(
+            MockChatGenerator(),
+            tools=[final_answer],
+            exit_conditions=["final_answer"],
+            hooks={"after_tool": [stop_for_budget], "on_exit": [always_continue]},
+        )
+        agent.chat_generator.run = MagicMock(
+            return_value={
+                "replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "a"})])]
+            }
+        )
+        result = agent.run(messages=[ChatMessage.from_user("q")])
+        assert agent.chat_generator.run.call_count == 1
+        assert result["exit_reason"] == "budget_exceeded"
+        assert result["last_message"].text == "keep going"
+
+    def test_exit_reason_is_max_steps_on_the_last_step(self):
+        agent = _agent(MockChatGenerator(), tools=[save], max_agent_steps=1, hooks={"after_tool": [stop_for_budget]})
+        agent.chat_generator.run = MagicMock(
+            return_value={"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]}
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert result["exit_reason"] == "max_agent_steps"
+
+    def test_later_before_llm_hooks_still_run(self):
+        fired = []
+
+        def record(state: State) -> None:
+            fired.append(1)
+
+        agent = _agent(MockChatGenerator(), hooks={"before_llm": [stop_for_budget, hook(record)]})
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert fired == [1]
+        assert agent.chat_generator.run.call_count == 0
+        assert result["exit_reason"] == "budget_exceeded"
+
+    def test_after_run_hooks_still_run(self):
+        agent = _agent(
+            MockChatGenerator(), tools=[save], hooks={"after_tool": [stop_for_budget], "after_run": [write_report]}
+        )
+        agent.chat_generator.run = MagicMock(
+            return_value={"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]}
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert result["exit_reason"] == "budget_exceeded"
+        assert result["last_message"].text == "REPORT"
+
+    def test_on_exit_hooks_do_not_run(self):
+        fired = []
+
+        def record(state: State) -> None:
+            fired.append(1)
+
+        agent = _agent(
+            MockChatGenerator(), tools=[save], hooks={"after_tool": [stop_for_budget], "on_exit": [hook(record)]}
+        )
+        agent.chat_generator.run = MagicMock(
+            return_value={"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]}
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert result["exit_reason"] == "budget_exceeded"
+        assert fired == []
+
+
 class TestHookReuseAcrossHookPoints:
     def test_same_hook_under_two_hook_points(self):
         counter = []
@@ -578,6 +697,17 @@ class TestAgentHooksAsync:
         await agent.run_async(messages=[ChatMessage.from_user("hi")])
         assert fired == [1]
 
+    @pytest.mark.asyncio
+    async def test_mid_step_stop_skips_the_next_call(self):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [stop_for_budget]})
+        agent.chat_generator.run_async = AsyncMock(
+            return_value={"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]}
+        )
+        result = await agent.run_async(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run_async.await_count == 1
+        assert result["exit_reason"] == "budget_exceeded"
+        assert result["last_message"].tool_call_result is not None
+
     @pytest.mark.asyncio
     async def test_after_tool_hook_rewrites_results(self):
         agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [offload_tool_results]})

test/hooks/budget/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
+#
+# SPDX-License-Identifier: Apache-2.0

test/hooks/budget/test_hooks.py

@@ -0,0 +1,114 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Annotated, Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from haystack.components.agents import Agent
+from haystack.components.agents.state import State
+from haystack.components.generators.chat import MockChatGenerator
+from haystack.dataclasses import ChatMessage, ChatRole, ToolCall
+from haystack.hooks.budget import TokenBudgetHook
+from haystack.hooks.budget.hooks import _FINAL_MESSAGE_TEXT
+from haystack.tools import tool
+
+pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning")
+
+
+@tool
+def fetch(topic: Annotated[str, "the topic to fetch"]) -> str:
+    """Fetch a document about a topic."""
+    return "DATA"
+
+
+def _fetch_reply(total_tokens: int) -> dict:
+    message = ChatMessage.from_assistant(
+        tool_calls=[ToolCall("fetch", {"topic": "x"})], meta={"usage": {"total_tokens": total_tokens}}
+    )
+    return {"replies": [message]}
+
+
+def _state(usage: dict) -> State:
+    schema = {
+        "token_usage": {"type": dict[str, Any]},
+        "stop_run": {"type": str},
+        "messages": {"type": list[ChatMessage]},
+    }
+    return State(schema=schema, data={"token_usage": usage})
+
+
+class TestTokenBudgetHook:
+    @pytest.mark.parametrize(
+        "usage",
+        [
+            {"total_tokens": 100},
+            {"prompt_tokens": 60, "completion_tokens": 40},
+            {"input_tokens": 60, "output_tokens": 40},
+        ],
+        ids=["total_tokens", "openai-style", "anthropic-style"],
+    )
+    def test_stops_when_usage_reaches_the_budget(self, usage, caplog):
+        state = _state(usage)
+        with caplog.at_level(logging.WARNING):
+            TokenBudgetHook(max_total_tokens=100).run(state)
+        assert state.data["stop_run"] == "token_budget_exceeded"
+        assert state.data.get("messages") is None
+        assert "token budget of 100 (100 used)" in caplog.text
+
+    @pytest.mark.parametrize("usage", [{"total_tokens": 99}, {}], ids=["under-budget", "no-usage-reported"])
+    def test_does_not_stop_below_the_budget(self, usage):
+        state = _state(usage)
+        TokenBudgetHook(max_total_tokens=100).run(state)
+        assert state.data.get("stop_run") is None
+
+    def test_adds_a_final_message(self):
+        state = _state({"total_tokens": 100})
+        TokenBudgetHook(max_total_tokens=100, add_final_message=True).run(state)
+        assert state.data["messages"][-1].text == _FINAL_MESSAGE_TEXT
+        assert state.data["messages"][-1].is_from(ChatRole.ASSISTANT)
+
+    def test_non_positive_budget_raises(self):
+        with pytest.raises(ValueError, match="max_total_tokens"):
+            TokenBudgetHook(max_total_tokens=0)
+
+    def test_to_dict_from_dict_roundtrip(self):
+        hook = TokenBudgetHook(max_total_tokens=5000, add_final_message=True)
+        restored = TokenBudgetHook.from_dict(hook.to_dict())
+        assert restored.max_total_tokens == 5000
+        assert restored.add_final_message is True
+
+    def test_stops_an_agent_run_when_the_budget_is_spent(self):
+        agent = Agent(
+            chat_generator=MockChatGenerator(),
+            tools=[fetch],
+            hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100)]},
+        )
+        agent.warm_up()
+        agent.chat_generator.run = MagicMock(
+            side_effect=[_fetch_reply(60), _fetch_reply(60), {"replies": [ChatMessage.from_assistant("done")]}]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 2
+        assert result["tool_call_counts"]["fetch"] == 2
+        assert result["exit_reason"] == "token_budget_exceeded"
+
+    def test_stops_a_text_only_loop_kept_alive_by_continue_run(self):
+        class KeepIterating:
+            def run(self, state: State) -> None:
+                state.set("continue_run", True)
+
+        agent = Agent(
+            chat_generator=MockChatGenerator(),
+            hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100)], "on_exit": [KeepIterating()]},
+        )
+        agent.warm_up()
+        agent.chat_generator.run = MagicMock(
+            return_value={"replies": [ChatMessage.from_assistant("draft", meta={"usage": {"total_tokens": 60}})]}
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 2
+        assert result["exit_reason"] == "token_budget_exceeded"