From a5c50a565d75e4584ffeb1845b55891d3ad7c19a Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Thu, 6 Aug 2026 08:57:12 +0200 Subject: [PATCH] Fix so that if a user restores state from a before run hook then the step counter is correctly picked up from the restored state --- haystack/components/agents/agent.py | 88 ++++++++++--------- ...ter-before-run-hooks-39cb6a98cfa5f4f4.yaml | 5 ++ test/components/agents/test_agent_hooks.py | 22 +++++ test/components/agents/test_utils.py | 11 ++- test/hooks/test_invocation.py | 18 ++-- 5 files changed, 91 insertions(+), 53 deletions(-) create mode 100644 releasenotes/notes/sync-agent-step-count-after-before-run-hooks-39cb6a98cfa5f4f4.yaml diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py index b44552b4636..f3f3147ca4e 100644 --- a/haystack/components/agents/agent.py +++ b/haystack/components/agents/agent.py @@ -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: @@ -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( @@ -794,7 +794,7 @@ 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)) @@ -802,15 +802,15 @@ def _select_tools(self, tools: ToolsType | list[str] | None = None) -> ToolsType 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)." @@ -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) 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 @@ -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) @@ -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) 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 @@ -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) @@ -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, @@ -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) @@ -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) @@ -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: @@ -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, @@ -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) @@ -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) @@ -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: @@ -1144,7 +1148,7 @@ 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: @@ -1152,5 +1156,5 @@ async def _continue_after_exit_hooks_async(self, exe_context: _ExecutionContext) 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) diff --git a/releasenotes/notes/sync-agent-step-count-after-before-run-hooks-39cb6a98cfa5f4f4.yaml b/releasenotes/notes/sync-agent-step-count-after-before-run-hooks-39cb6a98cfa5f4f4.yaml new file mode 100644 index 00000000000..b59faa68bcb --- /dev/null +++ b/releasenotes/notes/sync-agent-step-count-after-before-run-hooks-39cb6a98cfa5f4f4.yaml @@ -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. diff --git a/test/components/agents/test_agent_hooks.py b/test/components/agents/test_agent_hooks.py index 41cd07b1474..f038a98d59b 100644 --- a/test/components/agents/test_agent_hooks.py +++ b/test/components/agents/test_agent_hooks.py @@ -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"]) @@ -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): @@ -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: diff --git a/test/components/agents/test_utils.py b/test/components/agents/test_utils.py index 116401e7290..4bbf950014e 100644 --- a/test/components/agents/test_utils.py +++ b/test/components/agents/test_utils.py @@ -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 diff --git a/test/hooks/test_invocation.py b/test/hooks/test_invocation.py index 53c1d1908f5..d9fa6f3f143 100644 --- a/test/hooks/test_invocation.py +++ b/test/hooks/test_invocation.py @@ -45,37 +45,41 @@ class TestRunHooks: def test_runs_all_hooks_for_hook_point_in_order(self): log: list = [] hooks = {"before_llm": [RecordingHook("a", log), RecordingHook("b", log)]} - _run_hooks(hooks, "before_llm", State(schema={})) + _run_hooks(hooks=hooks, hook_point="before_llm", state=State(schema={})) assert log == [("run", "a"), ("run", "b")] def test_only_runs_the_given_hook_point(self): log: list = [] hooks = {"before_llm": [RecordingHook("a", log)], "on_exit": [RecordingHook("b", log)]} - _run_hooks(hooks, "on_exit", State(schema={})) + _run_hooks(hooks=hooks, hook_point="on_exit", state=State(schema={})) assert log == [("run", "b")] def test_no_hooks_for_hook_point_is_noop(self): - _run_hooks({}, "before_llm", State(schema={})) # does not raise + _run_hooks(hooks={}, hook_point="before_llm", state=State(schema={})) # does not raise class TestRunHooksAsync: @pytest.mark.asyncio async def test_awaits_run_async_when_present(self): log: list = [] - await _run_hooks_async({"before_llm": [AsyncRecordingHook("a", log)]}, "before_llm", State(schema={})) + await _run_hooks_async( + hooks={"before_llm": [AsyncRecordingHook("a", log)]}, hook_point="before_llm", state=State(schema={}) + ) assert log == [("run_async", "a")] @pytest.mark.asyncio async def test_falls_back_to_run_when_no_run_async(self): log: list = [] - await _run_hooks_async({"before_llm": [RecordingHook("a", log)]}, "before_llm", State(schema={})) + await _run_hooks_async( + hooks={"before_llm": [RecordingHook("a", log)]}, hook_point="before_llm", state=State(schema={}) + ) assert log == [("run", "a")] @pytest.mark.asyncio async def test_falls_back_to_run_in_worker_thread(self): hook = ThreadRecordingHook() event_loop_thread_id = threading.get_ident() - await _run_hooks_async({"before_llm": [hook]}, "before_llm", State(schema={})) + await _run_hooks_async(hooks={"before_llm": [hook]}, hook_point="before_llm", state=State(schema={})) assert hook.thread_id is not None assert hook.thread_id != event_loop_thread_id @@ -83,5 +87,5 @@ async def test_falls_back_to_run_in_worker_thread(self): async def test_runs_in_order_mixing_sync_and_async(self): log: list = [] hooks = {"before_llm": [AsyncRecordingHook("a", log), RecordingHook("b", log)]} - await _run_hooks_async(hooks, "before_llm", State(schema={})) + await _run_hooks_async(hooks=hooks, hook_point="before_llm", state=State(schema={})) assert log == [("run_async", "a"), ("run", "b")]