From ca4b6ca1b4a003edf35a551010c717d52e3ce3c7 Mon Sep 17 00:00:00 2001
From: anakin87 <stefanofiorucci@gmail.com>
Date: Wed, 19 Aug 2026 17:12:58 +0200
Subject: [PATCH 1/6] feat: stop_run state key + TokenBudgetHook

---
 haystack/components/agents/agent.py        |  27 ++--
 haystack/hooks/budget/__init__.py          |  15 +++
 haystack/hooks/budget/hooks.py             |  78 +++++++++++
 test/components/agents/test_agent.py       |   2 +
 test/components/agents/test_agent_hooks.py | 144 +++++++++++++++++++++
 test/hooks/budget/__init__.py              |   3 +
 test/hooks/budget/test_hooks.py            |  78 +++++++++++
 7 files changed, 339 insertions(+), 8 deletions(-)
 create mode 100644 haystack/hooks/budget/__init__.py
 create mode 100644 haystack/hooks/budget/hooks.py
 create mode 100644 test/hooks/budget/__init__.py
 create mode 100644 test/hooks/budget/test_hooks.py

diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py
index 7dc4dd583a9..4c920d19929 100644
--- a/haystack/components/agents/agent.py
+++ b/haystack/components/agents/agent.py
@@ -82,6 +82,7 @@
 # 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 end the run after the current step completes.
 # - `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 @@
 #   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},
