Skip to content

fix: Fix State restoration in Agent from a before_run hook #12253

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sjrl merged 1 commit into from
Aug 6, 2026
Merged

fix: Fix State restoration in Agent from a before_run hook #12253

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 46 additions & 42 deletions haystack/components/agents/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -573,7 +573,7 @@ def _warm_up_tools(self) -> None:
"""Warm up the configured tools once."""
if not self._tools_warmed_up:
if self.tools:
warm_up_tools(self.tools)
warm_up_tools(tools=self.tools)
self._tools_warmed_up = True

def _warm_up_hooks(self) -> None:
Expand Down Expand Up @@ -737,8 +737,8 @@ def _initialize_fresh_execution(
if all(m.is_from(ChatRole.SYSTEM) for m in messages):
logger.warning("All messages provided to the Agent component are system messages. This is not recommended.")

selected_tools = self._select_tools(tools)
flat_tools = flatten_tools_or_toolsets(selected_tools)
selected_tools = self._select_tools(tools=tools)
flat_tools = flatten_tools_or_toolsets(tools=selected_tools)
# Validate tool support once for the run (covers both init-time and runtime tools)
if flat_tools and not self._chat_generator_supports_tools:
raise TypeError(
Expand Down Expand Up @@ -794,23 +794,23 @@ def _select_tools(self, tools: ToolsType | list[str] | None = None) -> ToolsType
# Toolsets are spawned into per-run copies (see _spawn_tools / _select_tools_by_name) so concurrent runs
# sharing the same configured Toolset don't corrupt each other's run-scoped state.
if tools is None:
return _spawn_tools(self.tools)
return _spawn_tools(tools=self.tools)

if isinstance(tools, list) and all(isinstance(t, str) for t in tools):
return _select_tools_by_name(self.tools, cast(list[str], tools))

if isinstance(tools, Toolset):
# Per-run tools are not covered by the Agent's own warm_up(), so warm them up here.
# warm_up() is expected to be idempotent, so re-warming on every run is cheap.
warm_up_tools(tools)
return _spawn_tools(tools)
warm_up_tools(tools=tools)
return _spawn_tools(tools=tools)

if isinstance(tools, list):
selected = cast(list[Tool | Toolset], tools) # mypy can't narrow the Union type from isinstance check
# Per-run tools are not covered by the Agent's own warm_up(), so warm them up here.
# warm_up() is expected to be idempotent, so re-warming on every run is cheap.
warm_up_tools(selected)
return _spawn_tools(selected)
warm_up_tools(tools=selected)
return _spawn_tools(tools=selected)

raise TypeError(
"tools must be a list of Tool and/or Toolset objects, a Toolset, or a list of tool names (strings)."
Expand Down Expand Up @@ -871,11 +871,13 @@ def run(
**kwargs,
)

with self._create_agent_span(exe_context.tools) as span:
with self._create_agent_span(tools=exe_context.tools) as span:
span.set_content_tag("haystack.agent.input", agent_inputs)
_run_hooks(self.hooks, BEFORE_RUN, exe_context.state)
_run_hooks(hooks=self.hooks, hook_point=BEFORE_RUN, state=exe_context.state)
# A before_run hook can restore a saved State, so resume the execution counter from its step count.
exe_context.counter = exe_context.state.data.get("step_count", 0)
Comment on lines +877 to +878

@sjrl sjrl Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the new line that fixes the original issue

while exe_context.counter < self.max_agent_steps:
if not self._run_step(exe_context, span):
if not self._run_step(exe_context=exe_context, agent_span=span):
break
else:
# Reached only when the loop ends without a `break`. A `break` means a step already set its own
Expand All @@ -885,8 +887,8 @@ def run(
max_agent_steps=self.max_agent_steps,
)
exe_context.state.set("exit_reason", _EXIT_REASON_MAX_STEPS)
_run_hooks(self.hooks, AFTER_RUN, exe_context.state)
result = _public_outputs(exe_context.state)
_run_hooks(hooks=self.hooks, hook_point=AFTER_RUN, state=exe_context.state)
result = _public_outputs(state=exe_context.state)
if msgs := result.get("messages"):
result["last_message"] = msgs[-1]
span.set_content_tag("haystack.agent.output", result)
Expand Down Expand Up @@ -952,11 +954,13 @@ async def run_async(
**kwargs,
)

with self._create_agent_span(exe_context.tools) as span:
with self._create_agent_span(tools=exe_context.tools) as span:
span.set_content_tag("haystack.agent.input", agent_inputs)
await _run_hooks_async(self.hooks, BEFORE_RUN, exe_context.state)
await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_RUN, state=exe_context.state)
# A before_run hook can restore a saved State, so resume the execution counter from its step count.
exe_context.counter = exe_context.state.data.get("step_count", 0)
Comment on lines +960 to +961

@sjrl sjrl Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same new line in the async method

while exe_context.counter < self.max_agent_steps:
if not await self._run_step_async(exe_context, span):
if not await self._run_step_async(exe_context=exe_context, agent_span=span):
break
else:
# Reached only when the loop ends without a `break`. A `break` means a step already set its own
Expand All @@ -966,8 +970,8 @@ async def run_async(
max_agent_steps=self.max_agent_steps,
)
exe_context.state.set("exit_reason", _EXIT_REASON_MAX_STEPS)
await _run_hooks_async(self.hooks, AFTER_RUN, exe_context.state)
result = _public_outputs(exe_context.state)
await _run_hooks_async(hooks=self.hooks, hook_point=AFTER_RUN, state=exe_context.state)
result = _public_outputs(state=exe_context.state)
if msgs := result.get("messages"):
result["last_message"] = msgs[-1]
span.set_content_tag("haystack.agent.output", result)
Expand All @@ -982,12 +986,12 @@ def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) ->
) as step_span:
# Re-flatten the tools every step so dynamic toolsets (e.g. SearchableToolset) surface tools discovered in
# earlier steps. Validate names here so duplicates fail before starting the step.
current_tools = flatten_tools_or_toolsets(exe_context.tools)
_check_duplicate_tool_names(current_tools)
current_tools = flatten_tools_or_toolsets(tools=exe_context.tools)
_check_duplicate_tool_names(tools=current_tools)
# Expose the current tools to hooks (e.g. ConfirmationHook) via State.
exe_context.state.set("tools", current_tools, handler_override=replace_values)

