feat: add stop_run State key and TokenBudgetHook

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,13 +82,15 @@
 # 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
 #   (e.g. a `before_llm` hook that triggers compaction once the context grows too large). Kept internal rather than
 #   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 @@ 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 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 @@ 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 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 @@ 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,
@@ -1041,6 +1049,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,

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)

haystack/hooks/budget/hooks.py

@@ -0,0 +1,100 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# 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 @@ 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},
@@ -248,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"),
@@ -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]",