chore: import upstream snapshot with attribution

test/components/agents/test_agent_hitl.py

@@ -0,0 +1,385 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Annotated, Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from haystack import component
+from haystack.components.agents import Agent
+from haystack.components.generators.chat import OpenAIChatGenerator
+from haystack.dataclasses import ChatMessage, ToolCall
+from haystack.human_in_the_loop import (
+    AlwaysAskPolicy,
+    BlockingConfirmationStrategy,
+    ConfirmationHook,
+    ConfirmationUIResult,
+    NeverAskPolicy,
+    SimpleConsoleUI,
+)
+from haystack.human_in_the_loop.types import ConfirmationStrategy, ConfirmationUI
+from haystack.tools import Tool, Toolset, create_tool_from_function
+
+
+class MockUserInterface(ConfirmationUI):
+    def __init__(self, ui_result: ConfirmationUIResult) -> None:
+        self.ui_result = ui_result
+
+    def get_user_confirmation(
+        self, tool_name: str, tool_description: str, tool_params: dict[str, Any]
+    ) -> ConfirmationUIResult:
+        return self.ui_result
+
+
+def addition_tool(a: int, b: int) -> int:
+    return a + b
+
+
+@pytest.fixture
+def tools() -> list[Tool]:
+    tool = create_tool_from_function(
+        function=addition_tool, name="addition_tool", description="A tool that adds two integers together."
+    )
+    return [tool]
+
+
+@pytest.fixture
+def confirmation_strategies() -> dict[str, ConfirmationStrategy]:
+    return {
+        "addition_tool": BlockingConfirmationStrategy(
+            confirmation_policy=NeverAskPolicy(), confirmation_ui=SimpleConsoleUI()
+        )
+    }
+
+
+@pytest.fixture
+def confirmation_hook(confirmation_strategies) -> ConfirmationHook:
+    return ConfirmationHook(confirmation_strategies=confirmation_strategies)
+
+
+class TestAgent:
+    def test_confirmation_hook_accepted_on_before_tool(self, tools, confirmation_hook):
+        agent = Agent(chat_generator=MockChatGenerator(), tools=tools, hooks={"before_tool": [confirmation_hook]})
+        assert agent.hooks["before_tool"] == [confirmation_hook]
+
+    @pytest.mark.parametrize("bad_point", ["before_llm", "on_exit"])
+    def test_confirmation_hook_rejected_on_wrong_hook_point(self, tools, confirmation_hook, bad_point):
+        with pytest.raises(ValueError, match="only supports"):
+            Agent(chat_generator=MockChatGenerator(), tools=tools, hooks={bad_point: [confirmation_hook]})
+
+    def test_to_dict(self, tools, confirmation_hook, monkeypatch):
+        monkeypatch.setenv("OPENAI_API_KEY", "test")
+        agent = Agent(chat_generator=OpenAIChatGenerator(), tools=tools, hooks={"before_tool": [confirmation_hook]})
+        agent_dict = agent.to_dict()
+        assert agent_dict == {
+            "type": "haystack.components.agents.agent.Agent",
+            "init_parameters": {
+                "chat_generator": {
+                    "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
+                    "init_parameters": {
+                        "model": "gpt-5-mini",
+                        "streaming_callback": None,
+                        "api_base_url": None,
+                        "organization": None,
+                        "generation_kwargs": {},
+                        "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True},
+                        "timeout": None,
+                        "max_retries": None,
+                        "tools": None,
+                        "tools_strict": False,
+                        "http_client_kwargs": None,
+                    },
+                },
+                "tools": [
+                    {
+                        "type": "haystack.tools.tool.Tool",
+                        "data": {
+                            "name": "addition_tool",
+                            "description": "A tool that adds two integers together.",
+                            "parameters": {
+                                "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
+                                "required": ["a", "b"],
+                                "type": "object",
+                            },
+                            "function": "test_agent_hitl.addition_tool",
+                            "async_function": None,
+                            "outputs_to_string": None,
+                            "inputs_from_state": None,
+                            "outputs_to_state": None,
+                        },
+                    }
+                ],
+                "system_prompt": None,
+                "exit_conditions": ["text"],
+                "state_schema": {},
+                "max_agent_steps": 100,
+                "streaming_callback": None,
+                "raise_on_tool_invocation_failure": False,
+                "tool_concurrency_limit": 4,
+                "tool_streaming_callback_passthrough": False,
+                "required_variables": None,
+                "user_prompt": None,
+                "hooks": {
+                    "before_tool": [
+                        {
+                            "type": "haystack.human_in_the_loop.hooks.ConfirmationHook",
+                            "init_parameters": {
+                                "confirmation_strategies": {
+                                    "addition_tool": {
+                                        "type": "haystack.human_in_the_loop.strategies.BlockingConfirmationStrategy",
+                                        "init_parameters": {
+                                            "confirmation_policy": {
+                                                "type": "haystack.human_in_the_loop.policies.NeverAskPolicy",
+                                                "init_parameters": {},
+                                            },
+                                            "confirmation_ui": {
+                                                "type": "haystack.human_in_the_loop.user_interfaces.SimpleConsoleUI",
+                                                "init_parameters": {},
+                                            },
+                                            "reject_template": "Tool execution for '{tool_name}' was rejected by "
+                                            "the user.",
+                                            "modify_template": "The parameters for tool '{tool_name}' were updated "
+                                            "by the user to:\n{final_tool_params}",
+                                            "user_feedback_template": "With user feedback: {feedback}",
+                                        },
+                                    }
+                                }
+                            },
+                        }
+                    ]
+                },
+            },
+        }
+
+    def test_from_dict(self, tools, confirmation_hook, monkeypatch):
+        monkeypatch.setenv("OPENAI_API_KEY", "test")
+        agent = Agent(chat_generator=OpenAIChatGenerator(), tools=tools, hooks={"before_tool": [confirmation_hook]})
+        deserialized_agent = Agent.from_dict(agent.to_dict())
+        assert deserialized_agent.to_dict() == agent.to_dict()
+        assert isinstance(deserialized_agent.chat_generator, OpenAIChatGenerator)
+        assert len(deserialized_agent.tools) == 1
+        assert deserialized_agent.tools[0].name == "addition_tool"
+        assert deserialized_agent.tool_concurrency_limit == agent.tool_concurrency_limit
+        assert deserialized_agent.tool_streaming_callback_passthrough == agent.tool_streaming_callback_passthrough
+        hook = deserialized_agent.hooks["before_tool"][0]
+        assert isinstance(hook, ConfirmationHook)
+        assert isinstance(hook.confirmation_strategies["addition_tool"], BlockingConfirmationStrategy)
+        assert isinstance(hook.confirmation_strategies["addition_tool"].confirmation_policy, NeverAskPolicy)
+        assert isinstance(hook.confirmation_strategies["addition_tool"].confirmation_ui, SimpleConsoleUI)
+
+    @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
+    @pytest.mark.integration
+    def test_run_blocking_confirmation_strategy_modify(self, tools):
+        confirmation_hook = ConfirmationHook(
+            confirmation_strategies={
+                "addition_tool": BlockingConfirmationStrategy(
+                    confirmation_policy=AlwaysAskPolicy(),
+                    confirmation_ui=MockUserInterface(
+                        ConfirmationUIResult(action="modify", new_tool_params={"a": 2, "b": 3})
+                    ),
+                )
+            }
+        )
+        agent = Agent(
+            chat_generator=OpenAIChatGenerator(model="gpt-4o-mini"),
+            tools=tools,
+            hooks={"before_tool": [confirmation_hook]},
+        )
+        result = agent.run([ChatMessage.from_user("What is 2+2?")])
+
+        assert isinstance(result["last_message"], ChatMessage)
+        assert result["last_message"].text is not None
+        assert "5" in result["last_message"].text
+        # Auto-populated run-metadata outputs: at least one tool call plus a final answer.
+        assert result["step_count"] >= 2
+        assert result["tool_call_counts"]["addition_tool"] >= 1
+        assert result["token_usage"]["prompt_tokens"] > 0
+        assert result["token_usage"]["completion_tokens"] > 0
+        assert result["token_usage"]["total_tokens"] > 0
+
+    @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
+    @pytest.mark.integration
+    @pytest.mark.asyncio
+    async def test_run_async_blocking_confirmation_strategy_modify(self, tools):
+        confirmation_hook = ConfirmationHook(
+            confirmation_strategies={
+                "addition_tool": BlockingConfirmationStrategy(
+                    confirmation_policy=AlwaysAskPolicy(),
+                    confirmation_ui=MockUserInterface(
+                        ConfirmationUIResult(action="modify", new_tool_params={"a": 2, "b": 3})
+                    ),
+                )
+            }
+        )
+        agent = Agent(
+            chat_generator=OpenAIChatGenerator(model="gpt-4o-mini"),
+            tools=tools,
+            hooks={"before_tool": [confirmation_hook]},
+        )
+        result = await agent.run_async([ChatMessage.from_user("What is 2+2?")])
+
+        assert isinstance(result["last_message"], ChatMessage)
+        assert result["last_message"].text is not None
+        assert "5" in result["last_message"].text
+        # Auto-populated run-metadata outputs: at least one tool call plus a final answer.
+        assert result["step_count"] >= 2
+        assert result["tool_call_counts"]["addition_tool"] >= 1
+        assert result["token_usage"]["prompt_tokens"] > 0
+        assert result["token_usage"]["completion_tokens"] > 0
+        assert result["token_usage"]["total_tokens"] > 0
+
+
+@component
+class MockChatGenerator:
+    @component.output_types(replies=list[ChatMessage])
+    def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
+        return {"replies": [ChatMessage.from_assistant("Hello")]}
+
+
+def _producer() -> dict[str, str]:
+    return {"value": "PRODUCED"}
+
+
+def _consumer(value: Annotated[str, "the shared value"]) -> str:
+    return f"consumed:{value}"
+
+
+# `producer` overwrites the `shared` state key; `consumer` reads it via inputs_from_state. Called together in one
+# step, the batched executor must run producer first so consumer sees the fresh value.
+producer_tool = Tool(
+    name="producer",
+    description="Produce a value into state.",
+    parameters={"type": "object", "properties": {}, "required": []},
+    function=_producer,
+    outputs_to_state={"shared": {"source": "value"}},
+)
+consumer_tool = Tool(
+    name="consumer",
+    description="Consume the shared value.",
+    parameters={"type": "object", "properties": {"value": {"type": "string"}}, "required": ["value"]},
+    function=_consumer,
+    inputs_from_state={"shared": "value"},
+)
+
+
+class TestConfirmationStrategyToolArgPrep:
+    def test_confirmed_dependent_tool_runs_with_fresh_state(self):
+        """
+        When a confirmation strategy is configured, a tool that reads State a same-step tool writes must still run
+        with the freshly-produced value, not the stale step-start value.
+        """
+        confirmation_hook = ConfirmationHook(
+            confirmation_strategies={
+                "consumer": BlockingConfirmationStrategy(
+                    confirmation_policy=NeverAskPolicy(), confirmation_ui=SimpleConsoleUI()
+                )
+            }
+        )
+        agent = Agent(
+            chat_generator=MockChatGenerator(),
+            tools=[producer_tool, consumer_tool],
+            state_schema={"shared": {"type": str}},
+            hooks={"before_tool": [confirmation_hook]},
+        )
+        agent.warm_up()
+        # Step 1: the model calls producer and consumer together (consumer relies on inputs_from_state for `value`).
+        # Step 2: a plain text reply ends the run.
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {
+                    "replies": [
+                        ChatMessage.from_assistant(tool_calls=[ToolCall("producer", {}), ToolCall("consumer", {})])
+                    ]
+                },
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+
+        # `shared` starts stale; producer overwrites it to "PRODUCED" before consumer runs.
+        result = agent.run(messages=[ChatMessage.from_user("go")], shared="OLD")
+
+        consumer_results = [
+            m.tool_call_result.result
+            for m in result["messages"]
+            if m.tool_call_result is not None and m.tool_call_result.origin.tool_name == "consumer"
+        ]
+        assert consumer_results == ["consumed:PRODUCED"]
+
+
+def _echo(x: Annotated[int, "the value to echo"]) -> dict[str, int]:
+    return {"echoed": x}
+
+
+def _echo_tool(name: str) -> Tool:
+    return Tool(
+        name=name,
+        description=f"Echo tool {name}.",
+        parameters={"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]},
+        function=_echo,
+    )
+
+
+def _single_tool_hook(tool_name: str, ui_result: ConfirmationUIResult) -> ConfirmationHook:
+    return ConfirmationHook(
+        confirmation_strategies={
+            tool_name: BlockingConfirmationStrategy(
+                confirmation_policy=AlwaysAskPolicy(), confirmation_ui=MockUserInterface(ui_result)
+            )
+        }
+    )
+
+
+class TestMultipleConfirmationHooks:
+    def test_multiple_hooks_each_targeting_a_different_tool(self):
+        """
+        Three before_tool ConfirmationHooks, each owning a different tool of a three-call batch, compose correctly:
+        a tool a hook does not own passes through unchanged, so reject/modify decisions from all three hooks land on
+        the right calls in the final, re-read pending message.
+        """
+        foo, bar, baz = _echo_tool("foo"), _echo_tool("bar"), _echo_tool("baz")
+        hooks = {
+            "before_tool": [
+                _single_tool_hook("foo", ConfirmationUIResult(action="reject")),
+                _single_tool_hook("bar", ConfirmationUIResult(action="modify", new_tool_params={"x": 99})),
+                _single_tool_hook("baz", ConfirmationUIResult(action="modify", new_tool_params={"x": 77})),
+            ]
+        }
+        agent = Agent(chat_generator=MockChatGenerator(), tools=[foo, bar, baz], hooks=hooks)
+        agent.warm_up()
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {
+                    "replies": [
+                        ChatMessage.from_assistant(
+                            tool_calls=[ToolCall("foo", {"x": 1}), ToolCall("bar", {"x": 2}), ToolCall("baz", {"x": 3})]
+                        )
+                    ]
+                },
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+
+        result = agent.run(messages=[ChatMessage.from_user("go")])
+
+        # The full transcript, in order: foo is rejected (its call + an error tool result appear before the
+        # surviving batch), bar and baz are modified (each with an explanation message) and then executed.
+        assert result["messages"] == [
+            ChatMessage.from_user("go"),
+            ChatMessage.from_assistant(tool_calls=[ToolCall("foo", {"x": 1})]),
+            ChatMessage.from_tool(
+                tool_result="Tool execution for 'foo' was rejected by the user.",
+                origin=ToolCall("foo", {"x": 1}),
+                error=True,
+            ),
+            ChatMessage.from_user("The parameters for tool 'bar' were updated by the user to:\n{'x': 99}"),
+            ChatMessage.from_user("The parameters for tool 'baz' were updated by the user to:\n{'x': 77}"),
+            ChatMessage.from_assistant(tool_calls=[ToolCall("bar", {"x": 99}), ToolCall("baz", {"x": 77})]),
+            ChatMessage.from_tool(tool_result='{"echoed": 99}', origin=ToolCall("bar", {"x": 99}), error=False),
+            ChatMessage.from_tool(tool_result='{"echoed": 77}', origin=ToolCall("baz", {"x": 77}), error=False),
+            ChatMessage.from_assistant("done"),
+        ]
+        # Only the two confirmed tools actually executed.
+        assert result["tool_call_counts"] == {"foo": 0, "bar": 1, "baz": 1}

test/components/agents/test_agent_hooks.py

@@ -0,0 +1,504 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Annotated, Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from haystack import component
+from haystack.components.agents import Agent
+from haystack.components.agents.state import State, replace_values
+from haystack.dataclasses import ChatMessage, ToolCall
+from haystack.hooks import hook
+from haystack.tools import Tool, Toolset, tool
+
+
+@component
+class MockChatGenerator:
+    @component.output_types(replies=list[ChatMessage])
+    def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
+        return {"replies": [ChatMessage.from_assistant("Hello")]}
+
+    @component.output_types(replies=list[ChatMessage])
+    async def run_async(
+        self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
+    ) -> dict[str, Any]:
+        return {"replies": [ChatMessage.from_assistant("Hello")]}
+
+
+@tool
+def save(content: Annotated[str, "what to save"]) -> str:
+    """Save content."""
+    return "saved"
+
+
+@tool
+def final_answer(answer: Annotated[str, "the answer"]) -> str:
+    """Provide the final answer."""
+    return answer
+
+
+# --- Module-level hooks (importable so they can be serialized) ---
+
+
+@hook
+def record_a(state: State) -> None:
+    state.set("trace", ["a"])
+
+
+@hook
+def record_b(state: State) -> None:
+    state.set("trace", ["b"])
+
+
+@hook
+def build_context(state: State) -> None:
+    if state.get("step_count") == 0:
+        state.set("messages", [ChatMessage.from_system("INJECTED")])
+
+
+@hook
+def record_tool_calls(state: State) -> None:
+    # Realistic before_tool use case: audit the tool calls the model is about to run.
+    pending = state.data["messages"][-1].tool_calls
+    state.set("trace", [tc.tool_name for tc in pending])
+
+
+@hook
+def replace_pending_save_call(state: State) -> None:
+    messages = state.data["messages"]
+    replacement = ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "changed"})])
+    state.set("messages", [*messages[:-1], replacement], handler_override=replace_values)
+
+
+@hook
+def replace_pending_call_with_non_tool_message(state: State) -> None:
+    messages = state.data["messages"]
+    replacement = ChatMessage.from_assistant("Skipping tool execution.")
+    state.set("messages", [*messages[:-1], replacement], handler_override=replace_values)
+
+
+@hook
+def record_after_tool(state: State) -> None:
+    state.set("trace", ["after_tool"])
+
+
+@hook
+def offload_tool_results(state: State) -> None:
+    # Realistic after_tool use case: replace freshly produced tool results with a compact pointer, preserving origin.
+    rewritten = [
+        ChatMessage.from_tool(tool_result="OFFLOADED", origin=m.tool_call_result.origin, error=m.tool_call_result.error)
+        if m.tool_call_result is not None
+        else m
+        for m in state.data["messages"]
+    ]
+    state.set("messages", rewritten, handler_override=replace_values)
+
+
+@hook
+def require_save(state: State) -> None:
+    if state.get("tool_call_counts", {}).get("save", 0) == 0:
+        state.set("messages", [ChatMessage.from_system("You must call save before finishing.")])
+        state.set("continue_run", True)
+
+
+@hook
+def always_continue(state: State) -> None:
+    state.set("messages", [ChatMessage.from_system("keep going")])
+    state.set("continue_run", True)
+
+
+@hook
+def critique(state: State) -> None:
+    # Push back on the first final answer to force one more loop.
+    if state.get("tool_call_counts", {}).get("final_answer", 0) < 2:
+        state.set("messages", [ChatMessage.from_user("Please expand.")])
+        state.set("continue_run", True)
+
+
+class LifecycleHook:
+    """A class-based hook that records its lifecycle calls (e.g. opening/closing a client)."""
+
+    def __init__(self) -> None:
+        self.warmed = 0
+        self.warmed_async = 0
+        self.closed = 0
+        self.closed_async = 0
+
+    def run(self, state: State) -> None:
+        pass
+
+    def warm_up(self) -> None:
+        self.warmed += 1
+
+    async def warm_up_async(self) -> None:
+        self.warmed_async += 1
+
+    def close(self) -> None:
+        self.closed += 1
+
+    async def close_async(self) -> None:
+        self.closed_async += 1
+
+
+def _agent(generator, **kwargs):
+    agent = Agent(chat_generator=generator, **kwargs)
+    agent.warm_up()
+    return agent
+
+
+class TestAgentHooksValidation:
+    def test_invalid_hook_point_raises(self):
+        with pytest.raises(ValueError):
+            Agent(chat_generator=MockChatGenerator(), hooks={"bogus": [record_a]})
+
+    def test_non_callable_object_raises(self):
+        with pytest.raises(TypeError, match="must have a callable 'run"):
+            Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [object()]})
+
+    def test_unwrapped_function_hints_at_hook_decorator(self):
+        def my_hook(state: State) -> None:
+            pass
+
+        with pytest.raises(TypeError, match="@hook decorator"):
+            Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [my_hook]})
+
+    def test_continue_run_is_a_reserved_state_schema_key(self):
+        with pytest.raises(ValueError):
+            Agent(chat_generator=MockChatGenerator(), state_schema={"continue_run": {"type": bool}})
+
+    def test_continue_run_is_not_exposed_as_output(self):
+        agent = _agent(MockChatGenerator())
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert "continue_run" not in result
+
+
+class TestBeforeLlmHook:
+    def test_step_0_only_context_injected_once_across_loops(self):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_llm": [build_context]})
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        # build_context only injects on step_count == 0, so it appears once despite the two-loop run
+        injected = [m for m in result["messages"] if m.is_from("system") and m.text == "INJECTED"]
+        assert len(injected) == 1
+        assert agent.chat_generator.run.call_count == 2
+
+    def test_hooks_run_in_list_order(self):
+        agent = _agent(
+            MockChatGenerator(), state_schema={"trace": {"type": list}}, hooks={"before_llm": [record_a, record_b]}
+        )
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert result["trace"] == ["a", "b"]
+
+
+class TestBeforeToolHook:
+    def test_runs_before_tools_when_model_calls_a_tool(self):
+        agent = _agent(
+            MockChatGenerator(),
+            tools=[save],
+            state_schema={"trace": {"type": list}},
+            hooks={"before_tool": [record_tool_calls]},
+        )
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        # before_tool sees the pending tool call and does not fire on the final text-only step
+        assert result["trace"] == ["save"]
+        assert result["tool_call_counts"]["save"] == 1
+
+    def test_rereads_last_state_message_after_hook(self):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_tool": [replace_pending_save_call]})
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "original"})])]},
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        tool_messages = [m for m in result["messages"] if m.tool_call_result is not None]
+        assert len(tool_messages) == 1
+        assert tool_messages[0].tool_call_result.origin.arguments == {"content": "changed"}
+
+    def test_skips_tool_execution_when_hook_leaves_last_message_without_tool_calls(self):
+        agent = _agent(
+            MockChatGenerator(), tools=[save], hooks={"before_tool": [replace_pending_call_with_non_tool_message]}
+        )
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 2
+        assert result["tool_call_counts"]["save"] == 0
+        assert not [m for m in result["messages"] if m.tool_call_result is not None]
+        assert [m.text for m in result["messages"]][-2:] == ["Skipping tool execution.", "done"]
+
+
+class TestAfterToolHook:
+    def test_runs_after_tools_and_not_on_text_only_exit(self):
+        agent = _agent(
+            MockChatGenerator(),
+            tools=[save],
+            state_schema={"trace": {"type": list}},
+            hooks={"after_tool": [record_after_tool]},
+        )
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        # Fires once for the single tool step; does not fire on the final text-only step (no tools ran there).
+        assert result["trace"] == ["after_tool"]
+
+    def test_does_not_run_when_no_tools_are_called(self):
+        agent = _agent(
+            MockChatGenerator(),
+            tools=[save],
+            state_schema={"trace": {"type": list}},
+            hooks={"after_tool": [record_after_tool]},
+        )
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert result.get("trace", []) == []
+
+    def test_rewrites_tool_results_before_next_llm_call(self):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [offload_tool_results]})
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        # The rewritten result is what the next LLM call sees...
+        second_call_messages = agent.chat_generator.run.call_args_list[1].kwargs["messages"]
+        assert [m.tool_call_result.result for m in second_call_messages if m.tool_call_result is not None] == [
+            "OFFLOADED"
+        ]
+        # ...and what ends up in the final conversation.
+        assert [m.tool_call_result.result for m in result["messages"] if m.tool_call_result is not None] == [
+            "OFFLOADED"
+        ]
+
+    def test_rewriting_results_does_not_change_exit_matching(self):
+        # Exit conditions match on tool name + error, not result content, so offloading the result still exits.
+        agent = _agent(
+            MockChatGenerator(),
+            tools=[final_answer],
+            exit_conditions=["final_answer"],
+            hooks={"after_tool": [offload_tool_results]},
+        )
+        agent.chat_generator.run = MagicMock(
+            return_value={
+                "replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "42"})])]
+            }
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 1
+        assert [m.tool_call_result.result for m in result["messages"] if m.tool_call_result is not None] == [
+            "OFFLOADED"
+        ]
+
+
+class TestOnExitHook:
+    def test_mandatory_tool_forces_extra_step(self):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"on_exit": [require_save]})
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant("Done")]},  # tries to exit without saving
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
+                {"replies": [ChatMessage.from_assistant("All finished")]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 3
+        assert result["tool_call_counts"]["save"] == 1
+        assert result["last_message"].text == "All finished"
+
+    def test_readonly_hook_does_not_change_exit(self):
+        fired = []
+
+        def record(state: State) -> None:
+            fired.append(1)
+
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"on_exit": [hook(record)]})
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("Final")]})
+        agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 1
+        assert fired == [1]
+
+    def test_critique_on_tool_based_exit(self):
+        agent = _agent(
+            MockChatGenerator(), tools=[final_answer], exit_conditions=["final_answer"], hooks={"on_exit": [critique]}
+        )
+        agent.chat_generator.run = MagicMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "short"})])]},
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "longer"})])]},
+            ]
+        )
+        result = agent.run(messages=[ChatMessage.from_user("q")])
+        # critique forces one extra loop before the final answer is accepted
+        assert agent.chat_generator.run.call_count == 2
+        assert result["tool_call_counts"]["final_answer"] == 2
+
+    def test_max_agent_steps_bounds_always_cancel_hook(self):
+        agent = _agent(MockChatGenerator(), max_agent_steps=3, hooks={"on_exit": [always_continue]})
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("text")]})
+        agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 3
+
+    def test_continue_run_from_before_llm_hook_does_not_force_continuation(self):
+        def leak_continue(state: State) -> None:
+            state.set("continue_run", True)
+
+        def audit(state: State) -> None:
+            pass
+
+        agent = _agent(
+            MockChatGenerator(),
+            max_agent_steps=5,
+            hooks={"before_llm": [hook(leak_continue)], "on_exit": [hook(audit)]},
+        )
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("Final")]})
+        agent.run(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run.call_count == 1
+
+
+class TestHookReuseAcrossHookPoints:
+    def test_same_hook_under_two_hook_points(self):
+        counter = []
+
+        def count_call(state: State) -> None:
+            counter.append(1)
+
+        counting = hook(count_call)
+        agent = _agent(MockChatGenerator(), hooks={"before_llm": [counting], "on_exit": [counting]})
+        agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
+        agent.run(messages=[ChatMessage.from_user("hi")])
+        # once for before_llm, once for on_exit (which appends nothing, so the run still exits)
+        assert len(counter) == 2
+
+
+class TestAgentHooksSerde:
+    def test_to_dict_from_dict_roundtrip(self, monkeypatch):
+        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
+        from haystack.components.generators.chat import OpenAIChatGenerator
+
+        agent = Agent(
+            chat_generator=OpenAIChatGenerator(),
+            tools=[save],
+            hooks={"before_llm": [build_context], "on_exit": [require_save]},
+        )
+        restored = Agent.from_dict(agent.to_dict())
+        assert set(restored.hooks) == {"before_llm", "on_exit"}
+        assert restored.hooks["before_llm"][0].function is build_context.function
+        assert restored.hooks["on_exit"][0].function is require_save.function
+
+
+class TestAgentHooksAsync:
+    @pytest.mark.asyncio
+    async def test_sync_hook_runs_in_async_run(self):
+        fired = []
+
+        def record(state: State) -> None:
+            fired.append(1)
+
+        agent = _agent(MockChatGenerator(), hooks={"before_llm": [hook(record)]})
+        agent.chat_generator.run_async = AsyncMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
+        await agent.run_async(messages=[ChatMessage.from_user("hi")])
+        assert fired == [1]
+
+    @pytest.mark.asyncio
+    async def test_after_tool_hook_rewrites_results(self):
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [offload_tool_results]})
+        agent.chat_generator.run_async = AsyncMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
+                {"replies": [ChatMessage.from_assistant("done")]},
+            ]
+        )
+        result = await agent.run_async(messages=[ChatMessage.from_user("hi")])
+        assert [m.tool_call_result.result for m in result["messages"] if m.tool_call_result is not None] == [
+            "OFFLOADED"
+        ]
+
+    @pytest.mark.asyncio
+    async def test_async_on_exit_hook_forces_extra_step(self):
+        async def require_save_async(state: State) -> None:
+            if state.get("tool_call_counts", {}).get("save", 0) == 0:
+                state.set("messages", [ChatMessage.from_system("call save first")])
+                state.set("continue_run", True)
+
+        agent = _agent(MockChatGenerator(), tools=[save], hooks={"on_exit": [hook(require_save_async)]})
+        agent.chat_generator.run_async = AsyncMock(
+            side_effect=[
+                {"replies": [ChatMessage.from_assistant("Done")]},
+                {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
+                {"replies": [ChatMessage.from_assistant("All finished")]},
+            ]
+        )
+        result = await agent.run_async(messages=[ChatMessage.from_user("hi")])
+        assert agent.chat_generator.run_async.call_count == 3
+        assert result["tool_call_counts"]["save"] == 1
+
+
+class TestAgentHookLifecycle:
+    def test_init_does_not_warm_up(self):
+        h = LifecycleHook()
+        Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
+        assert h.warmed == 0
+
+    def test_warm_up_warms_hooks_once(self):
+        h = LifecycleHook()
+        agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
+        agent.warm_up()
+        agent.warm_up()
+        assert h.warmed == 1
+
+    def test_reused_hook_warmed_once(self):
+        h = LifecycleHook()
+        agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h], "on_exit": [h]})
+        agent.warm_up()
+        assert h.warmed == 1
+
+    def test_close_closes_hooks(self):
+        h = LifecycleHook()
+        agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
+        agent.close()
+        assert h.closed == 1
+
+
+class TestAgentHookLifecycleAsync:
+    @pytest.mark.asyncio
+    async def test_warm_up_async_prefers_async(self):
+        h = LifecycleHook()
+        agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
+        await agent.warm_up_async()
+        assert h.warmed_async == 1
+        assert h.warmed == 0
+
+    @pytest.mark.asyncio
+    async def test_close_async_prefers_async(self):
+        h = LifecycleHook()
+        agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
+        await agent.close_async()
+        assert h.closed_async == 1
+        assert h.closed == 0