_run_hooks(self.hooks, BEFORE_LLM, exe_context.state)
_run_hooks(hooks=self.hooks, hook_point=BEFORE_LLM, state=exe_context.state)
chat_generator_inputs = {
"messages": exe_context.state.data["messages"],
**exe_context.chat_generator_inputs,
Expand All @@ -1000,17 +1004,17 @@ def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) ->
llm_span.set_content_tag("haystack.agent.step.llm.output", result)
llm_messages = result["replies"]
exe_context.state.set("messages", llm_messages)
_record_llm_usage(exe_context.state, llm_messages)
_record_context_tokens(exe_context.state, llm_messages)
_record_llm_usage(state=exe_context.state, llm_messages=llm_messages)
_record_context_tokens(state=exe_context.state, llm_messages=llm_messages)

# Stop on the "no tool call" exit: no tools available, or a plain assistant text reply (see _is_text_exit).
if not current_tools or _is_text_exit(llm_messages):
if not current_tools or _is_text_exit(messages=llm_messages):
exe_context.counter += 1
exe_context.state.set("step_count", exe_context.counter)
exe_context.state.set("exit_reason", _EXIT_REASON_TEXT)
return self._continue_after_exit_hooks(exe_context)
return self._continue_after_exit_hooks(exe_context=exe_context)

_run_hooks(self.hooks, BEFORE_TOOL, exe_context.state)
_run_hooks(hooks=self.hooks, hook_point=BEFORE_TOOL, state=exe_context.state)
# Re-read the pending tool calls from State so that any rewrites a before_tool hook made (e.g.
# ConfirmationHook rejecting or modifying calls) are honored by the executor.
pending_tool_call_messages = _pending_tool_call_messages_from_state(exe_context.state)
Expand All @@ -1023,8 +1027,8 @@ def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) ->
}
tool_messages, exe_context.state = _run_tool(**tool_execution_inputs)
exe_context.state.set("messages", tool_messages)
_record_tool_calls(exe_context.state, tool_messages)
_run_hooks(self.hooks, AFTER_TOOL, exe_context.state)
_record_tool_calls(state=exe_context.state, tool_messages=tool_messages)
_run_hooks(hooks=self.hooks, hook_point=AFTER_TOOL, state=exe_context.state)