@@ -744,6 +746,7 @@ def _initialize_fresh_execution(
         state.set("tool_call_counts", {tool.name: 0 for tool in flat_tools})
         state.set("exit_reason", None)
         state.set("continue_run", False)
+        state.set("stop_run", None)
         state.set("tools", flat_tools)
         state.set("hook_context", hook_context or {})
 
@@ -838,9 +841,10 @@ def run(
               `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 supplied by a hook through
+              `stop_run`.
             - Any additional keys defined in the `state_schema`.
         """
         agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs}
@@ -864,9 +868,12 @@ def run(
             while exe_context.counter < self.max_agent_steps:
                 if not self._run_step(exe_context=exe_context, agent_span=span):
                     break
+                # A hook requested a stop: the step it was set in has fully completed, so end the run here
+                if (reason := exe_context.state.data.get("stop_run")) is not None:
+                    exe_context.state.set("exit_reason", reason)
+                    break
             else:
-                # Reached only when the loop ends without a `break`. A `break` means a step already set its own
-                # `exit_reason`, so this branch runs only when `max_agent_steps` is why the Agent stopped.
+                # The Agent reached `max_agent_steps`
                 logger.warning(
                     "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
                     max_agent_steps=self.max_agent_steps,
@@ -922,9 +929,10 @@ async def run_async(
               `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 supplied by a hook through
+              `stop_run`.
             - Any additional keys defined in the `state_schema`.
         """
         agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs}
@@ -948,9 +956,12 @@ async def run_async(
             while exe_context.counter < self.max_agent_steps:
                 if not await self._run_step_async(exe_context=exe_context, agent_span=span):
                     break
+                # A hook requested a stop: the step it was set in has fully completed, so end the run here
+                if (reason := exe_context.state.data.get("stop_run")) is not None:
+                    exe_context.state.set("exit_reason", reason)
+                    break
             else:
-                # Reached only when the loop ends without a `break`. A `break` means a step already set its own
-                # `exit_reason`, so this branch runs only when `max_agent_steps` is why the Agent stopped.
+                # The Agent reached `max_agent_steps`
                 logger.warning(
                     "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
                     max_agent_steps=self.max_agent_steps,
diff --git a/haystack/hooks/budget/__init__.py b/haystack/hooks/budget/__init__.py
new file mode 100644
index 00000000000..12c7426f395
--- /dev/null
+++ b/haystack/hooks/budget/__init__.py
@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# 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)
diff --git a/haystack/hooks/budget/hooks.py b/haystack/hooks/budget/hooks.py
new file mode 100644
index 00000000000..5e4b442e607
--- /dev/null
+++ b/haystack/hooks/budget/hooks.py
@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+from haystack.components.agents.state.state import State
+from haystack.core.serialization import default_from_dict, default_to_dict
+from haystack.utils.experimental import _experimental
+
+
+@_experimental
+class TokenBudgetHook:
+    """
+    Stop an Agent run when its token usage reaches a configured budget.
+
+    The hook runs at the `after_tool` hook point and checks the cumulative `total_tokens` recorded in the Agent state.
+    When the configured budget is reached, the Agent stops with the exit reason `"token_budget_exceeded"`.
+
+    <!-- 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={"after_tool": [TokenBudgetHook(max_total_tokens=100_000)]},
+    )
+
+    result = agent.run(messages=[...])
+    ```
+
+    The budget is checked after a tool step completes, so the final token usage may exceed it. Only calls made by the
+    Agent's chat generator contribute to `token_usage`; calls made by tools or other hooks are not included.
+    """
+
+    allowed_hook_points = ("after_tool",)
+
+    def __init__(self, *, max_total_tokens: int) -> None:
+        """
+        Create a token budget hook.
+
+        :param max_total_tokens: Maximum cumulative token usage before the Agent is 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
+
+    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.
+        """
+        total_tokens = state.data.get("token_usage", {}).get("total_tokens", 0)
+        if total_tokens >= self.max_total_tokens:
+            state.set("stop_run", "token_budget_exceeded")
+
+    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)
+
+    @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)
diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py
index 568c0486c05..16d6d82cdc4 100644
--- a/test/components/agents/test_agent.py
+++ b/test/components/agents/test_agent.py
@@ -221,6 +221,7 @@ def test_state_schema_resolution(self, weather_tool):
             "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},
@@ -1227,6 +1228,7 @@ def test_tracing_span_run(self, spying_tracer, weather_tool):
                     "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]",
diff --git a/test/components/agents/test_agent_hooks.py b/test/components/agents/test_agent_hooks.py
index 1ff08dd5069..0bf093c5eff 100644
--- a/test/components/agents/test_agent_hooks.py
+++ b/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.
@@ -194,6 +199,10 @@ def test_continue_run_is_not_exposed_as_output(self):
         result = agent.run(messages=[ChatMessage.from_user("hi")])
         assert "continue_run" not in result
 
+    def test_stop_run_is_a_reserved_state_schema_key(self):
+        with pytest.raises(ValueError):
+            Agent(chat_generator=MockChatGenerator(), state_schema={"stop_run": {"type": str}})
+
 
 class TestBeforeRunHook:
     def test_tools_are_available_before_the_first_step(self):
@@ -527,6 +536,130 @@ def audit(state: State) -> None:
         assert agent.chat_generator.run.call_count == 1
 
 
+class TestStopRun:
+    def test_stop_from_after_tool_ends_run_after_the_step(self):
+        agent = _agent(MockChatGenerator(), tools=[save], 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 agent.chat_generator.run.call_count == 1
+        assert result["exit_reason"] == "budget_exceeded"
+        assert result["last_message"].tool_call_result is not None
+        assert "stop_run" not in result
+
+    def test_stop_from_before_llm_completes_the_current_step(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 == 1
+        assert result["tool_call_counts"]["save"] == 1
+        assert result["exit_reason"] == "budget_exceeded"
+
+    def test_stop_from_before_tool_completes_the_current_step(self):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_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 agent.chat_generator.run.call_count == 1
+        assert result["tool_call_counts"]["save"] == 1
+        assert result["exit_reason"] == "budget_exceeded"
+
+    def test_stop_from_after_run_has_no_effect(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"
+
+    def test_tool_exit_in_the_same_step_wins_over_stop(self):
+        agent = _agent(
+            MockChatGenerator(),
+            tools=[final_answer],
+            exit_conditions=["final_answer"],
+            hooks={"after_tool": [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_stop_wins_over_continue_run_from_on_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_stop_on_the_last_allowed_step_keeps_the_stop_reason(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"] == "budget_exceeded"
+
+    def test_stop_wins_over_max_steps_when_continue_run_cancels_the_exit(self):
+        # always_continue cancels the tool exit on the only allowed step, so without the stop the run would end
+        # with exit_reason "max_agent_steps". The deliberate stop reason wins over the step budget.
+        agent = _agent(
+            MockChatGenerator(),
+            tools=[final_answer],
+            exit_conditions=["final_answer"],
+            max_agent_steps=1,
+            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 result["exit_reason"] == "budget_exceeded"
+
+    def test_after_run_hooks_run_on_a_stop_exit(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_on_a_stop_exit(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 +711,17 @@ def record(state: State) -> None:
         await agent.run_async(messages=[ChatMessage.from_user("hi")])
         assert fired == [1]
 
+    @pytest.mark.asyncio
+    async def test_stop_from_after_tool_ends_run_after_the_step(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]})
diff --git a/test/hooks/budget/__init__.py b/test/hooks/budget/__init__.py
new file mode 100644
index 00000000000..c1764a6e039
--- /dev/null
+++ b/test/hooks/budget/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
diff --git a/test/hooks/budget/test_hooks.py b/test/hooks/budget/test_hooks.py
new file mode 100644
index 00000000000..cc7c07799b5
--- /dev/null
+++ b/test/hooks/budget/test_hooks.py
@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Annotated
+from unittest.mock import MagicMock
+
+import pytest
+
+from haystack.components.agents import Agent
+from haystack.components.generators.chat import MockChatGenerator
+from haystack.dataclasses import ChatMessage, ToolCall
+from haystack.hooks.budget import TokenBudgetHook
+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 _agent(max_total_tokens: int) -> Agent:
+    agent = Agent(
+        chat_generator=MockChatGenerator(),
+        tools=[fetch],
+        hooks={"after_tool": [TokenBudgetHook(max_total_tokens=max_total_tokens)]},
+    )
+    agent.warm_up()
+    return agent
+
+
+class TestTokenBudgetHook:
+    def test_stops_the_run_when_the_budget_is_reached(self):
+        agent = _agent(max_total_tokens=100)
+        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["exit_reason"] == "token_budget_exceeded"
+        assert result["tool_call_counts"]["fetch"] == 2
+        # The stopped run keeps the whole conversation: every requested tool call has its result.
+        requested = [tc for m in result["messages"] for tc in m.tool_calls]
+        answered = [m.tool_call_result.origin for m in result["messages"] if m.tool_call_result is not None]
+        assert answered == requested
+
+    def test_run_ends_normally_under_budget(self):
+        agent = _agent(max_total_tokens=100)
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                _fetch_reply(10),
+                {"replies": [ChatMessage.from_assistant("done", meta={"usage": {"total_tokens": 10}})]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert result["exit_reason"] == "text"
+
+    def test_rejected_outside_after_tool(self):
+        with pytest.raises(ValueError, match="after_tool"):
+            Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100)]})
+
+    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):
+        restored = TokenBudgetHook.from_dict(TokenBudgetHook(max_total_tokens=5000).to_dict())
+        assert restored.max_total_tokens == 5000

From de59efef7c457096f6813f1f45f10fd1f4e02ba0 Mon Sep 17 00:00:00 2001
From: anakin87 <stefanofiorucci@gmail.com>
Date: Thu, 20 Aug 2026 15:07:21 +0200
Subject: [PATCH 2/6] change design

---
 .../pipeline-components/agents-1/hooks.mdx    |  1 +
 haystack/components/agents/agent.py           | 37 ++++++-----
 haystack/hooks/budget/hooks.py                | 21 ++++--
 haystack/hooks/protocol.py                    |  7 ++
 ...un-token-budget-hook-dafd7694063c5791.yaml | 14 ++++
 test/components/agents/test_agent_hooks.py    | 66 ++++++++++---------
 test/hooks/budget/test_hooks.py               | 50 ++++++++++++--
 7 files changed, 138 insertions(+), 58 deletions(-)
 create mode 100644 releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml

diff --git a/docs-website/docs/pipeline-components/agents-1/hooks.mdx b/docs-website/docs/pipeline-components/agents-1/hooks.mdx
index 29a047fb814..ee11eab7188 100644
--- a/docs-website/docs/pipeline-components/agents-1/hooks.mdx
+++ b/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 a string reason to gracefully stop the run (`state.set("stop_run", "my_reason")`). The Agent evaluates it at the step boundary (right after the `before_llm` hooks and before the LLM call), so a requested stop never spends another LLM call, and reports the reason as `exit_reason`. A natural exit in the same step (a text reply or a tool exit condition) keeps its own reason, and so does a run that stops because it ran out of `max_agent_steps` on that step. Setting it in an `after_run` hook has no effect.
 - `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.
diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py
index 4c920d19929..0d64d9ef66b 100644
--- a/haystack/components/agents/agent.py
+++ b/haystack/components/agents/agent.py
@@ -66,7 +66,8 @@
 logger = logging.getLogger(__name__)
 
 # `exit_reason` values the Agent sets when it stops without a tool exit condition: a tool-call-free reply, or the
-# `max_agent_steps` budget running out. A tool exit condition instead reports the tool's name.
+# `max_agent_steps` budget running out. A tool exit condition instead reports the tool's name, and a hook can
+# supply a custom reason through the `stop_run` state key.
 _EXIT_REASON_TEXT = "text"
 _EXIT_REASON_MAX_STEPS = "max_agent_steps"
 
@@ -82,7 +83,8 @@
 # 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 end the run after the current step completes.
+# - `stop_run`: set by a hook to gracefully stop the run: evaluated at the step boundary (right after the
+#   `before_llm` hooks, before the LLM call) and used as the `exit_reason` when the stop ends the run.
 # - `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
@@ -746,7 +748,6 @@ def _initialize_fresh_execution(
         state.set("tool_call_counts", {tool.name: 0 for tool in flat_tools})
         state.set("exit_reason", None)
         state.set("continue_run", False)
-        state.set("stop_run", None)
         state.set("tools", flat_tools)
         state.set("hook_context", hook_context or {})
 
@@ -843,8 +844,8 @@ def run(
             - "exit_reason": Why the Agent stopped, useful for routing the output downstream (e.g. with a
               `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),
-              `"max_agent_steps"` when the Agent hit `max_agent_steps`, or a custom reason supplied by a hook through
-              `stop_run`.
+              `"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}
@@ -868,12 +869,9 @@ def run(
             while exe_context.counter < self.max_agent_steps:
                 if not self._run_step(exe_context=exe_context, agent_span=span):
                     break
-                # A hook requested a stop: the step it was set in has fully completed, so end the run here
-                if (reason := exe_context.state.data.get("stop_run")) is not None:
-                    exe_context.state.set("exit_reason", reason)
-                    break
             else:
-                # The Agent reached `max_agent_steps`
+                # Reached only when the loop ends without a `break`. A `break` means a step already set its own
+                # `exit_reason`, so this branch runs only when `max_agent_steps` is why the Agent stopped.
                 logger.warning(
                     "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
                     max_agent_steps=self.max_agent_steps,
@@ -931,8 +929,8 @@ async def run_async(
             - "exit_reason": Why the Agent stopped, useful for routing the output downstream (e.g. with a
               `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),
-              `"max_agent_steps"` when the Agent hit `max_agent_steps`, or a custom reason supplied by a hook through
-              `stop_run`.
+              `"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}
@@ -956,12 +954,9 @@ async def run_async(
             while exe_context.counter < self.max_agent_steps:
                 if not await self._run_step_async(exe_context=exe_context, agent_span=span):
                     break
-                # A hook requested a stop: the step it was set in has fully completed, so end the run here
-                if (reason := exe_context.state.data.get("stop_run")) is not None:
-                    exe_context.state.set("exit_reason", reason)
-                    break
             else:
-                # The Agent reached `max_agent_steps`
+                # Reached only when the loop ends without a `break`. A `break` means a step already set its own
+                # `exit_reason`, so this branch runs only when `max_agent_steps` is why the Agent stopped.
                 logger.warning(
                     "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
                     max_agent_steps=self.max_agent_steps,
@@ -989,6 +984,10 @@ def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) ->
             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,
@@ -1052,6 +1051,10 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac
             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,
diff --git a/haystack/hooks/budget/hooks.py b/haystack/hooks/budget/hooks.py
index 5e4b442e607..3d97b5cee7b 100644
--- a/haystack/hooks/budget/hooks.py
+++ b/haystack/hooks/budget/hooks.py
@@ -5,6 +5,7 @@
 from typing import Any
 
 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.utils.experimental import _experimental
 
@@ -14,8 +15,8 @@ class TokenBudgetHook:
     """
     Stop an Agent run when its token usage reaches a configured budget.
 
-    The hook runs at the `after_tool` hook point and checks the cumulative `total_tokens` recorded in the Agent state.
-    When the configured budget is reached, the Agent stops with the exit reason `"token_budget_exceeded"`.
+    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"`.
 
     <!-- test-ignore -->
     ```python
@@ -26,17 +27,18 @@ class TokenBudgetHook:
     agent = Agent(
         chat_generator=OpenAIChatGenerator(),
         tools=[web_search],
-        hooks={"after_tool": [TokenBudgetHook(max_total_tokens=100_000)]},
+        hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100_000)]},
     )
 
     result = agent.run(messages=[...])
     ```
 
-    The budget is checked after a tool step completes, so the final token usage may exceed it. Only calls made by the
-    Agent's chat generator contribute to `token_usage`; calls made by tools or other hooks are not included.
+    The budget is checked before each LLM call, so no call is ever made with the budget already exceeded; the final
+    token usage can exceed the budget by at most the last step's usage. Only calls made by the Agent's chat generator
+    contribute to `token_usage`; calls made by tools or other hooks are not included.
     """
 
-    allowed_hook_points = ("after_tool",)
+    allowed_hook_points = ("before_llm",)
 
     def __init__(self, *, max_total_tokens: int) -> None:
         """
@@ -55,7 +57,12 @@ def run(self, state: State) -> None:
 
         :param state: Agent state containing the cumulative token usage.
         """
-        total_tokens = state.data.get("token_usage", {}).get("total_tokens", 0)
+        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:
             state.set("stop_run", "token_budget_exceeded")
 
diff --git a/haystack/hooks/protocol.py b/haystack/hooks/protocol.py
index 0259366894c..41668a441c7 100644
--- a/haystack/hooks/protocol.py
+++ b/haystack/hooks/protocol.py
@@ -26,6 +26,13 @@ class Hook(Protocol):
     `step_count`, `token_usage` and `tool_call_counts` are available; any additional keys defined in the Agent's
     `state_schema` are available too. The same hook object can be registered under multiple hook points.
 
+    Two reserved state keys drive control flow. An `on_exit` hook can set `continue_run` to keep the Agent running
+    instead of stopping on an exit condition. Any hook can set `stop_run` to a string reason to gracefully end the
+    run: the Agent evaluates it at the step boundary (right after the `before_llm` hooks, before the LLM call) and
+    reports it as `exit_reason`. A natural exit in the same step keeps its own reason, as does a run that stops
+    because it ran out of `max_agent_steps` on that step, and setting `stop_run` in an `after_run` hook has no
+    effect.
+
     Implement this protocol directly for stateful hooks (e.g. one wrapping a component), or use the `@hook` decorator to
     wrap a plain `(State) -> None` function.
 
diff --git a/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml b/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml
new file mode 100644
index 00000000000..f4d99ba1fab
--- /dev/null
+++ b/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml
@@ -0,0 +1,14 @@
+---
+features:
+  - |
+    Agent hooks can now gracefully stop a run by setting the ``stop_run`` control flag to a string reason
+    (``state.set("stop_run", "my_reason")``). The Agent evaluates the flag at the step boundary, right after the
+    ``before_llm`` hooks and before the LLM call, so a requested stop never spends another LLM call. The reason is
+    reported in the ``exit_reason`` output; a natural exit in the same step (a text reply or a tool exit condition)
+    keeps its own reason, and so does a run that stops because it ran out of ``max_agent_steps`` on that step.
+    Setting the flag in an ``after_run`` hook has no effect.
+  - |
+    Added ``TokenBudgetHook``, an experimental ``before_llm`` hook that gracefully stops an Agent run once the
+    cumulative token usage recorded in the Agent state reaches a configured budget, with the exit reason
+    ``"token_budget_exceeded"``. It reads the reported ``total_tokens`` when present and falls back to summing
+    input and output tokens across the usage key conventions of the different chat generators.
diff --git a/test/components/agents/test_agent_hooks.py b/test/components/agents/test_agent_hooks.py
index 0bf093c5eff..01e6b098af2 100644
--- a/test/components/agents/test_agent_hooks.py
+++ b/test/components/agents/test_agent_hooks.py
@@ -537,49 +537,51 @@ def audit(state: State) -> None:
 
 
 class TestStopRun:
-    def test_stop_from_after_tool_ends_run_after_the_step(self):
-        agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [stop_for_budget]})
+    @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["exit_reason"] == "budget_exceeded"
+        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_stop_from_before_llm_completes_the_current_step(self):
+    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 == 1
-        assert result["tool_call_counts"]["save"] == 1
+        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"
 
-    def test_stop_from_before_tool_completes_the_current_step(self):
-        agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_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 agent.chat_generator.run.call_count == 1
-        assert result["tool_call_counts"]["save"] == 1
-        assert result["exit_reason"] == "budget_exceeded"
+    def test_check_runs_after_all_before_llm_hooks(self):
+        fired = []
 
-    def test_stop_from_after_run_has_no_effect(self):
-        agent = _agent(MockChatGenerator(), hooks={"after_run": [stop_for_budget]})
+        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 result["exit_reason"] == "text"
+        assert fired == [1]
+        assert agent.chat_generator.run.call_count == 0
+        assert result["exit_reason"] == "budget_exceeded"
 
-    def test_tool_exit_in_the_same_step_wins_over_stop(self):
+    @pytest.mark.parametrize("hook_point", ["after_tool", "on_exit"])
+    def test_natural_exit_wins_in_the_same_step(self, hook_point):
         agent = _agent(
             MockChatGenerator(),
             tools=[final_answer],
             exit_conditions=["final_answer"],
-            hooks={"after_tool": [stop_for_budget]},
+            hooks={hook_point: [stop_for_budget]},
         )
         agent.chat_generator.run = MagicMock(
             return_value={
@@ -589,7 +591,7 @@ def test_tool_exit_in_the_same_step_wins_over_stop(self):
         result = agent.run(messages=[ChatMessage.from_user("q")])
         assert result["exit_reason"] == "final_answer"
 
-    def test_stop_wins_over_continue_run_from_on_exit(self):
+    def test_stop_wins_when_continue_run_cancels_the_exit(self):
         agent = _agent(
             MockChatGenerator(),
             tools=[final_answer],
@@ -606,17 +608,15 @@ def test_stop_wins_over_continue_run_from_on_exit(self):
         assert result["exit_reason"] == "budget_exceeded"
         assert result["last_message"].text == "keep going"
 
-    def test_stop_on_the_last_allowed_step_keeps_the_stop_reason(self):
+    def test_max_steps_wins_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"] == "budget_exceeded"
+        assert result["exit_reason"] == "max_agent_steps"
 
-    def test_stop_wins_over_max_steps_when_continue_run_cancels_the_exit(self):
-        # always_continue cancels the tool exit on the only allowed step, so without the stop the run would end
-        # with exit_reason "max_agent_steps". The deliberate stop reason wins over the step budget.
+    def test_max_steps_wins_over_a_cancelled_exit(self):
         agent = _agent(
             MockChatGenerator(),
             tools=[final_answer],
@@ -630,9 +630,15 @@ def test_stop_wins_over_max_steps_when_continue_run_cancels_the_exit(self):
             }
         )
         result = agent.run(messages=[ChatMessage.from_user("q")])
-        assert result["exit_reason"] == "budget_exceeded"
+        assert result["exit_reason"] == "max_agent_steps"
+
+    def test_stop_from_after_run_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"
 
-    def test_after_run_hooks_run_on_a_stop_exit(self):
+    def test_after_run_hooks_still_run(self):
         agent = _agent(
             MockChatGenerator(), tools=[save], hooks={"after_tool": [stop_for_budget], "after_run": [write_report]}
         )
@@ -643,7 +649,7 @@ def test_after_run_hooks_run_on_a_stop_exit(self):
         assert result["exit_reason"] == "budget_exceeded"
         assert result["last_message"].text == "REPORT"
 
-    def test_on_exit_hooks_do_not_run_on_a_stop_exit(self):
+    def test_on_exit_hooks_do_not_run(self):
         fired = []
 
         def record(state: State) -> None:
@@ -712,7 +718,7 @@ def record(state: State) -> None:
         assert fired == [1]
 
     @pytest.mark.asyncio
-    async def test_stop_from_after_tool_ends_run_after_the_step(self):
+    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"})])]}
diff --git a/test/hooks/budget/test_hooks.py b/test/hooks/budget/test_hooks.py
index cc7c07799b5..a329a323405 100644
--- a/test/hooks/budget/test_hooks.py
+++ b/test/hooks/budget/test_hooks.py
@@ -29,11 +29,16 @@ def _fetch_reply(total_tokens: int) -> dict:
     return {"replies": [message]}
 
 
+def _fetch_reply_without_total(usage: dict) -> dict:
+    message = ChatMessage.from_assistant(tool_calls=[ToolCall("fetch", {"topic": "x"})], meta={"usage": usage})
+    return {"replies": [message]}
+
+
 def _agent(max_total_tokens: int) -> Agent:
     agent = Agent(
         chat_generator=MockChatGenerator(),
         tools=[fetch],
-        hooks={"after_tool": [TokenBudgetHook(max_total_tokens=max_total_tokens)]},
+        hooks={"before_llm": [TokenBudgetHook(max_total_tokens=max_total_tokens)]},
     )
     agent.warm_up()
     return agent
@@ -54,6 +59,43 @@ def test_stops_the_run_when_the_budget_is_reached(self):
         answered = [m.tool_call_result.origin for m in result["messages"] if m.tool_call_result is not None]
         assert answered == requested
 
+    @pytest.mark.parametrize(
+        "usage",
+        [{"input_tokens": 40, "output_tokens": 20}, {"prompt_tokens": 40, "completion_tokens": 20}],
+        ids=["anthropic-style", "openai-style-without-total"],
+    )
+    def test_stops_the_run_when_usage_lacks_total_tokens(self, usage):
+        agent = _agent(max_total_tokens=100)
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                _fetch_reply_without_total(usage),
+                _fetch_reply_without_total(usage),
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 2
+        assert result["exit_reason"] == "token_budget_exceeded"
+
+    def test_stops_a_text_reply_loop_kept_alive_by_continue_run(self):
+        # A critique-style loop: the model replies plain text and an on_exit hook keeps the run going.
+        # The budget check runs before every LLM call, so it covers text steps too.
+        class KeepIterating:
+            def run(self, state):
+                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"
+
     def test_run_ends_normally_under_budget(self):
         agent = _agent(max_total_tokens=100)
         agent.chat_generator.run = MagicMock(
@@ -65,9 +107,9 @@ def test_run_ends_normally_under_budget(self):
         result = agent.run(messages=[ChatMessage.from_user("hi")])
         assert result["exit_reason"] == "text"
 
-    def test_rejected_outside_after_tool(self):
-        with pytest.raises(ValueError, match="after_tool"):
-            Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100)]})
+    def test_rejected_outside_before_llm(self):
+        with pytest.raises(ValueError, match="before_llm"):
+            Agent(chat_generator=MockChatGenerator(), hooks={"after_tool": [TokenBudgetHook(max_total_tokens=100)]})
 
     def test_non_positive_budget_raises(self):
         with pytest.raises(ValueError, match="max_total_tokens"):

From 0e9eecf288fa7a9ee657fb4094f289a3198d8ca9 Mon Sep 17 00:00:00 2001
From: anakin87 <stefanofiorucci@gmail.com>
Date: Thu, 20 Aug 2026 15:44:40 +0200
Subject: [PATCH 3/6] clean up and reorder tests

---
 .../pipeline-components/agents-1/hooks.mdx    |  2 +-
 haystack/components/agents/agent.py           |  6 +-
 haystack/hooks/budget/hooks.py                |  7 +-
 haystack/hooks/protocol.py                    |  7 --
 pydoc/hooks_api.yml                           |  2 +-
 test/components/agents/test_agent.py          |  2 +-
 test/components/agents/test_agent_hooks.py    | 72 +++++++------------
 7 files changed, 34 insertions(+), 64 deletions(-)

diff --git a/docs-website/docs/pipeline-components/agents-1/hooks.mdx b/docs-website/docs/pipeline-components/agents-1/hooks.mdx
index ee11eab7188..d1d51fa3a17 100644
--- a/docs-website/docs/pipeline-components/agents-1/hooks.mdx
+++ b/docs-website/docs/pipeline-components/agents-1/hooks.mdx
@@ -44,7 +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 a string reason to gracefully stop the run (`state.set("stop_run", "my_reason")`). The Agent evaluates it at the step boundary (right after the `before_llm` hooks and before the LLM call), so a requested stop never spends another LLM call, and reports the reason as `exit_reason`. A natural exit in the same step (a text reply or a tool exit condition) keeps its own reason, and so does a run that stops because it ran out of `max_agent_steps` on that step. Setting it in an `after_run` hook has no effect.
+- `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.
diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py
index 0d64d9ef66b..9ba87512bd2 100644
--- a/haystack/components/agents/agent.py
+++ b/haystack/components/agents/agent.py
@@ -66,8 +66,7 @@
 logger = logging.getLogger(__name__)
 
 # `exit_reason` values the Agent sets when it stops without a tool exit condition: a tool-call-free reply, or the
-# `max_agent_steps` budget running out. A tool exit condition instead reports the tool's name, and a hook can
-# supply a custom reason through the `stop_run` state key.
+# `max_agent_steps` budget running out. A tool exit condition instead reports the tool's name.
 _EXIT_REASON_TEXT = "text"
 _EXIT_REASON_MAX_STEPS = "max_agent_steps"
 
@@ -83,8 +82,7 @@
 # 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 gracefully stop the run: evaluated at the step boundary (right after the
-#   `before_llm` hooks, before the LLM call) and used as the `exit_reason` when the stop ends the run.
+# - `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
diff --git a/haystack/hooks/budget/hooks.py b/haystack/hooks/budget/hooks.py
index 3d97b5cee7b..1ad2a2fd7f2 100644
--- a/haystack/hooks/budget/hooks.py
+++ b/haystack/hooks/budget/hooks.py
@@ -18,6 +18,9 @@ class TokenBudgetHook:
     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
@@ -32,10 +35,6 @@ class TokenBudgetHook:
 
     result = agent.run(messages=[...])
     ```
-
-    The budget is checked before each LLM call, so no call is ever made with the budget already exceeded; the final
-    token usage can exceed the budget by at most the last step's usage. Only calls made by the Agent's chat generator
-    contribute to `token_usage`; calls made by tools or other hooks are not included.
     """
 
     allowed_hook_points = ("before_llm",)
diff --git a/haystack/hooks/protocol.py b/haystack/hooks/protocol.py
index 41668a441c7..0259366894c 100644
--- a/haystack/hooks/protocol.py
+++ b/haystack/hooks/protocol.py
@@ -26,13 +26,6 @@ class Hook(Protocol):
     `step_count`, `token_usage` and `tool_call_counts` are available; any additional keys defined in the Agent's
     `state_schema` are available too. The same hook object can be registered under multiple hook points.
 
-    Two reserved state keys drive control flow. An `on_exit` hook can set `continue_run` to keep the Agent running
-    instead of stopping on an exit condition. Any hook can set `stop_run` to a string reason to gracefully end the
-    run: the Agent evaluates it at the step boundary (right after the `before_llm` hooks, before the LLM call) and
-    reports it as `exit_reason`. A natural exit in the same step keeps its own reason, as does a run that stops
-    because it ran out of `max_agent_steps` on that step, and setting `stop_run` in an `after_run` hook has no
-    effect.
-
     Implement this protocol directly for stateful hooks (e.g. one wrapping a component), or use the `@hook` decorator to
     wrap a plain `(State) -> None` function.
 
diff --git a/pydoc/hooks_api.yml b/pydoc/hooks_api.yml
index 4fed6fb1452..ed2a254f67c 100644
--- a/pydoc/hooks_api.yml
+++ b/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/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",
diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py
index 16d6d82cdc4..7c368b0d6c9 100644
--- a/test/components/agents/test_agent.py
+++ b/test/components/agents/test_agent.py
@@ -249,7 +249,7 @@ def test_output_types(self, weather_tool, component_tool, monkeypatch):
             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"),
diff --git a/test/components/agents/test_agent_hooks.py b/test/components/agents/test_agent_hooks.py
index 01e6b098af2..2d2037a8b4c 100644
--- a/test/components/agents/test_agent_hooks.py
+++ b/test/components/agents/test_agent_hooks.py
@@ -199,10 +199,6 @@ def test_continue_run_is_not_exposed_as_output(self):
         result = agent.run(messages=[ChatMessage.from_user("hi")])
         assert "continue_run" not in result
 
-    def test_stop_run_is_a_reserved_state_schema_key(self):
-        with pytest.raises(ValueError):
-            Agent(chat_generator=MockChatGenerator(), state_schema={"stop_run": {"type": str}})
-
 
 class TestBeforeRunHook:
     def test_tools_are_available_before_the_first_step(self):
@@ -537,6 +533,17 @@ def audit(state: State) -> None:
 
 
 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]})
@@ -551,32 +558,14 @@ def test_mid_step_stop_skips_the_next_call(self, hook_point):
         assert result["exit_reason"] == "budget_exceeded"
         assert "stop_run" not in result
 
-    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"
-
-    def test_check_runs_after_all_before_llm_hooks(self):
-        fired = []
-
-        def record(state: State) -> None:
-            fired.append(1)
-
-        agent = _agent(MockChatGenerator(), hooks={"before_llm": [stop_for_budget, hook(record)]})
+    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 fired == [1]
-        assert agent.chat_generator.run.call_count == 0
-        assert result["exit_reason"] == "budget_exceeded"
+        assert result["exit_reason"] == "text"
 
     @pytest.mark.parametrize("hook_point", ["after_tool", "on_exit"])
-    def test_natural_exit_wins_in_the_same_step(self, hook_point):
+    def test_exit_reason_is_the_natural_exit_in_the_same_step(self, hook_point):
         agent = _agent(
             MockChatGenerator(),
             tools=[final_answer],
@@ -591,7 +580,7 @@ def test_natural_exit_wins_in_the_same_step(self, hook_point):
         result = agent.run(messages=[ChatMessage.from_user("q")])
         assert result["exit_reason"] == "final_answer"
 
-    def test_stop_wins_when_continue_run_cancels_the_exit(self):
+    def test_exit_reason_is_the_stop_when_continue_run_cancels_the_exit(self):
         agent = _agent(
             MockChatGenerator(),
             tools=[final_answer],
@@ -608,7 +597,7 @@ def test_stop_wins_when_continue_run_cancels_the_exit(self):
         assert result["exit_reason"] == "budget_exceeded"
         assert result["last_message"].text == "keep going"
 
-    def test_max_steps_wins_on_the_last_step(self):
+    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"})])]}
@@ -616,27 +605,18 @@ def test_max_steps_wins_on_the_last_step(self):
         result = agent.run(messages=[ChatMessage.from_user("hi")])
         assert result["exit_reason"] == "max_agent_steps"
 
-    def test_max_steps_wins_over_a_cancelled_exit(self):
-        agent = _agent(
-            MockChatGenerator(),
-            tools=[final_answer],
-            exit_conditions=["final_answer"],
-            max_agent_steps=1,
-            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 result["exit_reason"] == "max_agent_steps"
+    def test_later_before_llm_hooks_still_run(self):
+        fired = []
 
-    def test_stop_from_after_run_is_ignored(self):
-        agent = _agent(MockChatGenerator(), hooks={"after_run": [stop_for_budget]})
+        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 result["exit_reason"] == "text"
+        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(

From 95c6fc12e44fb0f506a6a0f6e668c1da3bdff6a5 Mon Sep 17 00:00:00 2001
From: anakin87 <stefanofiorucci@gmail.com>
Date: Thu, 20 Aug 2026 16:02:29 +0200
Subject: [PATCH 4/6] better tests for TokenBudgetHook

---
 ...un-token-budget-hook-dafd7694063c5791.yaml |  25 +++--
 test/hooks/budget/test_hooks.py               | 103 +++++++-----------
 2 files changed, 55 insertions(+), 73 deletions(-)

diff --git a/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml b/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml
index f4d99ba1fab..6dc0ef44b02 100644
--- a/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml
+++ b/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml
@@ -1,14 +1,19 @@
 ---
 features:
   - |
-    Agent hooks can now gracefully stop a run by setting the ``stop_run`` control flag to a string reason
-    (``state.set("stop_run", "my_reason")``). The Agent evaluates the flag at the step boundary, right after the
-    ``before_llm`` hooks and before the LLM call, so a requested stop never spends another LLM call. The reason is
-    reported in the ``exit_reason`` output; a natural exit in the same step (a text reply or a tool exit condition)
-    keeps its own reason, and so does a run that stops because it ran out of ``max_agent_steps`` on that step.
-    Setting the flag in an ``after_run`` hook has no effect.
+    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``, an experimental ``before_llm`` hook that gracefully stops an Agent run once the
-    cumulative token usage recorded in the Agent state reaches a configured budget, with the exit reason
-    ``"token_budget_exceeded"``. It reads the reported ``total_tokens`` when present and falls back to summing
-    input and output tokens across the usage key conventions of the different chat generators.
+    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)]},
+        )
diff --git a/test/hooks/budget/test_hooks.py b/test/hooks/budget/test_hooks.py
index a329a323405..1692dc9bbd2 100644
--- a/test/hooks/budget/test_hooks.py
+++ b/test/hooks/budget/test_hooks.py
@@ -2,12 +2,13 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Annotated
+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, ToolCall
 from haystack.hooks.budget import TokenBudgetHook
@@ -29,59 +30,58 @@ def _fetch_reply(total_tokens: int) -> dict:
     return {"replies": [message]}
 
 
-def _fetch_reply_without_total(usage: dict) -> dict:
-    message = ChatMessage.from_assistant(tool_calls=[ToolCall("fetch", {"topic": "x"})], meta={"usage": usage})
-    return {"replies": [message]}
+def _state(usage: dict) -> State:
+    schema = {"token_usage": {"type": dict[str, Any]}, "stop_run": {"type": str}}
+    return State(schema=schema, data={"token_usage": usage})
 
 
-def _agent(max_total_tokens: int) -> Agent:
-    agent = Agent(
-        chat_generator=MockChatGenerator(),
-        tools=[fetch],
-        hooks={"before_llm": [TokenBudgetHook(max_total_tokens=max_total_tokens)]},
+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"],
     )
-    agent.warm_up()
-    return agent
+    def test_stops_when_usage_reaches_the_budget(self, usage):
+        state = _state(usage)
+        TokenBudgetHook(max_total_tokens=100).run(state)
+        assert state.data["stop_run"] == "token_budget_exceeded"
 
+    @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
 
-class TestTokenBudgetHook:
-    def test_stops_the_run_when_the_budget_is_reached(self):
-        agent = _agent(max_total_tokens=100)
+    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):
+        restored = TokenBudgetHook.from_dict(TokenBudgetHook(max_total_tokens=5000).to_dict())
+        assert restored.max_total_tokens == 5000
+
+    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["exit_reason"] == "token_budget_exceeded"
         assert result["tool_call_counts"]["fetch"] == 2
-        # The stopped run keeps the whole conversation: every requested tool call has its result.
-        requested = [tc for m in result["messages"] for tc in m.tool_calls]
-        answered = [m.tool_call_result.origin for m in result["messages"] if m.tool_call_result is not None]
-        assert answered == requested
-
-    @pytest.mark.parametrize(
-        "usage",
-        [{"input_tokens": 40, "output_tokens": 20}, {"prompt_tokens": 40, "completion_tokens": 20}],
-        ids=["anthropic-style", "openai-style-without-total"],
-    )
-    def test_stops_the_run_when_usage_lacks_total_tokens(self, usage):
-        agent = _agent(max_total_tokens=100)
-        agent.chat_generator.run = MagicMock(
-            side_effect=[
-                _fetch_reply_without_total(usage),
-                _fetch_reply_without_total(usage),
-                {"replies": [ChatMessage.from_assistant("done")]},
-            ]
-        )
-        result = agent.run(messages=[ChatMessage.from_user("hi")])
-        assert agent.chat_generator.run.call_count == 2
         assert result["exit_reason"] == "token_budget_exceeded"
 
-    def test_stops_a_text_reply_loop_kept_alive_by_continue_run(self):
-        # A critique-style loop: the model replies plain text and an on_exit hook keeps the run going.
-        # The budget check runs before every LLM call, so it covers text steps too.
+    def test_stops_a_text_only_loop_kept_alive_by_continue_run(self):
         class KeepIterating:
-            def run(self, state):
+            def run(self, state: State) -> None:
                 state.set("continue_run", True)
 
         agent = Agent(
@@ -95,26 +95,3 @@ def run(self, state):
         result = agent.run(messages=[ChatMessage.from_user("hi")])
         assert agent.chat_generator.run.call_count == 2
         assert result["exit_reason"] == "token_budget_exceeded"
-
-    def test_run_ends_normally_under_budget(self):
-        agent = _agent(max_total_tokens=100)
-        agent.chat_generator.run = MagicMock(
-            side_effect=[
-                _fetch_reply(10),
-                {"replies": [ChatMessage.from_assistant("done", meta={"usage": {"total_tokens": 10}})]},
-            ]
-        )
-        result = agent.run(messages=[ChatMessage.from_user("hi")])
-        assert result["exit_reason"] == "text"
-
-    def test_rejected_outside_before_llm(self):
-        with pytest.raises(ValueError, match="before_llm"):
-            Agent(chat_generator=MockChatGenerator(), hooks={"after_tool": [TokenBudgetHook(max_total_tokens=100)]})
-
-    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):
-        restored = TokenBudgetHook.from_dict(TokenBudgetHook(max_total_tokens=5000).to_dict())
-        assert restored.max_total_tokens == 5000

From 774e784df8c58f3fa4af3f1e6a8d5d322130fb52 Mon Sep 17 00:00:00 2001
From: anakin87 <stefanofiorucci@gmail.com>
Date: Thu, 20 Aug 2026 18:40:13 +0200
Subject: [PATCH 5/6] add_final_message param

---
 haystack/hooks/budget/hooks.py  | 11 +++++++++--
 test/hooks/budget/test_hooks.py | 20 +++++++++++++++++---
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/haystack/hooks/budget/hooks.py b/haystack/hooks/budget/hooks.py
index 1ad2a2fd7f2..848d8620e80 100644
--- a/haystack/hooks/budget/hooks.py
+++ b/haystack/hooks/budget/hooks.py
@@ -7,8 +7,11 @@
 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
 
+_FINAL_MESSAGE_TEXT = "The Agent stopped because the token budget was exceeded."
+
 
 @_experimental
 class TokenBudgetHook:
@@ -39,16 +42,18 @@ class TokenBudgetHook:
 
     allowed_hook_points = ("before_llm",)
 
-    def __init__(self, *, max_total_tokens: int) -> None:
+    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:
         """
@@ -64,6 +69,8 @@ def run(self, state: State) -> None:
             total_tokens = _first_numeric(usage, _INPUT_TOKEN_KEYS) + _first_numeric(usage, _OUTPUT_TOKEN_KEYS)
         if total_tokens >= self.max_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]:
         """
@@ -71,7 +78,7 @@ def to_dict(self) -> dict[str, Any]:
 
         :returns: Serialized representation of the hook.
         """
-        return default_to_dict(self, max_total_tokens=self.max_total_tokens)
+        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":
diff --git a/test/hooks/budget/test_hooks.py b/test/hooks/budget/test_hooks.py
index 1692dc9bbd2..119efd0acf0 100644
--- a/test/hooks/budget/test_hooks.py
+++ b/test/hooks/budget/test_hooks.py
@@ -10,8 +10,9 @@
 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, ToolCall
+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")
@@ -31,7 +32,11 @@ def _fetch_reply(total_tokens: int) -> dict:
 
 
 def _state(usage: dict) -> State:
-    schema = {"token_usage": {"type": dict[str, Any]}, "stop_run": {"type": str}}
+    schema = {
+        "token_usage": {"type": dict[str, Any]},
+        "stop_run": {"type": str},
+        "messages": {"type": list[ChatMessage]},
+    }
     return State(schema=schema, data={"token_usage": usage})
 
 
@@ -49,6 +54,7 @@ def test_stops_when_usage_reaches_the_budget(self, usage):
         state = _state(usage)
         TokenBudgetHook(max_total_tokens=100).run(state)
         assert state.data["stop_run"] == "token_budget_exceeded"
+        assert state.data.get("messages") is None
 
     @pytest.mark.parametrize("usage", [{"total_tokens": 99}, {}], ids=["under-budget", "no-usage-reported"])
     def test_does_not_stop_below_the_budget(self, usage):
@@ -56,13 +62,21 @@ def test_does_not_stop_below_the_budget(self, 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):
-        restored = TokenBudgetHook.from_dict(TokenBudgetHook(max_total_tokens=5000).to_dict())
+        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(

From af83ca311719d42f75366d2146062e0a94770353 Mon Sep 17 00:00:00 2001
From: anakin87 <stefanofiorucci@gmail.com>
Date: Fri, 21 Aug 2026 10:51:35 +0200
Subject: [PATCH 6/6] add warning

---
 haystack/hooks/budget/hooks.py  | 9 +++++++++
 test/hooks/budget/test_hooks.py | 7 +++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/haystack/hooks/budget/hooks.py b/haystack/hooks/budget/hooks.py
index 848d8620e80..8a5a3225d11 100644
--- a/haystack/hooks/budget/hooks.py
+++ b/haystack/hooks/budget/hooks.py
@@ -4,12 +4,16 @@
 
 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."
 
 
@@ -68,6 +72,11 @@ def run(self, state: State) -> None:
         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)])
diff --git a/test/hooks/budget/test_hooks.py b/test/hooks/budget/test_hooks.py
index 119efd0acf0..473da1bbe46 100644
--- a/test/hooks/budget/test_hooks.py
+++ b/test/hooks/budget/test_hooks.py
@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 from typing import Annotated, Any
 from unittest.mock import MagicMock
 
@@ -50,11 +51,13 @@ class TestTokenBudgetHook:
         ],
         ids=["total_tokens", "openai-style", "anthropic-style"],
     )
-    def test_stops_when_usage_reaches_the_budget(self, usage):
+    def test_stops_when_usage_reaches_the_budget(self, usage, caplog):
         state = _state(usage)
-        TokenBudgetHook(max_total_tokens=100).run(state)
+        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):