test/components/agents/test_state_class.py

@@ -0,0 +1,561 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+from dataclasses import dataclass
+from typing import Dict, Generic, List, Optional, TypeVar, Union
+
+import pytest
+
+from haystack.components.agents.state.state import (
+    State,
+    _is_list_type,
+    _is_valid_type,
+    _schema_from_dict,
+    _schema_to_dict,
+    _validate_schema,
+    merge_lists,
+)
+from haystack.dataclasses import ChatMessage
+
+
+@pytest.fixture
+def basic_schema():
+    return {"numbers": {"type": list}, "metadata": {"type": dict}, "name": {"type": str}}
+
+
+def numbers_handler(current, new):
+    if current is None:
+        return sorted(set(new))
+    return sorted(set(current + new))
+
+
+@pytest.fixture
+def complex_schema():
+    return {"numbers": {"type": list, "handler": numbers_handler}, "metadata": {"type": dict}, "name": {"type": str}}
+
+
+def test_is_list_type():
+    assert _is_list_type(list) is True
+    assert _is_list_type(list[int]) is True
+    assert _is_list_type(list[str]) is True
+    assert _is_list_type(dict) is False
+    assert _is_list_type(int) is False
+    assert _is_list_type(Union[list[int], None]) is False
+    assert _is_list_type(list[int] | None) is False
+
+
+class TestMergeLists:
+    def test_merge_two_lists(self):
+        current = [1, 2, 3]
+        new = [4, 5, 6]
+        result = merge_lists(current, new)
+        assert result == [1, 2, 3, 4, 5, 6]
+        # Ensure original lists weren't modified
+        assert current == [1, 2, 3]
+        assert new == [4, 5, 6]
+
+    def test_append_to_list(self):
+        current = [1, 2, 3]
+        new = 4
+        result = merge_lists(current, new)
+        assert result == [1, 2, 3, 4]
+        assert current == [1, 2, 3]  # Ensure original wasn't modified
+
+    def test_create_new_list(self):
+        current = 1
+        new = 2
+        result = merge_lists(current, new)
+        assert result == [1, 2]
+
+    def test_replace_with_list(self):
+        current = 1
+        new = [2, 3]
+        result = merge_lists(current, new)
+        assert result == [1, 2, 3]
+
+
+class TestIsValidType:
+    def test_builtin_types(self):
+        assert _is_valid_type(str) is True
+        assert _is_valid_type(int) is True
+        assert _is_valid_type(dict) is True
+        assert _is_valid_type(list) is True
+        assert _is_valid_type(tuple) is True
+        assert _is_valid_type(set) is True
+        assert _is_valid_type(bool) is True
+        assert _is_valid_type(float) is True
+
+    def test_generic_types(self):
+        assert _is_valid_type(list[str]) is True
+        assert _is_valid_type(List[str]) is True
+        assert _is_valid_type(dict[str, int]) is True
+        assert _is_valid_type(Dict[str, int]) is True
+        assert _is_valid_type(list[dict[str, int]]) is True
+        assert _is_valid_type(List[Dict[str, int]]) is True
+        assert _is_valid_type(dict[str, list[int]]) is True
+        assert _is_valid_type(Dict[str, List[int]]) is True
+
+    def test_custom_classes(self):
+        @dataclass
+        class CustomClass:
+            value: int
+
+        T = TypeVar("T")
+
+        class GenericCustomClass(Generic[T]):
+            pass
+
+        # Test regular and generic custom classes
+        assert _is_valid_type(CustomClass) is True
+        assert _is_valid_type(GenericCustomClass) is True
+        assert _is_valid_type(GenericCustomClass[int]) is True
+
+        # Test generic types with custom classes
+        assert _is_valid_type(list[CustomClass]) is True
+        assert _is_valid_type(List[CustomClass]) is True
+        assert _is_valid_type(dict[str, CustomClass]) is True
+        assert _is_valid_type(Dict[str, CustomClass]) is True
+        assert _is_valid_type(dict[str, GenericCustomClass[int]]) is True
+        assert _is_valid_type(Dict[str, GenericCustomClass[int]]) is True
+
+    def test_invalid_types(self):
+        # Test regular values
+        assert _is_valid_type(42) is False
+        assert _is_valid_type("string") is False
+        assert _is_valid_type([1, 2, 3]) is False
+        assert _is_valid_type({"a": 1}) is False
+        assert _is_valid_type(True) is False
+
+        # Test class instances
+        @dataclass
+        class SampleClass:
+            value: int
+
+        instance = SampleClass(42)
+        assert _is_valid_type(instance) is False
+
+        # Test callable objects
+        assert _is_valid_type(len) is False
+        assert _is_valid_type(lambda x: x) is False
+        assert _is_valid_type(print) is False
+
+    def test_union_and_optional_types(self):
+        # Test basic Union types
+        assert _is_valid_type(Union[str, int]) is True
+        assert _is_valid_type(Union[str, None]) is True
+        assert _is_valid_type(Union[list[int], dict[str, str]]) is True
+
+        # Test Optional types (which are Union[T, None])
+        assert _is_valid_type(Optional[str]) is True
+        assert _is_valid_type(Optional[list[int]]) is True
+        assert _is_valid_type(Optional[dict[str, list]]) is True
+
+        # Test that Union itself is not a valid type (only instantiated Unions are)
+        assert _is_valid_type(Union) is False
+
+        # Test PEP 604 union types (X | Y syntax)
+        assert _is_valid_type(str | int) is True
+        assert _is_valid_type(str | None) is True
+        assert _is_valid_type(list[int] | dict[str, str]) is True
+
+        # Test PEP 604 Optional-like types (X | None syntax)
+        assert _is_valid_type(list[int] | None) is True
+        assert _is_valid_type(dict[str, list] | None) is True
+
+    def test_nested_generic_types(self):
+        assert _is_valid_type(list[list[dict[str, list[int]]]]) is True
+        assert _is_valid_type(dict[str, list[dict[str, set]]]) is True
+        assert _is_valid_type(dict[str, Optional[list[int]]]) is True
+        assert _is_valid_type(list[Union[str, dict[str, list[int]]]]) is True
+        # PEP 604 nested types
+        assert _is_valid_type(dict[str, list[int] | None]) is True
+        assert _is_valid_type(list[str | dict[str, list[int]]]) is True
+
+    def test_edge_cases(self):
+        # Test None and NoneType
+        assert _is_valid_type(None) is False
+        assert _is_valid_type(type(None)) is True
+
+        # Test functions and methods
+        def sample_func():
+            pass
+
+        assert _is_valid_type(sample_func) is False
+        assert _is_valid_type(type(sample_func)) is True
+
+        # Test modules
+        assert _is_valid_type(inspect) is False
+
+        # Test type itself
+        assert _is_valid_type(type) is True
+
+    @pytest.mark.parametrize(
+        "test_input,expected",
+        [
+            (str, True),
+            (int, True),
+            (list[int], True),
+            (dict[str, int], True),
+            (List[int], True),
+            (Dict[str, int], True),
+            (Union[str, int], True),
+            (Optional[str], True),
+            # PEP 604 union types
+            (str | int, True),
+            (str | None, True),
+            (list[int] | None, True),
+            (42, False),
+            ("string", False),
+            ([1, 2, 3], False),
+            (lambda x: x, False),
+        ],
+    )
+    def test_parametrized_cases(self, test_input, expected):
+        assert _is_valid_type(test_input) is expected
+
+
+class TestState:
+    def test_validate_schema_valid(self, basic_schema):
+        # Should not raise any exceptions
+        _validate_schema(basic_schema)
+
+    def test_validate_schema_invalid_type(self):
+        invalid_schema = {"test": {"type": "not_a_type"}}
+        with pytest.raises(ValueError, match="must be a Python type"):
+            _validate_schema(invalid_schema)
+
+    def test_validate_schema_missing_type(self):
+        invalid_schema = {"test": {"handler": lambda x, y: x + y}}
+        with pytest.raises(ValueError, match="missing a 'type' entry"):
+            _validate_schema(invalid_schema)
+
+    def test_validate_schema_invalid_handler(self):
+        invalid_schema = {"test": {"type": list, "handler": "not_callable"}}
+        with pytest.raises(ValueError, match="must be callable or None"):
+            _validate_schema(invalid_schema)
+
+    def test_validate_schema_with_messages(self):
+        class ChatMessageSubclass(ChatMessage):
+            pass
+
+        schema_with_messages = {"messages": {"type": List[ChatMessage]}}
+        _validate_schema(schema_with_messages)
+
+        schema_with_messages_subclass = {"messages": {"type": List[ChatMessageSubclass]}}
+        _validate_schema(schema_with_messages_subclass)
+
+    def test_state_initialization(self, basic_schema):
+        # Test empty initialization
+        state = State(basic_schema)
+        assert state.data == {}
+
+        # Test initialization with data
+        initial_data = {"numbers": [1, 2, 3], "name": "test"}
+        state = State(basic_schema, initial_data)
+        assert state.data["numbers"] == [1, 2, 3]
+        assert state.data["name"] == "test"
+
+    def test_state_get(self, basic_schema):
+        state = State(basic_schema, {"name": "test"})
+        assert state.get("name") == "test"
+        assert state.get("non_existent") is None
+        assert state.get("non_existent", "default") == "default"
+
+    def test_state_set_basic(self, basic_schema):
+        state = State(basic_schema)
+
+        # Test setting new values
+        state.set("numbers", [1, 2])
+        assert state.get("numbers") == [1, 2]
+
+        # Test updating existing values
+        state.set("numbers", [3, 4])
+        assert state.get("numbers") == [1, 2, 3, 4]
+
+    def test_state_set_with_handler(self, complex_schema):
+        state = State(complex_schema)
+
+        # Test custom handler for numbers
+        state.set("numbers", [3, 2, 1])
+        assert state.get("numbers") == [1, 2, 3]
+
+        state.set("numbers", [6, 5, 4])
+        assert state.get("numbers") == [1, 2, 3, 4, 5, 6]
+
+    def test_state_set_with_handler_override(self, basic_schema):
+        state = State(basic_schema)
+
+        # Custom handler that concatenates strings
+        custom_handler = lambda current, new: f"{current}-{new}" if current else new
+
+        state.set("name", "first")
+        state.set("name", "second", handler_override=custom_handler)
+        assert state.get("name") == "first-second"
+
+    def test_state_has(self, basic_schema):
+        state = State(basic_schema, {"name": "test"})
+        assert state.has("name") is True
+        assert state.has("non_existent") is False
+
+    def test_state_empty_schema(self):
+        state = State({})
+        assert state.data == {}
+
+        # Instead of comparing the entire schema directly, check structure separately
+        assert "messages" in state.schema
+        assert state.schema["messages"]["type"] == list[ChatMessage]
+        assert callable(state.schema["messages"]["handler"])
+
+        with pytest.raises(ValueError, match="Key 'any_key' not found in schema"):
+            state.set("any_key", "value")
+
+    def test_state_none_values(self, basic_schema):
+        state = State(basic_schema)
+        state.set("name", None)
+        assert state.get("name") is None
+        state.set("name", "value")
+        assert state.get("name") == "value"
+
+    def test_state_merge_lists(self, basic_schema):
+        state = State(basic_schema)
+        state.set("numbers", "not_a_list")
+        assert state.get("numbers") == ["not_a_list"]
+        state.set("numbers", [1, 2])
+        assert state.get("numbers") == ["not_a_list", 1, 2]
+
+    def test_state_nested_structures(self):
+        schema = {
+            "complex": {
+                "type": dict[str, list[int]],
+                "handler": lambda current, new: (
+                    {k: current.get(k, []) + new.get(k, []) for k in set(current.keys()) | set(new.keys())}
+                    if current
+                    else new
+                ),
+            }
+        }
+
+        state = State(schema)
+        state.set("complex", {"a": [1, 2], "b": [3, 4]})
+        state.set("complex", {"b": [5, 6], "c": [7, 8]})
+
+        expected = {"a": [1, 2], "b": [3, 4, 5, 6], "c": [7, 8]}
+        assert state.get("complex") == expected
+
+    def test_schema_to_dict(self, basic_schema):
+        expected_dict = {"numbers": {"type": "list"}, "metadata": {"type": "dict"}, "name": {"type": "str"}}
+        result = _schema_to_dict(basic_schema)
+        assert result == expected_dict
+
+    def test_schema_to_dict_with_handlers(self, complex_schema):
+        expected_dict = {
+            "numbers": {"type": "list", "handler": "test_state_class.numbers_handler"},
+            "metadata": {"type": "dict"},
+            "name": {"type": "str"},
+        }
+        result = _schema_to_dict(complex_schema)
+        assert result == expected_dict
+
+    def test_schema_from_dict(self, basic_schema):
+        schema_dict = {"numbers": {"type": "list"}, "metadata": {"type": "dict"}, "name": {"type": "str"}}
+        result = _schema_from_dict(schema_dict)
+        assert result == basic_schema
+
+    def test_schema_from_dict_with_handlers(self, complex_schema):
+        schema_dict = {
+            "numbers": {"type": "list", "handler": "test_state_class.numbers_handler"},
+            "metadata": {"type": "dict"},
+            "name": {"type": "str"},
+        }
+        result = _schema_from_dict(schema_dict)
+        assert result == complex_schema
+
+    def test_state_mutability(self):
+        state = State({"my_list": {"type": list}}, {"my_list": [1, 2]})
+
+        my_list = state.get("my_list")
+        my_list.append(3)
+
+        assert state.get("my_list") == [1, 2]
+
+    def test_state_to_dict(self):
+        # we test dict, a python type and a haystack dataclass
+        state_schema = {
+            "numbers": {"type": int},
+            "messages": {"type": list[ChatMessage]},
+            "dict_of_lists": {"type": dict},
+        }
+
+        data = {
+            "numbers": 1,
+            "messages": [ChatMessage.from_user(text="Hello, world!")],
+            "dict_of_lists": {"numbers": [1, 2, 3]},
+        }
+        state = State(state_schema, data)
+        state_dict = state.to_dict()
+        assert state_dict["schema"] == {
+            "numbers": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"},
+            "messages": {
+                "type": "list[haystack.dataclasses.chat_message.ChatMessage]",
+                "handler": "haystack.components.agents.state.state_utils.merge_lists",
+            },
+            "dict_of_lists": {"type": "dict", "handler": "haystack.components.agents.state.state_utils.replace_values"},
+        }
+        assert state_dict["data"] == {
+            "serialization_schema": {
+                "type": "object",
+                "properties": {
+                    "numbers": {"type": "integer"},
+                    "messages": {"type": "array", "items": {"type": "haystack.dataclasses.chat_message.ChatMessage"}},
+                    "dict_of_lists": {
+                        "type": "object",
+                        "properties": {"numbers": {"type": "array", "items": {"type": "integer"}}},
+                    },
+                },
+            },
+            "serialized_data": {
+                "numbers": 1,
+                "messages": [{"role": "user", "meta": {}, "name": None, "content": [{"text": "Hello, world!"}]}],
+                "dict_of_lists": {"numbers": [1, 2, 3]},
+            },
+        }
+
+    def test_state_from_dict(self):
+        state_dict = {
+            "schema": {
+                "numbers": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"},
+                "messages": {
+                    "type": "list[haystack.dataclasses.chat_message.ChatMessage]",
+                    "handler": "haystack.components.agents.state.state_utils.merge_lists",
+                },
+                "dict_of_lists": {
+                    "type": "dict",
+                    "handler": "haystack.components.agents.state.state_utils.replace_values",
+                },
+            },
+            "data": {
+                "serialization_schema": {
+                    "type": "object",
+                    "properties": {
+                        "numbers": {"type": "integer"},
+                        "messages": {
+                            "type": "array",
+                            "items": {"type": "haystack.dataclasses.chat_message.ChatMessage"},
+                        },
+                        "dict_of_lists": {
+                            "type": "object",
+                            "properties": {"numbers": {"type": "array", "items": {"type": "integer"}}},
+                        },
+                    },
+                },
+                "serialized_data": {
+                    "numbers": 1,
+                    "messages": [{"role": "user", "meta": {}, "name": None, "content": [{"text": "Hello, world!"}]}],
+                    "dict_of_lists": {"numbers": [1, 2, 3]},
+                },
+            },
+        }
+        state = State.from_dict(state_dict)
+        # Check types are correctly converted
+        assert state.schema["numbers"]["type"] == int
+        assert state.schema["dict_of_lists"]["type"] == dict
+        # Check handlers are functions, not comparing exact functions as they might be different references
+        assert callable(state.schema["numbers"]["handler"])
+        assert callable(state.schema["messages"]["handler"])
+        assert callable(state.schema["dict_of_lists"]["handler"])
+        # Check data is correct
+        assert state.data["numbers"] == 1
+        assert state.data["messages"] == [ChatMessage.from_user(text="Hello, world!")]
+        assert state.data["dict_of_lists"] == {"numbers": [1, 2, 3]}
+
+    def test_state_to_dict_typing_list(self):
+        # we test dict, a python type and a haystack dataclass
+        state_schema = {
+            "numbers": {"type": int},
+            "messages": {"type": List[ChatMessage]},
+            "dict_of_lists": {"type": dict},
+        }
+
+        data = {
+            "numbers": 1,
+            "messages": [ChatMessage.from_user(text="Hello, world!")],
+            "dict_of_lists": {"numbers": [1, 2, 3]},
+        }
+        state = State(state_schema, data)
+        state_dict = state.to_dict()
+        assert state_dict["schema"] == {
+            "numbers": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"},
+            "messages": {
+                "type": "typing.List[haystack.dataclasses.chat_message.ChatMessage]",
+                "handler": "haystack.components.agents.state.state_utils.merge_lists",
+            },
+            "dict_of_lists": {"type": "dict", "handler": "haystack.components.agents.state.state_utils.replace_values"},
+        }
+        assert state_dict["data"] == {
+            "serialization_schema": {
+                "type": "object",
+                "properties": {
+                    "numbers": {"type": "integer"},
+                    "messages": {"type": "array", "items": {"type": "haystack.dataclasses.chat_message.ChatMessage"}},
+                    "dict_of_lists": {
+                        "type": "object",
+                        "properties": {"numbers": {"type": "array", "items": {"type": "integer"}}},
+                    },
+                },
+            },
+            "serialized_data": {
+                "numbers": 1,
+                "messages": [{"role": "user", "meta": {}, "name": None, "content": [{"text": "Hello, world!"}]}],
+                "dict_of_lists": {"numbers": [1, 2, 3]},
+            },
+        }
+
+    def test_state_from_dict_typing_list(self):
+        state_dict = {
+            "schema": {
+                "numbers": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"},
+                "messages": {
+                    "type": "typing.List[haystack.dataclasses.chat_message.ChatMessage]",
+                    "handler": "haystack.components.agents.state.state_utils.merge_lists",
+                },
+                "dict_of_lists": {
+                    "type": "dict",
+                    "handler": "haystack.components.agents.state.state_utils.replace_values",
+                },
+            },
+            "data": {
+                "serialization_schema": {
+                    "type": "object",
+                    "properties": {
+                        "numbers": {"type": "integer"},
+                        "messages": {
+                            "type": "array",
+                            "items": {"type": "haystack.dataclasses.chat_message.ChatMessage"},
+                        },
+                        "dict_of_lists": {
+                            "type": "object",
+                            "properties": {"numbers": {"type": "array", "items": {"type": "integer"}}},
+                        },
+                    },
+                },
+                "serialized_data": {
+                    "numbers": 1,
+                    "messages": [{"role": "user", "meta": {}, "name": None, "content": [{"text": "Hello, world!"}]}],
+                    "dict_of_lists": {"numbers": [1, 2, 3]},
+                },
+            },
+        }
+        state = State.from_dict(state_dict)
+        # Check types are correctly converted
+        assert state.schema["numbers"]["type"] == int
+        assert state.schema["dict_of_lists"]["type"] == dict
+        # Check handlers are functions, not comparing exact functions as they might be different references
+        assert callable(state.schema["numbers"]["handler"])
+        assert callable(state.schema["messages"]["handler"])
+        assert callable(state.schema["dict_of_lists"]["handler"])
+        # Check data is correct
+        assert state.data["numbers"] == 1
+        assert state.data["messages"] == [ChatMessage.from_user(text="Hello, world!")]
+        assert state.data["dict_of_lists"] == {"numbers": [1, 2, 3]}