exe_context.counter += 1
exe_context.state.set("step_count", exe_context.counter)
Expand All @@ -1035,7 +1039,7 @@ def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) ->
)
if exit_condition_tool is not None:
exe_context.state.set("exit_reason", exit_condition_tool)
return self._continue_after_exit_hooks(exe_context)
return self._continue_after_exit_hooks(exe_context=exe_context)
return True

async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: tracing.Span) -> bool:
Expand All @@ -1045,12 +1049,12 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac
) as step_span:
# Re-flatten the tools every step so dynamic toolsets (e.g. SearchableToolset) surface tools discovered in
# earlier steps. Validate names here so duplicates fail before starting the step.
current_tools = flatten_tools_or_toolsets(exe_context.tools)
_check_duplicate_tool_names(current_tools)
current_tools = flatten_tools_or_toolsets(tools=exe_context.tools)
_check_duplicate_tool_names(tools=current_tools)
# Expose the current tools to hooks (e.g. ConfirmationHook) via State.
exe_context.state.set("tools", current_tools, handler_override=replace_values)

await _run_hooks_async(self.hooks, BEFORE_LLM, exe_context.state)
await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_LLM, state=exe_context.state)
chat_generator_inputs = {
"messages": exe_context.state.data["messages"],
**exe_context.chat_generator_inputs,
Expand All @@ -1065,17 +1069,17 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac
llm_span.set_content_tag("haystack.agent.step.llm.output", result)
llm_messages = result["replies"]
exe_context.state.set("messages", llm_messages)
_record_llm_usage(exe_context.state, llm_messages)
_record_context_tokens(exe_context.state, llm_messages)
_record_llm_usage(state=exe_context.state, llm_messages=llm_messages)
_record_context_tokens(state=exe_context.state, llm_messages=llm_messages)

# Stop on the "no tool call" exit: no tools available, or a plain assistant text reply (see _is_text_exit).
if not current_tools or _is_text_exit(llm_messages):
if not current_tools or _is_text_exit(messages=llm_messages):
exe_context.counter += 1
exe_context.state.set("step_count", exe_context.counter)
exe_context.state.set("exit_reason", _EXIT_REASON_TEXT)
return await self._continue_after_exit_hooks_async(exe_context)
return await self._continue_after_exit_hooks_async(exe_context=exe_context)

await _run_hooks_async(self.hooks, BEFORE_TOOL, exe_context.state)
await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_TOOL, state=exe_context.state)
# Re-read the pending tool calls from State so that any rewrites a before_tool hook made (e.g.
# ConfirmationHook rejecting or modifying calls) are honored by the executor.
pending_tool_call_messages = _pending_tool_call_messages_from_state(exe_context.state)
Expand All @@ -1088,8 +1092,8 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac
}
tool_messages, exe_context.state = await _run_tool_async(**tool_execution_inputs)
exe_context.state.set("messages", tool_messages)
_record_tool_calls(exe_context.state, tool_messages)
await _run_hooks_async(self.hooks, AFTER_TOOL, exe_context.state)
_record_tool_calls(state=exe_context.state, tool_messages=tool_messages)
await _run_hooks_async(hooks=self.hooks, hook_point=AFTER_TOOL, state=exe_context.state)

exe_context.counter += 1
exe_context.state.set("step_count", exe_context.counter)
Expand All @@ -1100,7 +1104,7 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac
)
if exit_condition_tool is not None:
exe_context.state.set("exit_reason", exit_condition_tool)
return await self._continue_after_exit_hooks_async(exe_context)
return await self._continue_after_exit_hooks_async(exe_context=exe_context)
return True

def _check_exit_conditions(self, llm_messages: list[ChatMessage], tool_messages: list[ChatMessage]) -> str | None:
Expand Down Expand Up @@ -1144,13 +1148,13 @@ def _continue_after_exit_hooks(self, exe_context: _ExecutionContext) -> bool:
if not self.hooks.get(ON_EXIT):
return False
exe_context.state.set("continue_run", False)
_run_hooks(self.hooks, ON_EXIT, exe_context.state)
_run_hooks(hooks=self.hooks, hook_point=ON_EXIT, state=exe_context.state)
return _consume_continue_run(exe_context.state)

async def _continue_after_exit_hooks_async(self, exe_context: _ExecutionContext) -> bool:
"""Async version of `_continue_after_exit_hooks`."""
if not self.hooks.get(ON_EXIT):
return False
exe_context.state.set("continue_run", False)
await _run_hooks_async(self.hooks, ON_EXIT, exe_context.state)
await _run_hooks_async(hooks=self.hooks, hook_point=ON_EXIT, state=exe_context.state)
return _consume_continue_run(exe_context.state)
5 changes: 5 additions & 0 deletions releasenotes/notes/sync-agent-step-count-after-before-run-hooks-39cb6a98cfa5f4f4.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
fixes:
- |
Keep an Agent's execution counter in sync with ``step_count`` restored by a ``before_run`` hook, so restarted
Agents continue from the saved step instead of resetting the count.
22 changes: 22 additions & 0 deletions test/components/agents/test_agent_hooks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,11 @@ def record_before_run(state: State) -> None:
state.set("trace", ["before_run"])


@hook
def restore_step_count(state: State) -> None:
state.set("step_count", 5)


@hook
def record_on_exit(state: State) -> None:
state.set("trace", ["on_exit"])
Expand Down Expand Up @@ -225,6 +230,14 @@ def test_hooks_run_in_list_order(self):
result = agent.run(messages=[ChatMessage.from_user("hi")])
assert result["trace"] == ["a", "b"]

def test_restored_step_count_is_used_by_execution_context(self):
agent = _agent(MockChatGenerator(), max_agent_steps=10, hooks={"before_run": [restore_step_count]})
agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})

result = agent.run(messages=[ChatMessage.from_user("hi")])

assert result["step_count"] == 6


class TestAfterRunHook:
def test_runs_once_on_text_exit(self):
Expand Down Expand Up @@ -580,6 +593,15 @@ async def test_sync_before_run_and_after_run_hooks_run_in_async_run(self):
result = await agent.run_async(messages=[ChatMessage.from_user("hi")])
assert result["trace"] == ["before_run", "after_run"]

@pytest.mark.asyncio
async def test_restored_step_count_is_used_by_execution_context(self):
agent = _agent(MockChatGenerator(), max_agent_steps=10, hooks={"before_run": [restore_step_count]})
agent.chat_generator.run_async = AsyncMock(return_value={"replies": [ChatMessage.from_assistant("done")]})

result = await agent.run_async(messages=[ChatMessage.from_user("hi")])

assert result["step_count"] == 6

@pytest.mark.asyncio
async def test_async_after_run_hook_runs_when_max_agent_steps_is_exhausted(self):
async def write_report_async(state: State) -> None:
Expand Down
11 changes: 7 additions & 4 deletions test/components/agents/test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,20 +232,23 @@ def test_records_latest_reply_usage_replacing_previous_value(self):
state = self._state()
state.set("context_tokens", 999)
_record_context_tokens(
state, [ChatMessage.from_assistant("Hi", meta={"usage": {"prompt_tokens": 12, "completion_tokens": 3}})]
state=state,
llm_messages=[
ChatMessage.from_assistant("Hi", meta={"usage": {"prompt_tokens": 12, "completion_tokens": 3}})
],
)
assert state.get("context_tokens") == 15

def test_no_messages_leaves_value_untouched(self):
state = self._state()
state.set("context_tokens", 42)
_record_context_tokens(state, [])
_record_context_tokens(state=state, llm_messages=[])
assert state.get("context_tokens") == 42

def test_missing_or_empty_usage_leaves_value_untouched(self):
state = self._state()
_record_context_tokens(state, [ChatMessage.from_assistant("no usage here")])
_record_context_tokens(state, [ChatMessage.from_assistant("empty", meta={"usage": {}})])
_record_context_tokens(state=state, llm_messages=[ChatMessage.from_assistant("no usage here")])
_record_context_tokens(state=state, llm_messages=[ChatMessage.from_assistant("empty", meta={"usage": {}})])
assert state.get("context_tokens") == 0


Expand Down
Loading
Loading